Chat Threads, ConversationAgent (WIP)

backend/agent.py

@@ -53,6 +53,90 @@ class Agent:
         if inspect.isawaitable(result):
             result = await result
         return result
+    
+    # Returns all resource values matching the kind and key
+    def get_resources(self, kind, key):
+        self.app
+
+
+class DebugAgent(Agent):
+    identifier = "debug"
+    schema_fields = []
+    tools = []
+
+    async def received_message(self, message):
+        chat_view = self.app.get_chat_view(message.chat_id)
+        response = chat_view.debug_string()
+
+        self.app.send_chat_message(
+            chat_id=message.chat_id,
+            sender=self.id,
+            content=response,
+            thread_id=chat_view.current_thread_id,
+        )
+
+class ConversationAgent(Agent):
+    identifier = "conversation"
+    schema_fields = []
+    tools = []
+
+    def __init__(self, app, data):
+        Agent.__init__(self, app, data)
+
+        # Constant system prompt for developing purposes 
+        self.system_prompt = "You are a friendly AI having a casual conversation with the user."
+
+    async def prompt(self, content):
+        # TODO: Advanced features
+        messages = [
+            LlamaMessage("system", self.system_prompt),
+            LlamaMessage("user", content)
+        ]
+        response = await replicate_llama_completion(self.app, messages)
+        return response
+
+    async def received_message(self, message):
+        def sender_name(message):
+            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
+
+        # TODO: before retrieving the chat context, acquire a lock for it
+
+        saved_context = ThreadContext(self.app.get_agent_resource(self.id, 'last-thread-context', message.chat_id))
+        chat_view = self.app.get_chat_view(message.chat_id)
+
+        if message.thread_id != saved_context.current_thread_id:
+            # the incoming message is not from the same thread that the agent was last aware of
+            # store details about the old thread in the agent's long term memory
+            old_thread = chat_view.thread_map.get(message.thread_id)
+            self.process_thread(old_thread)
+
+            # TODO: This is a good place to store the new context, its fine that the agent hasnt responded
+
+        current_thread = chat_view.thread_map.get(message.thread_id)
+
+        chat_summary = "\n".join(f"[{sender_name(message)}]: {message.content}" for message in current_thread.messages)
+        messages = [
+            LlamaMessage("system", f"{self.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}")
+        ]
+
+        response = await replicate_llama_completion(self.app, messages)
+        self.app.send_chat_message(
+            chat_id=message.chat_id,
+            sender=self.id,
+            content=response,
+        )
+
+        # TODO: drop the chat context lock
+    
+    # Compute thread summnary and concrete memories, then save the progress done
+    def process_thread(self, thread):
+        # TODO: extract memories from the thread messages, then store them in the agent's vector memory db
+        # TODO: somewhere here it must be known if the thread was seen before, the same memory should not be stored multiple times
+        pass
 
 class ToolUsingAgent:
     identifier = "tool-using"
@@ -245,6 +329,13 @@ class CoderV1Agent(Llama8bAgent):
     async def gral_update_file(self, message_history: list, filepath: str, file_content: str):
         return await gral.gral_update_file(self.app, self.config['system_prompt'], message_history, filepath, file_content)
 
+# Used by ConversationAgent to keep track of the last chat thread it participated in
+class ThreadContext:
+    def __init__(self, data):
+        if data:
+            self.current_thread_id = data.get('current_thread_id')
+        else:
+            self.current_thread_id = None
 
 class ParsedResponse:
     def __init__(self):

backend/app.py

@@ -12,6 +12,7 @@ import replicate
 import database as db
 import agent as agent_mod
 from server import start_server
+from chat import ChatView
 
 MESSAGE_PAGE_SIZE = 32
 
@@ -36,9 +37,11 @@ class App:
         # `Agent` is the kind 'basic', will be class `BasicAgent` in the future
         self.agent_classes = [
             agent_mod.BasicAgent,
+            agent_mod.ConversationAgent,
             agent_mod.Llama8bAgent,
             agent_mod.CoderV1Agent,
             agent_mod.ToolUsingAgent,
+            agent_mod.DebugAgent,
         ]
 
     def sql_session(self):
