From 49df72f2a105634de0b9a9171c58d642180219e3 Mon Sep 17 00:00:00 2001 From: Shivering Ocelot Date: Mon, 12 May 2025 12:55:50 -0500 Subject: [PATCH] backend: implement agents and connect them with chats --- backend/agent.py | 93 +++++++++++++++++++++++++++++++++++ backend/app.py | 62 +++++++++++++++++++---- backend/database.py | 5 +- backend/replicate.py | 0 backend/requirements.txt | 5 +- backend/server.py | 17 ++++--- backend/stream.py | 103 +++++++++++++++++++++++++++++++++++++++ run.sh | 22 +++++++-- testing/suite.js | 46 +++++++++++------ 9 files changed, 313 insertions(+), 40 deletions(-) create mode 100644 backend/agent.py delete mode 100644 backend/replicate.py create mode 100644 backend/stream.py diff --git a/backend/agent.py b/backend/agent.py new file mode 100644 index 0000000..f02542f --- /dev/null +++ b/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() \ No newline at end of file diff --git a/backend/app.py b/backend/app.py index 8ad8c53..778736c 100644 --- a/backend/app.py +++ b/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) \ No newline at end of file diff --git a/backend/database.py b/backend/database.py index 8ad1968..ed7fdc9 100644 --- a/backend/database.py +++ b/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. diff --git a/backend/replicate.py b/backend/replicate.py deleted file mode 100644 index e69de29..0000000 diff --git a/backend/requirements.txt b/backend/requirements.txt index 00a2728..a3da446 100644 --- a/backend/requirements.txt +++ b/backend/requirements.txt @@ -1,2 +1,5 @@ Quart==0.20.0 -sqlalchemy \ No newline at end of file +sqlalchemy +nest_asyncio +replicate +chromadb \ No newline at end of file diff --git a/backend/server.py b/backend/server.py index cabea02..1c2fa8b 100644 --- a/backend/server.py +++ b/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) diff --git a/backend/stream.py b/backend/stream.py new file mode 100644 index 0000000..0546a8f --- /dev/null +++ b/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 diff --git a/run.sh b/run.sh index 77d5b0c..7a60845 100644 --- a/run.sh +++ b/run.sh @@ -1,10 +1,17 @@ #!/bin/bash +source .venv/bin/activate + +if [ -f .env ]; then + source .env + export REPLICATE_API_TOKEN +fi + case "$1" in "") + mkdir -p local-data cd backend - source ../.venv/bin/activate - python app.py + python app.py ../test-data ;; "test") # Setup test data directory @@ -13,7 +20,6 @@ case "$1" in # Start backend in the background cd backend - source ../.venv/bin/activate python app.py ../test-data & BACKEND_PID=$! echo "Started backend: $BACKEND_PID" @@ -31,7 +37,6 @@ case "$1" in "debug") # Debug doesn't work very well... I don't use it. cd backend - source ../.venv/bin/activate python app.py & cd ../frontend @@ -51,6 +56,13 @@ case "$1" in done ;; *) - echo "Error: unknown command '$1'" + if [ -d "$1" ]; then + data_dir=../$1 + mkdir -p local-data + cd backend + python app.py $data_dir + else + echo "Error: unknown command or directory '$1'" + fi ;; esac \ No newline at end of file diff --git a/testing/suite.js b/testing/suite.js index 6d2a3ae..395cfe3 100644 --- a/testing/suite.js +++ b/testing/suite.js @@ -1,10 +1,10 @@ const WebSocket = require('ws'); -const ws = new WebSocket('ws://localhost:5000/ws'); +const USE_LLM = false; +const ws = new WebSocket('ws://localhost:5000/ws'); var nonceCounter = 1; var nonceCallbacks = {}; - var messageQueue = []; var messageResolverQueue = []; @@ -73,9 +73,18 @@ async function runTest() { }}) ); - console.log("Get Agents", - await send('get-agents') - ); + let agents = (await send('get-agents')).agents; + console.log("Get Agents", agents); + + if (USE_LLM) { + let promptResponse = await send('send-agent-prompt', { + body: { + agentId: agents[0].id, + content: "Hello World!", + } + }); + console.log("Prompt Response", promptResponse); + } await testChats(); @@ -86,14 +95,23 @@ async function testChats () { let response = await send('get-chats') let chat = response.chats[0]; - let msgRes = await send('send-chat-message', { - body: { - chatId: chat.id, - content: "Hello World!", - } - }); - console.log('message sent', msgRes); + if (USE_LLM) { + let msgRes = await send('send-chat-message', { + body: { + chatId: chat.id, + content: "Whats up!", + } + }); + console.log('message sent', msgRes); + console.log('Incoming:', await waitMessage()); - let nextMessage = await waitMessage(); - console.log('next incoming message', nextMessage); + // Multi turn conversation + await send('send-chat-message', { + body: { + chatId: chat.id, + content: "Spell the initials of my first message.", + } + }); + console.log('Incoming:', await waitMessage()); + } } \ No newline at end of file