message edit/update functionality

backend/agent.py

@@ -208,10 +208,12 @@ class ToolUsingAgent:
             output += f"[{action.toolset_name}.{action.tool_name}]:{action_result}\n"
         return output
 
-class Llama8bAgent(Agent):
+class LlamaAgent(Agent):
     identifier = "llama-8b"
+    description = "Uses the llama 3 instruct models. Model sizes: '8b', '30b'"
     schema_fields = [
         AgentField('system_prompt', 'string', default="You are a helpful AI assistant."),
+        AgentField('model_size', 'string', default="8b"),
     ]
     tools = []
 
@@ -228,13 +230,14 @@ class Llama8bAgent(Agent):
 
     async def received_message(self, message):
         def sender_name(message):
-            sender = message['sender']
+            sender = message.sender
             is_digit = isinstance(sender, (int, str)) and str(sender).isdigit()
             is_self = is_digit and str(sender) == str(self.id)
             return self.name if is_self else sender
 
-        chat_context = self.app.get_chat_context(message.chat_id)
-        chat_summary = "\n".join(f"[{sender_name(message)}]: {message['content']}" for message in chat_context.messages)
+        chat_view = self.app.get_chat_view(message.chat_id)
+        thread_messages = chat_view.get_thread(message.thread_id).messages
+        chat_summary = "\n".join(f"[{sender_name(message)}]: {message.content}" for message in thread_messages)
         messages = [
             LlamaMessage("system", f"{self.config['system_prompt']}\nYou are given the context of the conversation, respond with your next message in the conversation."),
             LlamaMessage("user", f"Recent Messages:\n{chat_summary}")
@@ -248,15 +251,17 @@ class Llama8bAgent(Agent):
             thread_id=message.thread_id,
         )
 
-class BasicAgent(Llama8bAgent):
+class BasicAgent(LlamaAgent):
     identifier = "basic"
-    schema_fields = Llama8bAgent.schema_fields
-    tools = Llama8bAgent.tools
+    schema_fields = LlamaAgent.schema_fields
+    tools = LlamaAgent.tools
 
-class CoderV1Agent(Llama8bAgent):
+class CoderV1Agent(LlamaAgent):
     identifier = "coder-v1"
+    description = "Optimized for coding tasks. Based on the llama agent:\n" + LlamaAgent.description
     schema_fields = [
         AgentField('system_prompt', 'string', default="You are an expert Software Engineer profficient in multiple programming languages."),
+        AgentField('model_size', 'string', default="30b"),
     ]
     # TODO: add gral functions to the tool list
     tools = []

backend/app.py

@@ -38,7 +38,7 @@ class App(ChatMixin):
         self.agent_classes = [
             agent_mod.BasicAgent,
             ConversationAgent,
-            agent_mod.Llama8bAgent,
+            agent_mod.LlamaAgent,
             agent_mod.CoderV1Agent,
             agent_mod.ToolUsingAgent,
             agent_mod.DebugAgent,

backend/app_mixins/chat.py

@@ -91,6 +91,7 @@ class ChatMixin:
             message = db.Message()
             message.chat_id = chat_id
             message.time = int(datetime.now().timestamp())
+            message.last_edit = message.time
             message.sender = sender
             message.kind = 'text'  # default kind
             message.content = content
@@ -120,6 +121,36 @@ class ChatMixin:
             'content': message.content,
         }
 
+    # Updates a single chat message.
+    # Most fields are optional. If they're not passed a value, they wont be updated.
+    # When kind is set, subkind is always set too.
+    def update_chat_message(self, message_id, content=None, kind=None, subkind=None, format=None):
+        with self.sql_session() as session:
+            message = session.execute(
+                select(db.Message).where(db.Message.id == message_id)
+            ).scalars().first()
+
+            if not message:
+                raise Exception(f"Message not found with id {message_id}")
+
+            summary = []
+
+            if content != None:
+                message.content = content
+                summary.append('content')
+            if kind != None:
+                message.kind = kind
+                message.subkind = subkind
+                summary.append('kind')
+                summary.append('subkind')
+            if format != None:
+                message.format = format
+                summary.append('format')
+            #message.last_edit = int(datetime.now().timestamp())
+
+            session.commit()
+            print(f"Updated message #{message.id} fields: {summary}")
+
     def get_chat_view(self, chat_id):
         return ChatView(self, chat_id)
 

backend/database.py

@@ -80,10 +80,12 @@ class Message(Base):
     chat_id: Mapped[int] = mapped_column(ForeignKey("chats.id"))
     thread_id: Mapped[Optional[int]] = mapped_column(ForeignKey("threads.id"))
 
-    # Timestamp of this message
-    time: Mapped[int]
     # Identifier of the sender: user / system / user-$ID / agent-$ID
     sender: Mapped[str]
+    # Timestamp of this message
+    time: Mapped[int]
+    # Timestamp of this message's last edit
+    #last_edit: Mapped[int]
 
     # Predefined set of message types:
     # message, log, memory, file, note

backend/server.py

@@ -190,6 +190,16 @@ def start_server(app):
                     )
                     await send({'message': "message sent", 'messageId': message['id'], 'nonce': nonce})
 