@@ -172,8 +175,22 @@ class App:
                 "time": message.time,
                 "sender": message.sender,
                 "kind": message.kind,
-                "content": message.content
+                "content": message.content,
+                "threadId": message.thread_id,
             } for message in messages]
+    
+    def create_chat_thread(self, chat_id, name=None):
+        with self.sql_session() as session:
+            thread = db.Chat()
+            thread.chat_id = chat_id
+            thread.name = name
+
+            session.add(thread)
+            session.commit()
+            print(f"Created thread #{thread.id}: {thread.name}")
+
+            thread_id = thread.id
+        return thread_id
 
     def verify_user_credentials(self, username, password):
         with self.sql_session() as session:
@@ -217,7 +234,7 @@ class App:
             print(f"Added agent #{agent_id} to chat #{chat_id}")
 
     # If an integer is passed in sender, it is interpreted as an agent id
-    def send_chat_message(self, chat_id, content, sender='user'):
+    def send_chat_message(self, chat_id, content, sender='user', thread_id=None):
         with self.sql_session() as session:
             message = db.Message()
             message.chat_id = chat_id
@@ -225,6 +242,7 @@ class App:
             message.sender = sender
             message.kind = 'text'  # default kind
             message.content = content
+            message.thread_id = thread_id
 
             session.add(message)
             session.commit()
@@ -252,6 +270,9 @@ class App:
     def get_chat_context(self, chat_id):
         return ChatContext(self.get_chat_messages(chat_id))
 
+    def get_chat_view(self, chat_id):
+        return ChatView(self, chat_id)
+
     def get_agent_control(self, agent_id):
         with self.sql_session() as session:
             agent = session.execute(
@@ -310,7 +331,8 @@ class App:
                 'time': message.time,
                 'sender': message.sender,
                 'kind': message.kind,
-                'content': message.content
+                'content': message.content,
+                'threadId': message.thread_id,
             }
             for channel in self.user_connections[user_id]:
                 asyncio.create_task(channel.put(user_message))
@@ -337,6 +359,45 @@ class App:
             picture=None,
             members=[agent_id],
         )
+    
+    def create_chat_thread(self, chat_id, name=None):
+        with self.sql_session() as session:
+            thread = db.Thread()
+            thread.chat_id = chat_id
+            thread.name = name
+
+            session.add(thread)
+            session.commit()
+            print(f"Created thread #{thread.id}")
+
+            return thread.id
+    
+    def get_agent_resource(self, agent_id, kind, key):
+        with self.sql_session() as session:
+            stored_value = session.execute(select(db.AgentResource.value).where(
+                db.AgentResource.agent_id == agent_id and db.AgentResource.kind == kind and db.AgentResource.key == key
+            )).scalar()
+            return stored_value
+
+    def save_agent_resource(self, agent_id, kind, key, value):
+        with self.sql_session() as session:
+            existing_resource = session.execute(select(db.AgentResource).where(
+                db.AgentResource.agent_id == agent_id and db.AgentResource.kind == kind and db.AgentResource.key == key
+            )).scalars().first()
+
+            if existing_resource:
+                existing_resource.value = value
+            else:
+                new_resource = db.AgentResource()
+                new_resource.agent_id = agent_id
+                new_resource.kind = kind
+                new_resource.key = key
+                new_resource.value = value
+
+                session.add(new_resource)
+                session.commit()
+                print(f"Saved agent resource: {agent_id} - {kind} - {key} = {value}")
+
 
 if __name__ == "__main__":
     if len(sys.argv) > 1:

backend/chat.py

