Full memory storing and retrieval for conversation agent. Update gral prompts. Add custom inference model.

backend/agents/conversation.py

@@ -1,9 +1,15 @@
 import re
 import json
 import inspect
+from string import Template
 
 from inference import LlamaMessage, replicate_llama_completion
-from agent import Agent
+from agent import Agent, AgentField
+
+chat_system_prompt_template = Template("""$role
+The first message contains information about you and a conversation that you're participating in.
+Respond only with your next message in the conversation.
+""".strip())
 
 memory_system_prompt = """
 You are an expert in simulating human-like memory for fictional characters. Given a conversation between two characters and a detailed description of one character's personality, extract the memories that this character is most likely to retain.
@@ -13,7 +19,35 @@ The extracted memories should reflect:
 - What they found surprising, confusing, or meaningful
 - What they might revisit or ruminate on later
 Do not summarize the conversation. Instead, highlight key moments or phrases that the character would mentally hold onto, framed from their perspective. Avoid information stated by yourself, you don't need to remember what you already know, instead you can identify the other's characters reactions to what your character says if it's notable or salient.
-"""
+""".strip()
+
+topic_system_prompt = """
+You analyze a conversation between two characters. Identify the main topic as a short phrase. Then extract up to N subtopics that are likely to be important, related, or may emerge later in the conversation.
+Return your response as a JSON object with this format:
+{"topic": "Main topic", "subtopics": ["First subtopic", "Second subtopic", "..."]}
+""".strip()
+
+memory_compress_system_prompt = """
+You are an assistant that receives a list of memory items and a current conversation; you must identify and return only the most important and relevant memories that help understand the conversation or move it forward, removing duplicates or very similar entries.
+Output a JSON list of strings with each selected memory, for example:
+["First memory","Second memory"]
+""".strip()
+
+chat_prompt_template = Template("""
+Recent Memories:
+$memories
+
+Recent Conversation:
+$conversation
+""".strip())
+
+topic_prompt_template = Template("""
+# Parameters
+Subtopic count: $subtopic_count
+
+# Conversation History
+$conversation
+""".strip())
 
 memory_prompt_template = """
 # Your Character
@@ -28,7 +62,17 @@ If there are no relevant memories to extract, it is ok to respond simply with an
 
 # Conversation History
 [conversation]
-"""
+""".strip()
+
+memory_compress_prompt_template = Template("""
+# Memories
+Following are all the known memories. Each memory is in its own line.
+
+$memories
+
+# Conversaton
+$conversation
+""".strip())
 
 # Used to keep track of the last chat thread it participated in
 class ThreadContext:
@@ -40,14 +84,16 @@ class ThreadContext:
 
 class ConversationAgent(Agent):
     identifier = "conversation"
-    schema_fields = []
+    schema_fields = [
+        AgentField('system_prompt', 'string', default="You are a friendly AI having a casual conversation with the user."),
+    ]
     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."
+        self.system_prompt = self.config['system_prompt']
 
     async def prompt(self, content):
         # TODO: Advanced features
@@ -66,23 +112,29 @@ class ConversationAgent(Agent):
         chat_view = self.app.get_chat_view(message.chat_id)
         print("ChatView:\n" + chat_view.debug_string())
 
+        print("Saved context thread:", saved_context.current_thread_id)
+        print("Current thread:", message.thread_id)
+        print("Message", message.content)
+
         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.get_thread(saved_context.current_thread_id)
             await self.process_thread(old_thread)
             
-            # TODO: This is a good place to store the new context, its fine that the agent hasnt responded
+            new_context_data = {'current_thread_id': message.thread_id}
+            self.app.save_agent_resource(self.id, 'last-thread-context', message.chat_id, new_context_data)
 
         current_thread = chat_view.get_thread(message.thread_id)
-
-        # REMOVE THIS. Only for testing, while the function only prints...
-        await self.process_thread(current_thread)
+        memories = await self.retrieve_memories_for_thread(current_thread)
 
         chat_summary = self.serialize_thread_history(current_thread)
         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"Conversation History:\n{chat_summary}")
+            LlamaMessage("system", chat_system_prompt_template.substitute(role=self.system_prompt)),
+            LlamaMessage("user", chat_prompt_template.substitute(
+                conversation=chat_summary,
+                memories="\n".join(memories),
+            )),
         ]
 
         response = await replicate_llama_completion(self.app, messages)
@@ -105,6 +157,81 @@ class ConversationAgent(Agent):
         
         return "\n".join(f"[{sender_name(message)}]: {message.content}" for message in thread.messages)
     
