backend: implement agents and connect them with chats

backend/agent.py

@@ -0,0 +1,93 @@
+from stream import Stream
+import threading
+
+LLAMA3_MODEL = "meta/meta-llama-3-8b-instruct"
+
+# Instances of this class lives very briefly, created when needed and immediately disposed.
+# Subclasses must be careful in not setting up long term resources
+class AgentControl:
+    def __init__(self, app, agent_id):
+        self.app = app
+        self.agent_id = agent_id
+        self.system_prompt = "You are a useful and precise AI assistant. Respond concisely."
+
+    # Receives a single string prompt, and must respond with a string
+    async def prompt(self, content):
+        response = await replicate_llama_completion(self.app, [
+            LlamaMessage("system", self.system_prompt),
+            LlamaMessage("user", content)
+        ])
+        return response
+
+    async def received_message(self, message):
+        chat_context = self.app.get_chat_context(message.chat_id)
+        def sender_name(message):
+            sender = message['sender']
+            is_digit = isinstance(sender, (int, str)) and str(sender).isdigit()
+            return "assistant" if is_digit else sender
+        chat_summary = "\n".join(f"**{sender_name(message)}**: {message['content']}" for message in chat_context.messages)
+
+        prologue = "\nYou are given the entire context of the conversation, respond with your next message in the conversation."
+        print(chat_summary)
+
+        response = await replicate_llama_completion(self.app, [
+            LlamaMessage("system", self.system_prompt + prologue),
+            LlamaMessage("user", chat_summary)
+        ])
+
+        self.app.send_chat_message(
+            chat_id=message.chat_id,
+            sender=self.agent_id,
+            content=response,
+        )
+    
+    def create_from_data(app, agent_id, kind, data):
+        return AgentControl(app, agent_id)
+
+class LlamaMessage:
+    def __init__(self, role, content):
+        self.role = role
+        self.content = content
+
+async def replicate_llama_completion(app, messages):
+    # Llama 3 instruct
+    full_prompt = "<|begin_of_text|>"
+    for message in messages:
+        header_name = message.role if message.role != "agent" else "assistant"
+        full_prompt += f"<|start_header_id|>{header_name}<|end_header_id|>\n\n{message.content}<|eot_id|>"
+
+    full_prompt += f"<|start_header_id|>assistant<|end_header_id|>"
+    
+    content_stream = Stream(app)
+    thread = threading.Thread(
+        target=run_replicate_inference_loop,
+        kwargs={
+            'replicate_client': app.replicate_client,
+            'model': LLAMA3_MODEL,
+            'prompt': full_prompt,
+            'content_stream': content_stream,
+        }
+    )
+
+    thread.start()
+    response = await content_stream.content_wait()
+    
+    return response.lstrip()
+
+
+def run_replicate_inference_loop(replicate_client, model, prompt, content_stream):
+    print("START INFERENCE")
+    completion = replicate_client.run(
+        model,
+        input={
+            "prompt": prompt,
+            "temperature": 0.75,
+            "max_new_tokens": 500,
+            "system_prompt": "" # We handle system prompt in the formatted prompt
+        }
+    )
+    for event in completion:
+        if event:
+            content_stream.feed(str(event))
+    print("END INFERENCE")
+    content_stream.close()