Tortoise based database models
This commit is contained in:
Shivering Ocelot committed
1 parent
2c7b3d61d1
commit
cccdb7d231
12 files changed
+469
-382
No files matched your search
+155
-180
@@ -4,8 +4,8 @@ import json
|
||||
import sys
|
||||
import os
|
||||
|
||||
from sqlalchemy.orm import Session
|
||||
from sqlalchemy.sql.expression import select
|
||||
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
|
||||
@@ -25,8 +25,7 @@ class App(ChatMixin):
|
||||
self.remote_connections = {}
|
||||
self.streams = {}
|
||||
|
||||
# Initialize Clients
|
||||
self.sql_engine = db.create_sql_engine(self.data_directory)
|
||||
# Clients
|
||||
self.qdrant_client = QdrantClient(path=f"{self.data_directory}/qdrant_data")
|
||||
|
||||
if replicate_api_token:
|
||||
@@ -34,21 +33,14 @@ class App(ChatMixin):
|
||||
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,
|
||||
agents_mod.CharacterAgent,
|
||||
agents_mod.ExpertAgent,
|
||||
agents_mod.ScriptedAgent,
|
||||
|
||||
#agent_mod.LlamaAgent,
|
||||
#agent_mod.CoderV1Agent,
|
||||
#agent_mod.ToolUsingAgent,
|
||||
#agent_mod.DebugAgent,
|
||||
]
|
||||
|
||||
def sql_session(self):
|
||||
return Session(self.sql_engine)
|
||||
# --- Connection Management (Synchronous) ---
|
||||
|
||||
def add_user_connection(self, user_id, channel):
|
||||
if user_id not in self.user_connections:
|
||||
@@ -60,96 +52,87 @@ class App(ChatMixin):
|
||||
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
|
||||
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 channel_list not in user_connections[remotename]:
|
||||
|
||||
if remotename not in user_connections:
|
||||
user_connections[remotename] = []
|
||||
user_connections[remotename].append(channel)
|
||||
if channel not in 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
|
||||
# --- Agent Methods (Asynchronous) ---
|
||||
|
||||
session.add(chat)
|
||||
session.commit()
|
||||
main_chat_id = chat.id
|
||||
async def create_agent(self, user_id, name, agentname, description, picture, kind, data, main_chat_id=None):
|
||||
print("create_agent; user_id, name:", user_id, name)
|
||||
|
||||
agent = db.Agent()
|
||||
agent.user_id = user_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 db.Agent.get_or_none(id=agent_id)
|
||||
if not agent:
|
||||
raise Exception(f"Agent not found with id {agent_id}")
|
||||
|
||||
summary = []
|
||||
if name is not None:
|
||||
agent.name = name
|
||||
agent.agentname = agentname
|
||||
agent.main_chat_id = main_chat_id
|
||||
summary.append('name')
|
||||
if description is not None:
|
||||
agent.description = description
|
||||
summary.append('description')
|
||||
if picture is not None:
|
||||
agent.picture = picture
|
||||
agent.kind = kind
|
||||
agent.data = data or {}
|
||||
summary.append('picture')
|
||||
if agentname is not None:
|
||||
agent.agentname = agentname
|
||||
summary.append('agentname')
|
||||
|
||||
session.add(agent)
|
||||
session.commit()
|
||||
print(f"Created agent #{agent.id}: {agent.name}")
|
||||
if data:
|
||||
# Tortoise tracks changes in JSONField automatically
|
||||
agent.data.update(data)
|
||||
summary.extend([f"data.{k}" for k in data.keys()])
|
||||
|
||||
membership = db.Membership()
|
||||
membership.chat_id = main_chat_id
|
||||
membership.agent_id = agent.id
|
||||
session.add(membership)
|
||||
session.commit()
|
||||
await agent.save()
|
||||
print(f"Updated agent #{agent.id} fields: {summary}")
|
||||
|
||||
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]
|
||||
async def get_agents(self, user_id):
|
||||
agents = await db.Agent.filter(user_id=user_id).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]
|
||||
@@ -163,141 +146,133 @@ class App(ChatMixin):
|
||||
"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()
|
||||
# --- User Methods (Asynchronous) ---
|
||||
|
||||
if user and user.password == password:
|
||||
return user.id
|
||||
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
|
||||
|
||||
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()
|
||||
# --- Agent Control & Actions ---
|
||||
|
||||
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}")
|
||||
async def get_agent_control(self, agent_id):
|
||||
agent = await db.Agent.get_or_none(id=agent_id)
|
||||
if not agent:
|
||||
raise Exception(f"Agent not found with id {agent_id}")
|
||||
|
||||
control = agent_class(app, agent.data)
|
||||
control._set_basic_info(agent_id, agent.name)
|
||||
return control
|
||||
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
|
||||
|
||||
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}")
|
||||
async def get_agent_config_fields(self, agent_id):
|
||||
agent = await db.Agent.get_or_none(id=agent_id)
|
||||
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}")
|
||||
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 {}
|
||||
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 = self.get_agent_control(agent_id)
|
||||
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 = self.get_agent_control(agent_id)
|
||||
control = await 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(
|
||||
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=[],
|
||||
)
|
||||
agent_id = self.create_agent(
|
||||
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", # Should be 'intro'
|
||||
kind="conversation",
|
||||
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
|
||||
# --- Resource Management ---
|
||||
|
||||
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()
|
||||
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
|
||||
|
||||
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
|
||||
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,
|
||||
defaults={"value": value}
|
||||
)
|
||||
print(f"Saved agent resource: {agent_id} - {kind} - {key}")
|
||||
|
||||
session.add(new_resource)
|
||||
session.commit()
|
||||
print(f"Saved agent resource: {agent_id} - {kind} - {key} = {value}")
|
||||
# --- Main Entry Point ---
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
async def 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(
|
||||
# 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'),
|
||||
)
|
||||
start_server(app)
|
||||
|
||||
# 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())
|
||||
Reference in new issue
Block a user