From e9ebb1378a6bbf909b75c6a39dcc02e774dac6c3 Mon Sep 17 00:00:00 2001 From: Shivering Ocelot Date: Sat, 28 Jun 2025 00:59:56 -0500 Subject: [PATCH] Full memory storing and retrieval for conversation agent. Update gral prompts. Add custom inference model. --- backend/agents/conversation.py | 165 +++++++++++++++++++++++++++++---- backend/app.py | 24 +++-- backend/chat.py | 24 ++--- backend/gral.py | 19 ++-- backend/inference.py | 14 ++- 5 files changed, 196 insertions(+), 50 deletions(-) diff --git a/backend/agents/conversation.py b/backend/agents/conversation.py index 19b7835..3637671 100644 --- a/backend/agents/conversation.py +++ b/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) \ No newline at end of file + 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}") \ No newline at end of file diff --git a/backend/app.py b/backend/app.py index 0152328..abb5429 100644 --- a/backend/app.py +++ b/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) @@ -402,8 +406,8 @@ class App: new_resource.value = value session.add(new_resource) - session.commit() - print(f"Saved agent resource: {agent_id} - {kind} - {key} = {value}") + session.commit() + print(f"Saved agent resource: {agent_id} - {kind} - {key} = {value}") if __name__ == "__main__": diff --git a/backend/chat.py b/backend/chat.py index 9c343ce..493e747 100644 --- a/backend/chat.py +++ b/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: @@ -90,10 +94,6 @@ class ChatView: # This is the oldest message so far, 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: diff --git a/backend/gral.py b/backend/gral.py index 44338c2..5cae626 100644 --- a/backend/gral.py +++ b/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, ] diff --git a/backend/inference.py b/backend/inference.py index 27293ef..3b641bd 100644 --- a/backend/inference.py +++ b/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, }