300 lines
11 KiB
Python
300 lines
11 KiB
Python
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
|
|
from qdrant_client import QdrantClient
|
|
import replicate
|
|
|
|
import database as db
|
|
import agent as agent_mod
|
|
from server import start_server
|
|
from chat import ChatView, Message, MESSAGE_PAGE_SIZE
|
|
|
|
from agents.conversation import ConversationAgent
|
|
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 = {}
|
|
|
|
# Initialize Clients
|
|
self.sql_engine = db.create_sql_engine(self.data_directory)
|
|
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:
|
|
raise Exception("Horde only works with a replicate connection. Please provide an API Token.")
|
|
|
|
# `Agent` is the kind 'basic', will be class `BasicAgent` in the future
|
|
self.agent_classes = [
|
|
agent_mod.BasicAgent,
|
|
ConversationAgent,
|
|
agent_mod.Llama8bAgent,
|
|
agent_mod.CoderV1Agent,
|
|
agent_mod.ToolUsingAgent,
|
|
agent_mod.DebugAgent,
|
|
]
|
|
|
|
def sql_session(self):
|
|
return Session(self.sql_engine)
|
|
|
|
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]: # Check if the list is empty
|
|
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 channel_list not in user_connections[remotename]:
|
|
user_connections[remotename] = []
|
|
user_connections[remotename].append(channel)
|
|
|
|
def create_agent(self, user_id, name, agentname, description, picture, kind, data, main_chat_id=None):
|
|
with self.sql_session() as session:
|
|
if main_chat_id is None:
|
|
chat = db.Chat()
|
|
chat.user_id = user_id
|
|
chat.name = name
|
|
|
|
session.add(chat)
|
|
session.commit()
|
|
main_chat_id = chat.id
|
|
|
|
agent = db.Agent()
|
|
agent.user_id = user_id
|
|
agent.name = name
|
|
agent.agentname = agentname
|
|
agent.main_chat_id = main_chat_id
|
|
agent.description = description
|
|
agent.picture = picture
|
|
agent.kind = kind
|
|
agent.data = data
|
|
|
|
session.add(agent)
|
|
session.commit()
|
|
print(f"Created agent #{agent.id}: {agent.name}")
|
|
|
|
membership = db.Membership()
|
|
membership.chat_id = main_chat_id
|
|
membership.agent_id = agent.id
|
|
session.add(membership)
|
|
session.commit()
|
|
|
|
return agent.id
|
|
|
|
# Updates the agent's basic fields, and applies field updates in its data
|
|
def update_agent(self, agent_id, name=None, description=None, picture=None, agentname=None, data=None):
|
|
with self.sql_session() as session:
|
|
agent = session.execute(
|
|
select(db.Agent).where(db.Agent.id == agent_id)
|
|
).scalars().first()
|
|
|
|
if not agent:
|
|
raise Exception(f"Agent not found with id {agent_id}")
|
|
|
|
summary = []
|
|
|
|
if name != None:
|
|
agent.name = name
|
|
summary.append('name')
|
|
if description != None:
|
|
agent.description = description
|
|
summary.append('description')
|
|
if picture != None:
|
|
agent.picture = picture
|
|
summary.append('picture')
|
|
if agentname != None:
|
|
agent.agentname = agentname
|
|
summary.append('agentname')
|
|
|
|
if data:
|
|
# Assign a copy to trigger a value change
|
|
agent.data = agent.data.copy() if agent.data else {}
|
|
agent.data.update(data)
|
|
summary.extend([f"data.{name}" for name in data.keys()])
|
|
|
|
session.commit()
|
|
print(f"Updated agent #{agent.id} fields: {summary}")
|
|
|
|
def get_agents(self, user_id):
|
|
with self.sql_session() as session:
|
|
agents = session.execute(select(db.Agent).where(db.Agent.user_id == user_id)).scalars().all()
|
|
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),
|
|
}
|
|
|
|
def verify_user_credentials(self, username, password):
|
|
with self.sql_session() as session:
|
|
user = session.execute(
|
|
select(db.User).where(db.User.username == username)
|
|
).scalars().first()
|
|
|
|
if user and user.password == password:
|
|
return user.id
|
|
return None
|
|
|
|
def create_new_user(self, username, email, password):
|
|
with self.sql_session() as session:
|
|
result = session.execute(
|
|
select(db.User.id).where(db.User.email == email or db.User.username == username)
|
|
).first()
|
|
|
|
if result is None:
|
|
user = db.User()
|
|
user.email = email
|
|
user.username = username
|
|
user.password = password
|
|
|
|
session.add(user)
|
|
session.commit()
|
|
|
|
user_id = user.id
|
|
print(f"Created user #{user_id}: {username} ({email})")
|
|
else:
|
|
print(f"Signup error: '{email}' or '{username}' are already used")
|
|
return user_id
|
|
|
|
def get_agent_control(self, agent_id):
|
|
with self.sql_session() as session:
|
|
agent = session.execute(
|
|
select(db.Agent).where(db.Agent.id == agent_id)
|
|
).scalars().first()
|
|
|
|
if not agent:
|
|
raise Exception(f"Agent not found with id {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(app, agent.data)
|
|
control._set_basic_info(agent_id, agent.name)
|
|
return control
|
|
|
|
def get_agent_config_fields(self, agent_id):
|
|
with self.sql_session() as session:
|
|
agent = session.execute(
|
|
select(db.Agent).where(db.Agent.id == agent_id)
|
|
).scalars().first()
|
|
if not agent:
|
|
raise Exception(f"Agent not found with id {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 {}
|
|
fields = [{
|
|
'name': field.name,
|
|
'type': field.type,
|
|
'value': data.get(field.name, field.default)
|
|
} for field in agent_class.schema_fields]
|
|
|
|
return fields
|
|
return {}
|
|
|
|
async def send_agent_prompt(self, agent_id, message):
|
|
control = self.get_agent_control(agent_id)
|
|
response = await control.prompt(message)
|
|
return response
|
|
|
|
async def call_agent_action(self, agent_id, action, args):
|
|
control = self.get_agent_control(agent_id)
|
|
response = await control.execute_action(action, args)
|
|
return response
|
|
|
|
def generate_default_user_data(self, user_id):
|
|
chat_id = self.create_chat(
|
|
user_id=user_id,
|
|
name="Initial Chat",
|
|
description="First user chat with Intro.",
|
|
picture=None,
|
|
members=[],
|
|
)
|
|
agent_id = 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", # Should be 'intro'
|
|
data={}
|
|
)
|
|
|
|
def get_agent_resource(self, agent_id, kind, key):
|
|
with self.sql_session() as session:
|
|
stored_value = session.execute(select(db.AgentResource.value).where(
|
|
db.AgentResource.agent_id == agent_id and db.AgentResource.kind == kind and db.AgentResource.key == key
|
|
)).scalar()
|
|
return stored_value
|
|
|
|
def save_agent_resource(self, agent_id, kind, key, value):
|
|
with self.sql_session() as session:
|
|
existing_resource = session.execute(select(db.AgentResource).where(
|
|
db.AgentResource.agent_id == agent_id and db.AgentResource.kind == kind and db.AgentResource.key == key
|
|
)).scalars().first()
|
|
|
|
if existing_resource:
|
|
existing_resource.value = value
|
|
else:
|
|
new_resource = db.AgentResource()
|
|
new_resource.agent_id = agent_id
|
|
new_resource.kind = kind
|
|
new_resource.key = key
|
|
new_resource.value = value
|
|
|
|
session.add(new_resource)
|
|
session.commit()
|
|
print(f"Saved agent resource: {agent_id} - {kind} - {key} = {value}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
if len(sys.argv) > 1:
|
|
data_directory = sys.argv[1]
|
|
else:
|
|
print("No directory specified. Using current directory as default.")
|
|
data_directory = "."
|
|
|
|
app = App(
|
|
data_directory=data_directory,
|
|
replicate_api_token=os.getenv('REPLICATE_API_TOKEN'),
|
|
)
|
|
start_server(app) |