377 lines
13 KiB
Python
377 lines
13 KiB
Python
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()) |