@@ -0,0 +1,130 @@
+
+from sqlalchemy.orm import Session
+from sqlalchemy.sql.expression import select
+
+import database as db
+
+class Message:
+    def __init__(self, id, content, sender, time=None):
+        self.id = id
+        self.content = content
+        self.sender = sender
+        self.time = time
+    
+    def from_db_record(record):
+        return Message(record.id, record.content, record.sender, record.time)
+    
+    def __str__(self):
+        return f"Message(#{self.id} ({self.time if self.time != None else 'no-time'}) {self.sender}: {self.content.replace("\n", "\\n")})"
+
+class Thread:
+    def __init__(self, id, chat_id, name=None, description=None):
+        self.id = id
+        self.chat_id = chat_id
+        self.name = name
+        self.description = description
+
+        self.messages = []
+
+class ChatView:
+    def __init__(self, app, chat_id):
+        self.app = app
+        self.chat_id = chat_id
+
+        self.name = None
+        self.description = None
+
+        # Thread holding messages without a set thread id
+        self.default_thread = Thread(None, chat_id)
+        # Most recent thread objects, earliest to latest
+        self.threads = []
+        # Map of threads by their id
+        self.thread_map = {}
+        # Long term summary of this chat
+        self.summary = None
+
+        # Id of the thread currently in focus, by default the last thread appearing in the conversation
+        self.current_thread_id = None
+        
+        self._earliest_message_id = None
+
+        self._load_basic_data()
+        self._load_next_messages(20)
+    
+    def _load_basic_data(self):
+        with self.app.sql_session() as session:
+            record = session.execute(
+                select(db.Chat).where(db.Chat.id == self.chat_id)
+            ).scalars().first()
+
+            if not record:
+                raise Exception(f"Chat not found with id {self.chat_id}")
+            
+            self.name = record.name
+            self.description = record.description
+
+    def _load_next_messages(self, load_count):
+        with self.app.sql_session() as session:
+            records = session.execute(
+                select(db.Message).where(
+                    db.Message.chat_id == self.chat_id,
+                    self._earliest_message_id == None or db.Message.id <= self._earliest_message_id
+                ).order_by(db.Message.time.asc())
+            ).scalars().fetchmany(load_count)
+
+            print(f"Loaded messages. Threads: {bool(self.threads)}, records: {bool(records)}")
+            if not self.threads and records:
+                print(f"Last Record: ", records[-1])
+                self.current_thread_id = records[-1].thread_id
+
+            # Go over messages from newest to oldest
+            for record in reversed(records):
+                if record.thread_id == None:
+                    thread = self.default_thread
+                else:
+                    thread = self.thread_map.get(record.thread_id)
+                    if not thread:
+                        thread = Thread(record.thread_id, self.chat_id)
+                        self.thread_map[thread.id] = thread
+                        self.threads.insert(0, thread)
+                # This is the oldest message processed, add it to the start of the chat
+                message = Message(record.id, record.content, record.sender, record.time)
+                thread.messages.insert(0, message)
+            
+            # Update the earliest id for the next load operation
+            if records:
+                self._earliest_message_id = records[0].id
+    
+    def get_thread(self, thread_id):
+        if thread_id == None:
+            return self.default_thread
+        return self.thread_map.get(thread_id)
+    
+    def get_current_thread(self):
+        # TODO: if !thread
+        return self.get_thread(self.current_thread_id)
+    
+    def debug_string(self):
+        result_lines = []
+        def _print(line):
+            result_lines.append(line)
+
+        def print_thread(thread, title=None):
+            if title:
+                _print(title + ":")
+            elif thread.name:
+                _print(thread.name + ":")
+            else:
+                _print(f"Thread {thread.id}:")
+            for message in thread.messages:
+                _print("  " + str(message))
+        
+        print_thread(self.default_thread, "Default Thread")
+        for thread in self.threads:
+            print_thread(thread)
+
+        _print("Current thread id: " + str(self.current_thread_id))
+
+        return "\n".join(result_lines)
+
+

backend/database.py

@@ -89,3 +89,45 @@ class Message(Base):
     # Tells whether this message continues in the chat after the session when it was sent.
     persistent: Mapped[bool] = mapped_column(default=False)
 
