265 lines
9.7 KiB
Python
265 lines
9.7 KiB
Python
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
|
|
self.tool_classes = [Calc]
|
|
|
|
# --- 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()) |