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}")