backend: implement agents and connect them with chats

backend/app.py

@@ -1,18 +1,36 @@
-from database import create_sql_engine, User, Agent, Chat, Message, Membership
-from server import start_server
-from sqlalchemy.orm import Session
-from sqlalchemy.sql.expression import select
-from quart import Quart, Response, request, send_from_directory, jsonify, websocket
 from datetime import datetime
 import asyncio
 import json
 import sys
+import os
+
+from sqlalchemy.orm import Session
+from sqlalchemy.sql.expression import select
+from quart import Quart, Response, request, send_from_directory, jsonify, websocket
+import replicate
+
+from database import create_sql_engine, User, Agent, Chat, Message, Membership
+from server import start_server
+from agent import AgentControl
+
+MESSAGE_PAGE_SIZE = 32
+
+class ChatContext:
+    def __init__(self, messages, prelude=None):
+        self.messages=messages
+        self.prelude=prelude
 
 class App:
-    def __init__(self, data_directory):
+    def __init__(self, data_directory, replicate_api_token=None):
         self.data_directory = data_directory
         self.sql_engine = create_sql_engine(self.data_directory)
         self.user_connections = {}
+        self.streams = {}
+
+        if replicate_api_token:
+            self.replicate_client = replicate.Client(api_token=replicate_api_token)
+        else:
+            raise Exception("Horde only works with a replicate client. Please provide an API Token.")
 
     def sql_session(self):
         return Session(self.sql_engine)
@@ -81,7 +99,7 @@ class App:
     def get_chat_messages(self, chat_id, earliest_index=0):
         with self.sql_session() as session:
             messages = session.execute(
-                select(Message).where(Message.chat_id == chat_id).order_by(Message.id.desc())
+                select(Message).where(Message.chat_id == chat_id).order_by(Message.time.asc())
             ).scalars().fetchmany(MESSAGE_PAGE_SIZE)
 
             return [{
@@ -166,6 +184,26 @@ class App:
                 'content': message.content,
             }
 
+    def get_chat_context(self, chat_id):
+        return ChatContext(self.get_chat_messages(chat_id))
+
+    def get_agent_control(self, agent_id):
+        with self.sql_session() as session:
+            agent = session.execute(
+                select(Agent).where(Agent.id == agent_id)
+            ).scalars().first()
+
+            if not agent:
+                raise Exception(f"Agent not found with id {agent_id}")
+        
+            control = AgentControl.create_from_data(app, agent_id, agent.kind, agent.data)
+            return control
+
+    async def send_agent_prompt(self, agent_id, message):
+        control = self.get_agent_control(agent_id)
+        response = await control.prompt(message)
+        return response
+
     def notify_message_to_user(self, user_id, message):
         if user_id in self.user_connections:
             user_message = {
@@ -180,7 +218,8 @@ class App:
                 asyncio.create_task(channel.put(user_message))
 
     def notify_message_to_agent(self, agent_id, message):
-        pass # TODO
+        control = self.get_agent_control(agent_id)
+        asyncio.create_task(control.received_message(message))
 
     def generate_default_user_data(self, user_id):
         agent_id = self.create_agent(
@@ -210,5 +249,10 @@ if __name__ == "__main__":
         print("No directory specified. Using current directory as default.")
         data_directory = "."
 
-    app = App(data_directory)
+    print("api token", os.getenv('REPLICATE_API_TOKEN'))
+
+    app = App(
+        data_directory=data_directory,
+        replicate_api_token=os.getenv('REPLICATE_API_TOKEN'),
+    )
     start_server(app)