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()

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)

backend/database.py

@@ -57,8 +57,7 @@ class Chat(Base):
     ordering: Mapped[Optional[int]]
 
     # For indexing purposes, some information is taken from the last chat message.
-    # Preview is the (possibly cropped) content,
-    # time is the timestamp of the last event.
+    # Preview is the (possibly cropped) content, time is the timestamp of the message.
     last_message_preview: Mapped[Optional[str]]
     last_message_time: Mapped[Optional[int]]
 
@@ -81,7 +80,7 @@ class Message(Base):
 
     # Timestamp of this message
     time: Mapped[int]
-    # Identifier of the sender: user / system / agent-$ID
+    # Identifier of the sender: user / system / $AGENT_ID
     sender: Mapped[str]
 
     # Predefined set of message types, interpreted as 'text' when empty.

backend/requirements.txt

@@ -1,2 +1,5 @@
 Quart==0.20.0
-sqlalchemy
+sqlalchemy
+nest_asyncio
+replicate
+chromadb

backend/server.py

@@ -1,6 +1,7 @@
 from quart import Quart, Response, request, send_from_directory, websocket
 import json
 import asyncio
+import nest_asyncio
 
 def start_server(app):
     server = Quart(__name__, static_folder='../frontend/build/web')
@@ -49,20 +50,18 @@ def start_server(app):
             return
 
         send_queue = asyncio.Queue()
-        app.add_user_connection(user_id, send_queue)
-
         async def sender_task():
             nonlocal send_queue
             while True:
                 msg = await send_queue.get()
                 await websocket.send(json.dumps(msg))
-
         async def send(msg):
             await send_queue.put(msg)
 
-        asyncio.create_task(sender_task())
-
         try:
+            app.add_user_connection(user_id, send_queue)
+            asyncio.create_task(sender_task())
+
             while True:
                 req = json.loads(await websocket.receive())
                 action = req['action']
@@ -70,7 +69,6 @@ def start_server(app):
                 nonce = req.get('nonce', None)
 
                 if action == "ping":
-                    print("ping")
                     await send({'message': "pong"})
 
                 elif action == "create-agent":
@@ -115,9 +113,12 @@ def start_server(app):
                     )
                     await send({'message': "message sent", 'messageId': message['id'], 'nonce': nonce})
 
-                else:
-                    print("Unknown message: " + action)
+                elif action == "send-agent-prompt":
+                    response = await app.send_agent_prompt(body['agentId'], body['content'])
+                    await send({'message': "propmt sent", 'response': response, 'nonce': nonce})
 
+                else:
+                    await send({'error': 'unknown action: ' + action})
         finally:
             app.remove_user_connection(user_id, send_queue)
 

backend/stream.py

@@ -0,0 +1,103 @@
+import asyncio
+import nest_asyncio
+
+id_counter = 0
+
+def run_fn (fn_obj, arg):
+    if fn_obj == None:
+        return
+
+    # https://stackoverflow.com/questions/36076619/test-if-function-or-method-is-normal-or-asynchronous
+    if asyncio.iscoroutinefunction(fn_obj):
+        if arg == None:
+            asyncio.create_task(fn_obj())
+        else:
+            asyncio.create_task(fn_obj(arg))
+        return
+
+    if arg == None:
+        fn_obj()
+    else:
+        fn_obj(arg)
+
+def make_runnable(fn_obj):
+    """
+    Creates an async signal and a task that listen to the signal.
+    When it is received, fn_obj is called.
+    The task doesn't stop and keeps forwarding signals to the passed function.
+    There can also be arguments coming along with the signal, those are passed to the function.
+    """
+    signal_queue = asyncio.Queue()
+
+    async def signal_listener():
+        while True:
+            args = await signal_queue.get()
+            if args is None:  # Stop signal
+                break
+            run_fn(fn_obj, args[0] if args else None)
+
+    # Start the listener task in the background
+    asyncio.create_task(signal_listener())
+
+    # Return a method to put items on the signal queue
+    return lambda *args: signal_queue.put_nowait(args)
+
+class Stream:
+    def __init__(self, app):
+        global id_counter
+
+        self.id = id_counter
+        id_counter = id_counter + 1
+
+        self.open = True
+        self.content = ""
+        self.listeners = []
+
+        app.streams[self.id] = self
+
+    def listen (self, chunk_fn, end_fn):
+        nest_asyncio.apply()
+        #chunk_fn = make_runnable(chunk_fn)
+        #end_fn = make_runnable(end_fn)
+
+        if self.open:
+            self.listeners.append({
+                "chunk": chunk_fn,
+                "end": end_fn,
+            })
+
+            if self.content != "":
+                run_fn(chunk_fn, self.content)
+        else:
+            run_fn(chunk_fn, self.content)
+            run_fn(end_fn, None)
+
+    def feed (self, chunk):
+        if not self.open:
+            return
+
+        self.content = self.content + chunk
+        for listener in self.listeners:
+            run_fn(listener["chunk"], chunk)
+
+    def close (self):
+        self.open = False
+
+        for listener in self.listeners:
+            run_fn(listener["end"], None)
+
+    async def content_wait (self):
+        async_loop = asyncio.get_event_loop()
+        event = asyncio.Event()
+
+        def set_event():
+            event.set()
+
+        def noop (chunk): pass
+        def end_fn ():
+            async_loop.call_soon_threadsafe(set_event)
+
+        self.listen(chunk_fn=noop, end_fn=end_fn)
+        await event.wait()
+
+        return self.content