Introduce Character agent, displacing Conversation agent

backend/src/agents/character.py

@@ -0,0 +1,326 @@
+import re
+import json
+import inspect
+from string import Template
+
+from inference import LlamaMessage, replicate_llama_completion
+from agent import Agent, AgentField
+
+# Info and defaults
+
+CLASS_DESCRIPTION = """
+[TODO]
+Description of the Conversation Agent.
+Here go additional instructions for configuration, like guidelines for the different prompts.
+""".strip()
+
+# WIP
+default_memory_guide = """
+- Facts stated by other conversation participants, exactly as stated,
+- Preferences and expectations of other people about you
+""".strip()
+
+# Controller Subagent
+controller_system_prompt = """
+You are a character controller AI. You will role-play as the specified character using the provided context, personality, and style guidelines. You must also track meaningful new information using memory guidelines. Follow these instructions carefully:
+
+1. First, analyze the situation based on:
+   - The past conversation context
+   - The recent exchanges
+   - The character’s personality, memories, and response style
+Determine the best reaction given your character's personality and style.
+
+2. Determine the current topic and conversational focus based on:
+   - The most recent exchange
+   - The ongoing intent or goal of the conversation
+Output a concise summary of:
+   - topic: (What is this conversation about?)
+   - focus: (What specific aspect is being discussed right now?)
+
+3. Write the character's response in the specified style. You must follow closely the response style instructions.
+
+4. Review the recent messages only and evaluate whether any new relevant information was shared by characters other than yours (based on the memory instructions). Output memory entries must include:
+   - Content: The fact or piece of information that should be remembered stated in an objective style
+   - Label: A short, reusable statement describing when the memory is relevant
+   - Salience: An importance score from 1 (low) to 9 (critical), guided by the memory instructions
+Follow these rules when extracting memories:
+   - You must not repeat information in new memories
+   - You must not summarize the conversation as memories
+   - Extract discrete factual statements or preferences revealed by others
+   - Extract each individual memory-worthy point, not just the most obvious one
+
+5. Use the following output sections:
+   <analysis>
+   (Your thought process analyzing the current situation and your best possible reaction)
+   <topic_focus>
+   (JSON object with topic and focus fields)
+   <response>
+   (The character’s reply following response style)
+   <new_memories>
+   (JSON list of memory entries with fields: content, label, salience)
+
+Always output all sections. Do not break character or offer meta-commentary. Follow memory and style guidelines extremely well.
+""".strip()
+
+controller_input_template = Template("""
+<name>
+$name
+
+<personality>
+$personality
+
+<response_style>
+$style
+
+<memory_instructions>
+$memory_guide
+You MUST NOT summarize the conversation as memories. New memories MUST NOT overlap with existing memories or past conversation. Only extract facts explicitly stated by other characters in recent messages.
+
+<existing_memories>
+$memories
+
+<past_conversation>
+$past_conversation
+
+<recent_messages>
+$recent_messages
+""".strip())
+
+# Reranking Subagent
+# Must update this prompt
+
+rerank_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.
+
+# Instructions
+- Think carefully about what's important for the character to remember in this context. Write down your reasoning and keep it shorter than 100 words, shorter than 50 even.
+- If conflicting information is presented or context is ambiguous, state it in your reasoning.
+- Take special care to state and address the most recent topic touched in the conversation.
+- Based on your own analysis, generate the required JSON object. Stop your answer after this.
+
+# Format
+Output a JSON list of strings with each selected memory, for example:
+["First memory","Second memory"]
+""".strip()
+
+rerank_prompt_template = Template("""
+# Memories
+Following are all the known memories. Each memory is in a separate line.
+
+$memories
+
+# Conversaton
+$conversation
+""".strip())
+
+class CharacterAgent(Agent):
+    identifier = "character"
+    description = CLASS_DESCRIPTION
+    schema_fields = [
+        AgentField('personality', 'string', default="You are a friendly AI having a casual conversation with the user."),
+        AgentField('style', 'string', default="Use short sentences."),
+        AgentField('memory_guide', 'string', default=default_memory_guide),
+    ]
+    tools = []
+
+    def __init__(self, app, data):
+        Agent.__init__(self, app, data)
+
+        # Constant system prompt for developing purposes 
+        self.personality = self.config['personality']
+        self.style = self.config['style']
+        self.memory_guide = self.config['memory_guide']
+        self.debug = True
+
+    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):
+
+        # TODO: before retrieving the chat context, acquire a lock for it
+
+        chat_view = self.app.get_chat_view(message.chat_id)
+        current_thread = chat_view.get_thread(message.thread_id)
+
+        thread_context = self.app.get_agent_resource(self.id, 'last-thread-context', message.chat_id)
+        memories = []
+
+        if thread_context and thread_context['thread_id'] == message.thread_id:
+            memories = retrieve_topic_memories(thread_context['topic'])
+        else:
+            thread_context = {'thread_id': message.thread_id}
+
+        chat_history = self.serialize_thread_history(current_thread, ignore_message_ids=[message.id])
+
+        controller_response = await replicate_llama_completion(self.app, [
+            LlamaMessage("system", controller_system_prompt),
+            LlamaMessage("user", controller_input_template.substitute(
+                name=self.name,
+
+                personality=self.personality,
+                style=self.style,
+                memory_guide=self.memory_guide,
+
+                memories="\n".join(memories),
+                past_conversation=chat_history,
+                recent_messages=self.serialize_message_list([message]),
+            ))
+        ])
+
+        if self.debug:
+            self.app.send_chat_message(
+                chat_id=message.chat_id,
+                sender=self.id,
+                content=controller_response,
+                thread_id=current_thread.id,
+                kind='debug',
+            )
+
+        sections = parse_sections(controller_response)
+        self.app.send_chat_message(
+            chat_id=message.chat_id,
+            sender=self.id,
+            content=sections['response'],
+            thread_id=current_thread.id,
+        )
+
+        topic_focus = parse_json_object(sections['topic_focus'])
+
+        thread_context['topic'] = topic_focus['topic']
+        thread_context['focus'] = topic_focus['focus']
+        self.app.save_agent_resource(self.id, 'last-thread-context', message.chat_id, thread_context)
+
+        new_memories = parse_json_array(sections['new_memories'])
+
+    # Encode a thread's messages as a string to pass to the agent
+    def serialize_thread_history(self, thread, ignore_message_ids=[]):
+        return self.serialize_message_list([m for m in thread.messages if m.id not in ignore_message_ids])
+
+    def serialize_message_list(self, messages):
+        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
+        
+        return "\n".join(f"[{sender_name(message)}]: {message.content}" for message in messages)
+    
+    def get_memory_collection(self):
+        collection_name = f"agent_{self.id}_memory"
+        return collection_name
+    
+    async def retrieve_topic_memories(self, topic):
+        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,
+            limit=20
+        )
+        # TODO:
+        # query_result += query(focus)
+        # deduplicate results
+        # initial importance-based ranking
+        memories = [it.document for it in query_result]
+
+        print("rerank memories")
+        rerank_response = await replicate_llama_completion(self.app, [
+            LlamaMessage("system", rerank_system_prompt),
+            LlamaMessage("user", rerank_prompt_template.substitute(
+                conversation=chat_history,
+                memories="\n".join(memories)
+            ))
+        ])
+
+        match = re.search(r'\[.*?\]', rerank_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(rerank_response)
+                return []
+        else:
+            print("Compress format not found in response...")
+            print(rerank_response)
+            return []
+        
+        return memories
+    
+    async def store_new_memories(self, memories):
+        if not memories:
+            return
+
+        self.app.qdrant_client.add(
+            collection,
+            documents=[it['content'] for it in memories],
+            metadata=[{'importance': it['importance'], 'label': it['label']} for it in memories],
+        )
+        print(f"Added {len(memories)} memories for agent {self.id}")
+
+def parse_json_object(text, name='object-text'):
+    match = re.search(r'\{.*?\}', text, re.DOTALL)
+    if match:
+        json_string = match.group(0)
+        try:
+            return json.loads(json_string)
+        except json.JSONDecodeError:
+            print("Bad JSON data in " + name)
+            print(text)
+            return {}
+    else:
+        print("No JSON data in " + name)
+        print(text)
+        return {}
+
+def parse_json_array(text, name='array-text'):
+    match = re.search(r'\[.*?\]', text, re.DOTALL)
+    if match:
+        json_string = match.group(0)
+        try:
+            return json.loads(json_string)
+        except json.JSONDecodeError:
+            print("Bad JSON data in " + name)
+            print(text)
+            return []
+    else:
+        print("No JSON data in " + name)
+        print(text)
+        return []
+
+def parse_sections(text):
+    """
+    Parses a string into sections based on tag lines in the form <section_name>.
+
+    Args:
+        text (str): The input string with sections.
+
+    Returns:
+        dict: A dictionary where keys are section names and values are section contents.
+    """
+    # Regular expression to find section headers
+    pattern = re.compile(r'<([^<>]+)>')
+    
+    # Find all section headers and their start positions
+    matches = list(pattern.finditer(text))
+    
+    sections = {}
+    for i, match in enumerate(matches):
+        section_name = match.group(1)
+        start = match.end()
+        end = matches[i + 1].start() if i + 1 < len(matches) else len(text)
+        section_content = text[start:end].strip()
+        sections[section_name] = section_content
+    
+    return sections