from datetime import datetime import asyncio import json import time import sys import os import signal 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 drives import init_drives from remote import RemoteConnection, RemoteToolSpec, parse_function_specs from app_mixins import ChatMixin, AgentMixin class App(ChatMixin, AgentMixin): def __init__(self, data_directory, replicate_token=None, openrouter_token=None): self.data_directory = data_directory self.user_connections = {} # Connected remotes: (user_id, remotename) -> RemoteConnection. self.remotes = {} self.streams = {} self.event_listeners = [] # active websocket tasks, for manual cancelling upon server shutdown self.ws_tasks = [] self.tmp = tmp.State() # Clients self.qdrant_client = QdrantClient(path=f"{self.data_directory}/qdrant_data") self.replicate_client = replicate_token and replicate.Client(api_token=replicate_token) or None self.openrouter_client = openrouter_token and OpenRouter(api_key=openrouter_token) or None self.agent_classes = [ agents_mod.ExpertAgent, ] from features.calc import Calc from features.web import WebSearch, WebFetch self.tool_classes = [Calc, WebSearch, WebFetch] # --- 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: result = listener(name, data) # Kinda inspired by stream if asyncio.iscoroutinefunction(listener): asyncio.create_task(result) 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] # --- Remotes --- def connectRemote(self, user_id, name, function_specs, send): """Connect a remote served by a live session's send channel. Remotes are scoped to the user: connecting a name that is already connected fails (no overriding). The remote stays connected until disconnectRemote — the 'remote-disconnect' action or the close of the serving session.""" if not name: raise Exception("remotename is required") if (user_id, name) in self.remotes: raise Exception(f"remote '{name}' is already connected") connection = RemoteConnection( user_id, name, parse_function_specs(function_specs), send) self.remotes[(user_id, name)] = connection self.emit_event("remote-connected", {"remotename": name, "remote": connection.describe()}) return connection def disconnectRemote(self, user_id, name): """Disconnect a remote of this user; its pending calls fail. Idempotent: returns the connection, or None when no remote is connected under that name (used by both the manual disconnect action and session teardown).""" connection = self.remotes.pop((user_id, name), None) if connection is not None: connection.close() self.emit_event("remote-disconnected", {"remotename": name}) return connection def listRemotes(self, user_id): """All connected remotes of this user.""" return [conn.describe() for key, conn in self.remotes.items() if key[0] == user_id] async def callRemote(self, user_id, name, function, args): """Invoke a function of a connected remote and await the result. Raises on: unknown remote, unknown function, timeout, or an error reported by the remote.""" connection = self.remotes.get((user_id, name)) if connection is None: raise Exception(f"no connected remote named '{name}'") return await connection.call(function, args) def getRemoteTools(self, user_id, remote_specs): """Expand run-input remote specs into agent tools. Entries are remote names, or {'name': ..., 'prefix': ...} to prefix the tool names. Every function of the remote becomes one tool named after the function itself, prefixed only when requested. Remotes that are not connected are skipped with a warning: they can vanish at any time and the run should proceed without them. """ tools = [] for spec in remote_specs or []: if isinstance(spec, str): name, prefix = spec, None else: name, prefix = spec.get('name'), spec.get('prefix', None) connection = self.remotes.get((user_id, name)) if connection is None: print(f"WARNING: remote '{name}' requested by run input is not connected, skipping") continue for fn in connection.functions.values(): tools.append(fn.get_spec(connection, prefix=prefix)) return tools def resolve_remote_call(self, request_id, result=None, error=None): if result is not None: RemoteConnection.resolve(request_id, result); else: RemoteConnection.reject(request_id, error); # --- 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 # --- Resource Management --- 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: await db.Model.filter(id=model_id).update( name=name, kind=kind, value=value, ) model = await db.Model.get(id=model_id) print(f"Saved model: {model.id} - {name}") return model # --- Main Entry Point --- async def main(): # Get data directory from environment variable data_directory = os.getenv("HUB_DATA") if not data_directory: raise ValueError("HUB_DATA environment variable must be set") port = os.getenv("HUB_PORT", 17922) # 1. Initialize Database First await db.init_database(data_directory) print("Database initialized with Tortoise ORM.") # Drives as well await init_drives(data_directory) # 2. Initialize App horde_app = App( data_directory=data_directory, replicate_token=os.getenv('REPLICATE_API_TOKEN'), openrouter_token=os.getenv("OPENROUTER_API_KEY") ) # 3. Start Server # manually trigger server shutdown, avoids hypercorn installing signal handlers shutdown_event = asyncio.Event() # TODO: get port from env or don't pass it (actually server should handle none) server_task = asyncio.create_task(start_server(horde_app, port, shutdown_event)) # 4. Graceful Termination loop = asyncio.get_running_loop() def interrupt(): print("Interrupt Signal") shutdown_event.set() # Unblock websocket handlers parked in websocket.receive() for task in list(horde_app.ws_tasks): task.cancel() for sig in (signal.SIGINT, signal.SIGTERM): loop.add_signal_handler(sig, interrupt) try: await server_task finally: print("Server stopped") horde_app.qdrant_client.close() await Tortoise.close_connections() if __name__ == "__main__": asyncio.run(main())