+    thread_id: Mapped[Optional[int]] = mapped_column(ForeignKey("threads.id"))
+
+# Represents a temporary topic or stage in a conversation
+class Thread(Base):
+    __modelname__ = 'thread'
+    __tablename__ = 'threads'
+
+    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
+    chat_id: Mapped[int] = mapped_column(ForeignKey("chats.id"))
+
+    # Optional name and descriptions of the thread
+    name: Mapped[Optional[str]]
+    description: Mapped[Optional[str]]
+
+# A persistent object or value associated with an agent
+class AgentResource(Base):
+    __modelname__ = 'agent_resource'
+    __tablename__ = 'agent_resources'
+
+    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
+    agent_id: Mapped[int] = mapped_column(ForeignKey("agents.id"))
+
+    # Kind is the class of the resource
+    # Key is the specific identifier within the class
+    # Multiple resources can share key
+    kind: Mapped[str]
+    key: Mapped[str]
+    value = mapped_column(JSON, default=None)
+
+# A log entry of an agent, the use and meaning depends on the agent class
+class AgentLog(Base):
+    __modelname__ = 'agent_log'
+    __tablename__ = 'agent_logs'
+
+    id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
+    agent_id: Mapped[int] = mapped_column(ForeignKey("agents.id"))
+
+    # Kind is the class of the resource
+    # Key is the specific identifier within the class
+    # Multiple resources can share key
+    value: Mapped[str]
+

backend/server.py

@@ -180,9 +180,17 @@ def start_server(app):
                     message = app.send_chat_message(
                         chat_id=body['chatId'],
                         content=body['content'],
+                        thread_id=body.get('threadId'),
                     )
                     await send({'message': "message sent", 'messageId': message['id'], 'nonce': nonce})
 
+                elif action == "create-chat-thread":
+                    thread_id = app.create_chat_thread(
+                        chat_id=body['chatId'],
+                        name=body.get('name'),
+                    )
+                    await send({'message': "chat thread created", 'threadId': thread_id, 'nonce': nonce})
+
                 elif action == "send-agent-prompt":
                     response = await app.send_agent_prompt(body['agentId'], body['content'])
                     await send({'message': "propmt sent", 'response': response, 'nonce': nonce})

frontend/lib/chat_screen.dart

@@ -5,18 +5,32 @@ import 'styles.dart';
 import 'backend.dart';
 import 'dart:async';
 