+    def get_memory_collection(self):
+        collection_name = f"agent_{self.id}_memory"
+        return collection_name
+    
+    async def retrieve_memories_for_thread(self, thread):
+        print("identify topic")
+        chat_history = self.serialize_thread_history(thread)
+
+        topic_response = await replicate_llama_completion(self.app, [
+            LlamaMessage("system", topic_system_prompt),
+            LlamaMessage("user", topic_prompt_template.substitute(
+                conversation=chat_history,
+                subtopic_count=3
+            ))
+        ])
+
+        match = re.search(r'\{.*?\}', topic_response, re.DOTALL)
+        if match:
+            json_string = match.group(0)
+            try:
+                topic_data = json.loads(json_string)
+            except json.JSONDecodeError:
+                print("Bad JSON data in response...")
+                print(topic_response)
+                return []
+        else:
+            print("Topic format not found in response...")
+            print(topic_response)
+            return []
+
+        print("retrieve memories")
+        client = self.app.qdrant_client
+        collection = self.get_memory_collection()
+
+        if not client.collection_exists(collection):
+            return []
+        
+        query_result = client.query(
+            collection,
+            query_text=topic_data['topic'],
+            limit=20
+        )
+        for subtopic in topic_data['subtopics']:
+            query_result += client.query(
+                collection,
+                query_text=subtopic,
+                limit=6
+            )        
+        memories = [it.document for it in query_result]
+
+        print("compress memories")
+        compress_response = await replicate_llama_completion(self.app, [
+            LlamaMessage("system", memory_compress_system_prompt),
+            LlamaMessage("user", memory_compress_prompt_template.substitute(
+                conversation=chat_history,
+                memories="\n".join(memories)
+            ))
+        ])
+
+        match = re.search(r'\[.*?\]', compress_response, re.DOTALL)
+        if match:
+            json_string = match.group(0)
+            try:
+                memories = json.loads(json_string)
+            except json.JSONDecodeError:
+                print("Bad JSON data in response...")
+                print(compress_response)
+                return []
+        else:
+            print("Compress format not found in response...")
+            print(compress_response)
+            return []
+        
+        return memories
+    
     # Compute thread summnary and concrete memories, then save the progress done
     async def process_thread(self, thread):
         # TODO: extract memories from the thread messages, then store them in the agent's vector memory db
@@ -125,14 +252,20 @@ class ConversationAgent(Agent):
             json_string = match.group(0)
             try:
                 memories = json.loads(json_string)
-                print("Extracted Memories:", memories)
             except json.JSONDecodeError:
                 print("Bad JSON data in response...")
                 print(response)
+                return
         else:
             print("Memory format not found in response...")
             print(response)
+            return
 
-        memories = [line for line in response.split('\n') if line.strip()]
-        
-        print("New agent memories:", memories)
+        collection = self.get_memory_collection()
+        if memories:
+            self.app.qdrant_client.add(
+                collection,
+                documents=[it['memory'] for it in memories],
+                metadata=[{'importance': it['importance']} for it in memories],
+            )
+            print(f"Added {len(memories)} memories for agent {self.id}")

backend/app.py

@@ -13,12 +13,10 @@ import replicate
 import database as db
 import agent as agent_mod
 from server import start_server
-from chat import ChatView, Message
+from chat import ChatView, Message, MESSAGE_PAGE_SIZE
 
 from agents.conversation import ConversationAgent
 
-MESSAGE_PAGE_SIZE = 32
-
 class ChatContext:
     def __init__(self, messages, prelude=None):
         self.messages=messages
@@ -170,11 +168,17 @@ class App:
             session.commit()
             print(f"Deleted chat #{chat_id} and its memberships.")
 
-    def get_chat_messages(self, chat_id, earliest_index=0):
+    def get_chat_messages(self, chat_id, earliest_index=None):
         with self.sql_session() as session:
+            condition = (db.Message.chat_id == chat_id)
+            if earliest_index is not None:
+                condition = condition and db.Message.id < earliest_index
+
             messages = session.execute(
-                select(db.Message).where(db.Message.chat_id == chat_id).order_by(db.Message.time.asc())
-            ).scalars().fetchmany(MESSAGE_PAGE_SIZE)
+                select(db.Message).where(condition).order_by(db.Message.time.desc())
+            ).scalars().all()
+
+            messages.reverse()
 
             return [{
                 "id": message.id,
@@ -274,8 +278,8 @@ class App:
             'content': message.content,
         }
 
-    def get_chat_context(self, chat_id):
-        return ChatContext(self.get_chat_messages(chat_id))
+    def get_chat_context(self, chat_id, earliest_index=None):
+        return ChatContext(self.get_chat_messages(chat_id, earliest_index))
 
     def get_chat_view(self, chat_id):
         return ChatView(self, chat_id)

backend/chat.py

@@ -4,6 +4,8 @@ from sqlalchemy.sql.expression import select
 
 import database as db
 
+MESSAGE_PAGE_SIZE = 256
+
 class Message:
     def __init__(self, id, content, sender, time=None):
         self.id = id
@@ -49,7 +51,7 @@ class ChatView:
         self._earliest_message_id = None
 
         self._load_basic_data()
-        self._load_next_messages(20)
+        self._load_next_messages(MESSAGE_PAGE_SIZE)
     
     def _load_basic_data(self):
         with self.app.sql_session() as session:
@@ -65,20 +67,22 @@ class ChatView:
 
     def _load_next_messages(self, load_count):
         with self.app.sql_session() as session:
-            records = session.execute(
+            records_rev = 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())
+                ).order_by(db.Message.time.desc())
             ).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