+                elif action == "update-chat-message":
+                    message = app.update_chat_message(
+                        message_id=body['messageId'],
+                        content=body['content'],
+                        kind=body.get('kind'),
+                        subkind=body.get('subkind'),
+                        format=body.get('format'),
+                    )
+                    await send({'message': "message sent", 'messageId': body['messageId'], 'nonce': nonce})
+
                 elif action == "create-chat-thread":
                     thread_id = app.create_chat_thread(
                         chat_id=body['chatId'],

frontend/lib/chat_screen.dart

@@ -57,6 +57,7 @@ class _ChatScreenState extends State<ChatScreen> {
 
   // UI State
   Set<int> expandedThreads = {};
+  int? editingMessageId = null;
 
   // Internal Funcionality
   Offset? _longPressOffset;
@@ -153,11 +154,12 @@ class _ChatScreenState extends State<ChatScreen> {
     });
   }
 
-  void sendMessage(String content) async {
-    if (content.trim().isEmpty) return; // Don't send empty messages
+  void sendMessage() async {
+    final content =  _textController.text;
     _textController.clear();
+    if (content.trim().isEmpty) return; // Don't send empty messages
 
-    final newMessage = Message(sender: "user", content: content, threadId: currentThreadId);
+    if (editingMessageId == null) {
       Backend.sendAction("send-chat-message", {
         'chatId': widget.chatId,
         'content': content,
@@ -167,7 +169,24 @@ class _ChatScreenState extends State<ChatScreen> {
           print("Error sending chat message: ${response['error']}");
         }
       });
+
+      final newMessage = Message(sender: "user", content: content, threadId: currentThreadId);
       addMessage(newMessage);
+    } else {
+      Backend.sendAction("update-chat-message", {
+        'messageId': editingMessageId,
+        'content': content,
+      }).then((response) {
+        if (response['error'] != null) {
+          print("Error updating message: ${response['error']}");
+        }
+      });
+
+      Message message = _messages.firstWhere((element) => element.id == editingMessageId!);
+      setState(() {
+        message.content = content;
+      });
+    }
   }
 
   void createThread () async {
@@ -186,7 +205,7 @@ class _ChatScreenState extends State<ChatScreen> {
     return messageColor;
   }
 
-  void _showMessageContextMenu(BuildContext context, Offset position) {
+  void _showMessageContextMenu(BuildContext context, Message message, Offset position) {
     showMenu(
       context: context,
       position: RelativeRect.fromLTRB(position.dx, position.dy, position.dx + 1, position.dy + 1),
@@ -201,8 +220,10 @@ class _ChatScreenState extends State<ChatScreen> {
         PopupMenuItem(
           child: Text('Edit'),
           onTap: () {
-            // TODO: Implement Edit functionality
-            print('Edit clicked');
+            _textController.text = message.content;
+            setState(() {
+              editingMessageId = message.id;
+            });
           },
         ),
         PopupMenuItem(
@@ -282,11 +303,11 @@ class _ChatScreenState extends State<ChatScreen> {
                       _longPressOffset = details.globalPosition;
                     },
                     onLongPress: () {
-                      _showMessageContextMenu(context, _longPressOffset!);
+                      _showMessageContextMenu(context, message, _longPressOffset!);
                       _longPressOffset = null;
                     },
                     onSecondaryTapUp: (details) {
-                      _showMessageContextMenu(context, details.globalPosition);
+                      _showMessageContextMenu(context, message, details.globalPosition);
                     },
                     child: Padding(
                       padding: isWideScreen ? message.padding : EdgeInsets.zero,
@@ -299,7 +320,14 @@ class _ChatScreenState extends State<ChatScreen> {
                             color: getMessageColor(context),
                             borderRadius: BorderRadius.circular(12),
                           ),
-                          child: Text(message.content),
+                          child: SelectableText(
+                            message.content,
+                            toolbarOptions: ToolbarOptions(
+                              copy: true,
+                              selectAll: false,
+                            ),
+                            onTap: () {},
+                          ),
                         ),
                       ),
                     ),
@@ -315,13 +343,16 @@ class _ChatScreenState extends State<ChatScreen> {
             padding: const EdgeInsets.all(8.0),
             child: Column(
               children: [
+                if (editingMessageId!= null) ...[
+                  Text('Editing ${editingMessageId}'),
+                ],
                 Expanded(
                   child: KeyboardListener(
                     focusNode: _inputFocusNode,
                     onKeyEvent: (event) {
                       // TODO: Probably doing this in the input focus node would be cleaner
                       if (HardwareKeyboard.instance.isControlPressed && HardwareKeyboard.instance.isLogicalKeyPressed(LogicalKeyboardKey.enter)) {
-                        sendMessage(_textController.text);
+                        sendMessage();
                         //return KeyEventResult.handled;
                       }
                       //return KeyEventResult.ignored;
@@ -338,7 +369,7 @@ class _ChatScreenState extends State<ChatScreen> {
                   children: [
                     ElevatedButton(
                       onPressed: () {
-                        sendMessage(_textController.text);
+                        sendMessage();
                       },
                       child: Text('Send'),
                     ),
@@ -349,6 +380,17 @@ class _ChatScreenState extends State<ChatScreen> {
                       },
                       child: Text('Start New Topic'),
                     ),
+                    if (editingMessageId!= null) ...[
+                      ElevatedButton(
+                        onPressed: () {
+                          _textController.text = "";
+                          setState(() {
+                            editingMessageId = null;
+                          });
+                        },
+                        child: Text('Cancel Edit'),
+                      )
+                    ],
                   ],
                 ),
               ],