from datetime import datetime import asyncio import json import time import sys import os from tortoise import Tortoise from tortoise.expressions import Q from quart import Quart, Response, request, send_from_directory, jsonify, websocket from qdrant_client import QdrantClient import replicate from openrouter import OpenRouter import database as db import agent as agent_mod import agents as agents_mod import temporary as tmp from server import start_server from chat import ChatView, Message, MESSAGE_PAGE_SIZE from app_mixins import ChatMixin class App(ChatMixin): def __init__(self, data_directory, replicate_api_token=None): self.data_directory = data_directory self.user_connections = {} self.remote_connections = {} self.streams = {} self.event_listeners = [] self.tmp = tmp.State() # Clients self.qdrant_client = QdrantClient(path=f"{self.data_directory}/qdrant_data") if replicate_api_token: self.replicate_client = replicate.Client(api_token=replicate_api_token) else: self.replicate_client = None openrouter_api_key = os.getenv("OPENROUTER_API_KEY") if openrouter_api_key: self.openrouter_client = OpenRouter(api_key=openrouter_api_key) else: self.openrouter_client = None self.agent_classes = [ agents_mod.ExpertAgent, ] # --- Connection Management (Synchronous) --- def on_event(self, callback): """Register a callback for events.""" if callback not in self.event_listeners: self.event_listeners.append(callback) return lambda: self.event_listeners.remove(callback) def emit_event(self, name, data=None): """Emit an event to all listeners.""" for listener in self.event_listeners: try: listener(name, data) except Exception as e: print(f"Error in event listener: {name}, {data}") def add_user_connection(self, user_id, channel): if user_id not in self.user_connections: self.user_connections[user_id] = [] if channel not in self.user_connections[user_id]: self.user_connections[user_id].append(channel) def remove_user_connection(self, user_id, channel): if user_id in self.user_connections: if channel in self.user_connections[user_id]: self.user_connections[user_id].remove(channel) if not self.user_connections[user_id]: del self.user_connections[user_id] def add_remote_connection(self, user_id, remotename, channel): if user_id not in self.remote_connections: self.remote_connections[user_id] = {} user_connections = self.remote_connections[user_id] if remotename not in user_connections: user_connections[remotename] = [] if channel not in user_connections[remotename]: user_connections[remotename].append(channel) # --- Agent Methods (Asynchronous) --- async def get_agent(self, agent_id): if tmp.is_tmp_id(agent_id): agent = self.tmp.agents.get(agent_id) else: agent = await db.Agent.get_or_none(id=agent_id) if not agent: raise Exception(f"Agent not found with id {agent_id}") return agent async def create_agent(self, user_id, name, agentname, description, picture, kind, data, main_chat_id=None, temporary=False): print("create_agent; user_id, name:", user_id, name) if temporary: agent = tmp.Agent.create( self, user_id=user_id, name=name, agentname=agentname, main_chat_id=main_chat_id, description=description, picture=picture, kind=kind, data=data or {} ) if main_chat_id is not None: if tmp.is_tmp_id(main_chat_id): self.tmp.chats[main_chat_id].participants.append(agent.id) else: print(f"ERROR: cannot use {main_chat_id} as main chat for temporary agent {agentname or agent.id}") agent.main_chat_id = None return agent.id if main_chat_id is None: chat = await db.Chat.create( user_id=user_id, name=name ) main_chat_id = chat.id agent = await db.Agent.create( user_id=user_id, name=name, agentname=agentname, main_chat_id=main_chat_id, description=description, picture=picture, kind=kind, data=data or {} ) print(f"Created agent #{agent.id}: {agent.name}") print(f"... adding agent {repr(agent.id)} to chat {repr(main_chat_id)}") await db.Membership.create( chat_id=main_chat_id, agent_id=agent.id ) return agent.id async def update_agent(self, agent_id, name=None, description=None, picture=None, agentname=None, data=None): agent = await self.get_agent(agent_id) summary = [] if name is not None: agent.name = name summary.append('name') if description is not None: agent.description = description summary.append('description') if picture is not None: agent.picture = picture summary.append('picture') if agentname is not None: agent.agentname = agentname summary.append('agentname') if data: # Tortoise tracks changes in JSONField automatically agent.data.update(data) summary.extend([f"data.{k}" for k in data.keys()]) if db.is_instance(agent): await agent.save() print(f"Updated agent #{agent.id} fields: {summary}") async def get_agents(self, user_id): agents = await db.Agent.filter(user_id=user_id).all() if user_id in self.tmp.user_agents: agents += [self.tmp.agents[agent_id] for agent_id in self.tmp.user_agents[user_id]] return [{ "id": agent.id, "name": agent.name, "agentname": agent.agentname, "description": agent.description, "picture": agent.picture, "kind": agent.kind, } for agent in agents] def get_agent_kinds(self): return [kind.identifier for kind in self.agent_classes] def get_agent_kind_info(self, kind): agent_class = next((it for it in self.agent_classes if it.identifier == kind), None) if agent_class is None: raise Exception(f"invalid agent kind: {kind}") return { "name": agent_class.identifier, "description": getattr(agent_class, 'description', None), } # --- User Methods (Asynchronous) --- async def verify_user_id(self, user_id): existing = await db.User.filter(Q(id=user_id)).exists() return existing async def verify_user_credentials(self, username, password): user = await db.User.get_or_none(username=username) if user and user.password == password: return user.id return None async def create_new_user(self, username, email, password): existing = await db.User.filter(Q(email=email) | Q(username=username)).exists() if not existing: user = await db.User.create( email=email, username=username, password=password ) print(f"Created user #{user.id}: {username} ({email})") return user.id else: print(f"Signup error: '{email}' or '{username}' are already used") return None # --- Agent Control & Actions --- async def get_agent_control(self, agent_id): agent = await self.get_agent(agent_id) agent_class = next((kind for kind in self.agent_classes if kind.identifier == agent.kind), None) if agent_class is None: raise Exception(f"invalid agent kind: {agent.kind}") control = agent_class(self, agent.data) control._set_basic_info(agent_id, agent.name) return control async def get_agent_config_fields(self, agent_id): agent = await self.get_agent(agent_id) agent_class = next((it for it in self.agent_classes if it.identifier == agent.kind), None) if not agent_class: raise Exception(f"Invalid agent kind: {agent.kind}") data = agent.data or {} return [{ 'name': field.name, 'type': field.type, 'value': data.get(field.name, field.default) } for field in agent_class.schema_fields] async def send_agent_prompt(self, agent_id, instruction, input=None, options={}): control = await self.get_agent_control(agent_id) response = await control.prompt(instruction, input, **options) return response async def call_agent_action(self, agent_id, action, args): control = await self.get_agent_control(agent_id) response = await control.execute_action(action, args) return response async def generate_default_user_data(self, user_id): chat_id = await self.create_chat( user_id=user_id, name="Initial Chat", description="First user chat with Intro.", picture=None, members=[], ) await self.create_agent( user_id=user_id, name="Intro", agentname="intro", main_chat_id=chat_id, description="Introduces you to the application.", picture=None, kind="conversation", data={} ) # --- Resource Management --- async def get_agent_resource(self, agent_id, kind, key): resource = await db.AgentResource.get_or_none(agent_id=agent_id, kind=kind, key=key) return resource.value if resource else None async def save_agent_resource(self, agent_id, kind, key, value): await db.AgentResource.update_or_create( agent_id=agent_id, kind=kind, key=key, value=value, ) print(f"Saved agent resource: {agent_id} - {kind} - {key}") async def get_models(self, kind=None): return await db.Model.all() async def get_model(self, model_id): return await db.Model.get(id=model_id) # Serves for creating and updating models. # Deletes the model by the id if all other parameters are empty. async def save_model(self, model_id, kind, name, value): if model_id is None: model = await db.Model.create( name=name, kind=kind, value=value, ) else: if value is None and name is None and kind is None: await db.Model.filter(id=model_id).delete() print(f"Deleted model: {model_id}") return None else: model = await db.Model.update_or_create( id=model_id, name=name, kind=kind, value=value, ) print(f"Saved model: {model.id} - {name}") return model async def create_agent_run(self, agent_id): agent = await self.get_agent(agent_id) run = agent_mod.AgentRun.create( self, agent_id=agent_id, user_id=agent.user_id, ) return run.id async def get_agent_runs(self, agent_id=None, user_id=None): runs = list(self.tmp.runs.values()) if agent_id: runs = [r for r in runs if r.agent_id == agent_id] elif user_id: runs = [r for r in runs if r.user_id == user_id] return [r.serialize() for r in runs] # --- Main Entry Point --- async def main(): if len(sys.argv) > 1: data_directory = sys.argv[1] else: data_directory = "." # 1. Initialize Database First await db.init_database(data_directory) print("Database initialized with Tortoise ORM.") # 2. Initialize App horde_app = App( data_directory=data_directory, replicate_api_token=os.getenv('REPLICATE_API_TOKEN'), ) # 3. Start Server # Note: start_server must be an async function or handle an event loop await start_server(horde_app) if __name__ == "__main__": try: asyncio.run(main()) except KeyboardInterrupt: # Clean shutdown for Tortoise connections asyncio.run(Tortoise.close_connections())