-class Message {
-  final int? id;
-  final String sender;
-  final String text;
-  final bool isFromUser;
+abstract class ChatEntry {
+  String content;
+  bool isFromUser;
 
-  Message({required this.sender, required this.text, this.id}) : isFromUser = sender == "user";
+  ChatEntry({required this.content, required this.isFromUser});
 
   Alignment get alignment => isFromUser ? Alignment.centerRight : Alignment.centerLeft;
   EdgeInsets get padding => isFromUser ? EdgeInsets.zero : const EdgeInsets.only(left: 10);
 }
 
+class Message extends ChatEntry {
+  int? id;
+  String sender;
+  int? threadId;
+
+  Message({required this.sender, required content, this.id, this.threadId}) : super(content: content, isFromUser: sender == "user");
+}
+
+class ThreadHeader extends ChatEntry {
+  final int threadId;
+  final int messageCount;
+  
+  ThreadHeader({required this.threadId, required this.messageCount})
+    : super(content: "Thread ($messageCount messages)", isFromUser: false);
+}
+
 class ChatScreen extends StatefulWidget {
   final int chatId;
 
@@ -25,7 +39,9 @@ class ChatScreen extends StatefulWidget {
   static Future<List<Message>> getChatMessages(int chatId) async {
     final response = await Backend.sendAction("get-chat-messages", {'chatId': chatId});
     List<dynamic> messageObjects = response['messages'];
-    List<Message> messages = messageObjects.map((msg) => Message(sender: msg['sender'], text: msg['content'])).toList();
+    List<Message> messages = messageObjects.map((msg) =>
+      Message(sender: msg['sender'], content: msg['content'], id: msg['id'], threadId: msg['threadId'])
+    ).toList();
     return messages;
   }
 
@@ -34,8 +50,15 @@ class ChatScreen extends StatefulWidget {
 }
 
 class _ChatScreenState extends State<ChatScreen> {
+  // Chat State
   String? _title;
   List<Message> _messages = [];
+  int? currentThreadId;
+
+  // UI State
+  Set<int> expandedThreads = {};
+
+  // Internal Funcionality
   Offset? _longPressOffset;
   StreamSubscription? _chatSubscription;
   final TextEditingController _textController = TextEditingController();
@@ -46,7 +69,7 @@ class _ChatScreenState extends State<ChatScreen> {
     super.initState();
     _loadInitialData();
     _chatSubscription = Backend.chatMessages(widget.chatId).listen((data) {
-      final message = Message(sender: data['sender'], text: data['content']);
+      final message = Message(sender: data['sender'], content: data['content'], threadId: data['threadId']);
       if (!message.isFromUser) {
         addMessage(message);
       }
@@ -70,9 +93,58 @@ class _ChatScreenState extends State<ChatScreen> {
     setState(() {
       _messages.addAll(initialMessages);
       _title = chatInfo.name;
+
+      currentThreadId = initialMessages.last.threadId;
+      if (currentThreadId != null) {
+        expandedThreads.add(currentThreadId!);
+      }
     });
   }
 
+  List<ChatEntry> get displayEntries {
+    List<ChatEntry> result = [];
+    Map<int, List<Message>> threadGroups = {};
+
+    // First pass: Group messages by thread
+    for (var message in _messages) {
+      if (message.threadId != null) {
+        threadGroups.putIfAbsent(message.threadId!, () => []).add(message);
+      }
+    }
+
+    // Second pass: Build the display list with correct ordering
+    for (var message in _messages) {
+      if (message.threadId == null) {
+        // Non-threaded messages go directly to the result
+        result.add(message);
+      } else {
+        final threadId = message.threadId!;
+        final threadMessages = threadGroups[threadId]!;
+        
+        // If expanded, add messages as they appear
+        if (expandedThreads.contains(threadId)) {
+          // Add the thread header before the first message
+          if (message == threadMessages.first) {
+            result.add(ThreadHeader(threadId: threadId, messageCount: threadMessages.length));
+          }
+          result.add(message);
+        } else {
+          // Add collapsed thread header in place of the last message
+          if (message == threadMessages.last) {
+            result.add(ThreadHeader(threadId: threadId, messageCount: threadMessages.length));
+          }
+        }
+      }
+    }
+
+    // Add current thread mannually if it has no messages
+    if (currentThreadId != null && !threadGroups.containsKey(currentThreadId!)) {
+      result.add(ThreadHeader(threadId: currentThreadId!, messageCount: 0));
+    }
+
+    return result;
+  }
+
   void addMessage(Message message) {
     setState(() {
       _messages.add(message);
@@ -82,8 +154,12 @@ class _ChatScreenState extends State<ChatScreen> {
   void sendMessage(String content) {
     if (content.trim().isEmpty) return; // Don't send empty messages
     _textController.clear();
-    final newMessage = Message(sender: "user", text: content);
-    Backend.sendAction("send-chat-message", {'chatId': widget.chatId, 'content': content}).then((response) {
+    final newMessage = Message(sender: "user", content: content, threadId: currentThreadId);
+    Backend.sendAction("send-chat-message", {
+      'chatId': widget.chatId,
+      'content': content,
+      'threadId': currentThreadId,
+    }).then((response) {
       if (response['error'] != null) {
         print("Error sending chat message: ${response['error']}");
       }
@@ -91,6 +167,16 @@ class _ChatScreenState extends State<ChatScreen> {
     addMessage(newMessage);
   }
 
+  void createThread () async {
+    final response = await Backend.sendAction("create-chat-thread", {'chatId': widget.chatId});
+    if (response['threadId'] != null) {
+      setState(() {
+        currentThreadId = response['threadId'];
+        expandedThreads.add(currentThreadId!);
+      });
+    }
+  }
+
   Color getMessageColor(BuildContext context) {
     final backgroundColor = Theme.of(context).scaffoldBackgroundColor;
     final messageColor = Color.alphaBlend(Colors.white.withOpacity(0.1), backgroundColor);
@@ -116,6 +202,12 @@ class _ChatScreenState extends State<ChatScreen> {
             print('Edit clicked');
           },
         ),
+        PopupMenuItem(
+          child: Text('Start Topic'),
+          onTap: () {
+            print('TODO: Start Topic from message');
+          },
+        ),
         PopupMenuItem(
           child: Text('Delete'),
           onTap: () {
@@ -126,6 +218,36 @@ class _ChatScreenState extends State<ChatScreen> {
       ],
     );
   }
+  
+  Widget _buildThreadHeader(BuildContext context, ThreadHeader header) {
+    return GestureDetector(
+      onTap: () {
+        setState(() {
+          if (expandedThreads.contains(header.threadId)) {
+            expandedThreads.remove(header.threadId);
+          } else {
+            expandedThreads.add(header.threadId);
+          }
+        });
+      },
+      child: Container(
+        padding: EdgeInsets.all(8),
+        margin: EdgeInsets.symmetric(vertical: 4),
+        decoration: BoxDecoration(
+          color: getMessageColor(context),
+          borderRadius: BorderRadius.circular(4),
+        ),
+        child: Row(
+          children: [
+            Icon(expandedThreads.contains(header.threadId) 
+              ? Icons.expand_less 
+              : Icons.expand_more),
+            Text(header.content),
+          ],
+        ),
+      ),
+    );
+  }
 
   @override
   Widget build(BuildContext context) {
@@ -144,42 +266,51 @@ class _ChatScreenState extends State<ChatScreen> {
         children: [
           Expanded(
             child: ListView.builder(
-              itemCount: _messages.length,
+              itemCount: displayEntries.length,
               itemBuilder: (context, index) {
-                final message = _messages[index];
-                return GestureDetector(
-                  onLongPressStart: (details) {
-                    _longPressOffset = details.globalPosition;
-                  },
-                  onLongPress: () {
-                    _showMessageContextMenu(context, _longPressOffset!);
-                    _longPressOffset = null;
-                  },
-                  onSecondaryTapUp: (details) {
-                    _showMessageContextMenu(context, details.globalPosition);
-                  },
-                  child: Padding(
-                    padding: isWideScreen ? message.padding : EdgeInsets.zero,
-                    child: Align(
-                      alignment: isWideScreen ? Alignment.centerLeft : message.alignment,
-                      child: Container(
-                        margin: EdgeInsets.symmetric(vertical: 2, horizontal: 8),
-                        padding: EdgeInsets.all(10),
-                        decoration: BoxDecoration(
-                          color: getMessageColor(context),
-                          borderRadius: BorderRadius.circular(12),
+                final entry = displayEntries[index];
+                
+                if (entry is ThreadHeader) {
+                  return _buildThreadHeader(context, entry);
+                } else if (entry is Message) {
+                  final message = entry as Message;
+                  return GestureDetector(
+                    onLongPressStart: (details) {
+                      _longPressOffset = details.globalPosition;
+                    },
+                    onLongPress: () {
+                      _showMessageContextMenu(context, _longPressOffset!);
+                      _longPressOffset = null;
+                    },
+                    onSecondaryTapUp: (details) {
+                      _showMessageContextMenu(context, details.globalPosition);
+                    },
+                    child: Padding(
+                      padding: isWideScreen ? message.padding : EdgeInsets.zero,
+                      child: Align(
+                        alignment: isWideScreen ? Alignment.centerLeft : message.alignment,
+                        child: Container(
+                          margin: EdgeInsets.symmetric(vertical: 2, horizontal: 8),
+                          padding: EdgeInsets.all(10),
+                          decoration: BoxDecoration(
+                            color: getMessageColor(context),
+                            borderRadius: BorderRadius.circular(12),
+                          ),
+                          child: Text(message.content),
                         ),
-                        child: Text(message.text),
                       ),
                     ),
-                  ),
-                );
+                  );
+                } else {
+                  // Unknown entry type
+                  return Text(entry.content);
+                }
               },
             ),
           ),
           Padding(
             padding: const EdgeInsets.all(8.0),
-            child: Row(
+            child: Column(
               children: [
                 Expanded(
                   child: KeyboardListener(
@@ -200,12 +331,22 @@ class _ChatScreenState extends State<ChatScreen> {
                     ),
                   ),
                 ),
-                SizedBox(width: 8),
-                ElevatedButton(
-                  onPressed: () {
-                    sendMessage(_textController.text);
-                  },
-                  child: Text('Send'),
+                Row(
+                  children: [
+                    ElevatedButton(
+                      onPressed: () {
+                        sendMessage(_textController.text);
+                      },
+                      child: Text('Send'),
+                    ),
+                    ElevatedButton(
+                      onPressed: () {
+                        createThread();
+                        print("TODO: Start new topic");
+                      },
+                      child: Text('Start New Topic'),
+                    ),
+                  ],
                 ),
               ],
             ),