+            # Update the earliest id for the next load operation
+            if records_rev:
+                self._earliest_message_id = records_rev[0].id
+                if not self.threads:
+                    # current thread is the latest known
+                    self.current_thread_id = records_rev[-1].thread_id
 
             # Go over messages from newest to oldest
-            for record in reversed(records):
+            for record in records_rev:
                 if record.thread_id == None:
                     thread = self.default_thread
                 else:
@@ -91,10 +95,6 @@ class ChatView:
                 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

backend/gral.py

@@ -1,14 +1,15 @@
-from inference import LlamaMessage, replicate_llama_completion
+from inference import LlamaMessage, replicate_llama_completion, Models
 import re
 
 message_format_section = """
 When suggesting code to the user, keep the code as small snippets.
-If the update is meant to be applied to a file, include a line linking the file immediately before the code block like this: "[filename.txt]
+You have the capability to update files in your current environment with file tags before your code suggestions.
+If the update is meant to be applied on a file, include a file tag line immediately before the code block like this: "[filename.txt]
 ```
 updated segment
 ```"
-Omit the file link when the block is not related to a known file
-Code blocks must be under 30 lines. Assistant responses MUST NOT contain horizontal rules.
+Omit the file link when the block is not related to a known file,
+Code blocks must be under 30 lines. Your responses MUST NOT contain horizontal rules.
 When suggesting bash code meant to be directly executed, use the same syntax as above but marking the code block with "[eval]" instead.
 """
 
@@ -27,9 +28,9 @@ async def gral_message(app, system_prompt, message_history: list, file_map: dict
     
     if file_parts:
         file_lines = [
-            "This message contains all files relevant to the prompt: each file section in this message starts with a horizontal rule, followed by the file name in square brackets, followed by the file content.",
-            "Only this message contains horizontal rules, assistant responses must not include them.",
-            "The prompt and conversation start in the next message after this."
+            "This message contains all files relevant to the prompt: each file in this message is separated with horizontal rules \"---\", followed by the file name in square brackets, followed by the file content.",
+            "Only this message contains horizontal rules, your responses never has them.",
+            "Conversation start in the message following this one."
         ]
         file_lines.extend(file_parts)
         messages.append(
@@ -38,7 +39,7 @@ async def gral_message(app, system_prompt, message_history: list, file_map: dict
 
     messages.extend([LlamaMessage(msg['role'], msg['content']) for msg in message_history])
 
-    response = await replicate_llama_completion(app, messages)
+    response = await replicate_llama_completion(app, messages, Models.llama_70b)
     return response
 
 async def gral_update_file(app, system_prompt, message_history: list, filepath: str, file_content: str):
@@ -57,7 +58,7 @@ async def gral_update_file(app, system_prompt, message_history: list, filepath:
     # Main Instruction Message
     message_lines = [
         f"Write the updated contents of the file: {filepath}, as concluded in the conversation above.",
-        f"The raw file contents of the file {filepath} are included following the horizontal rule below."
+        f"The current content of the file {filepath} is pasted after the horizontal rule below."
         "\n---\n",
         file_content,
     ]

backend/inference.py

@@ -11,8 +11,16 @@ class LlamaMessage:
         self.role = role
         self.content = content
 
-async def replicate_llama_completion(app, messages):
-    # Llama 3 instruct
+
+class Models:
+    llama_8b = "meta/meta-llama-3-8b-instruct"
+    llama_70b = "meta/meta-llama-3-70b-instruct"
+
+async def replicate_llama_completion(app, messages, model=None):
+    if model is None:
+        model = Models.llama_8b
+
+    # Llama instruct
     full_prompt = "<|begin_of_text|>"
     for message in messages:
         header_name = message.role if message.role != "agent" else "assistant"
@@ -24,7 +32,7 @@ async def replicate_llama_completion(app, messages):
         target=run_replicate_inference_loop,
         kwargs={
             'replicate_client': app.replicate_client,
-            'model': LLAMA3_MODEL,
+            'model': model,
             'prompt': full_prompt,
             'content_stream': content_stream,
         }