Tortoise based database models

This commit is contained in:
Shivering Ocelot committed 2026-01-03 12:41:17 -05:00
1 parent 2c7b3d61d1
commit cccdb7d231
12 files changed
+469 -382

No files matched your search

+5 -5
View File
@@ -65,10 +65,10 @@ class DebugAgent(Agent):
tools = []
async def received_message(self, message):
chat_view = self.app.get_chat_view(message.chat_id)
chat_view = await self.app.get_chat_view(message.chat_id)
response = chat_view.debug_string()
self.app.send_chat_message(
await self.app.send_chat_message(
chat_id=message.chat_id,
sender=self.id,
content=response,
@@ -178,7 +178,7 @@ class ToolUsingAgent:
system_prologue="You are given the entire context of the conversation, respond with your next message in the conversation by using the '!!response' command."
)
self.app.send_chat_message(
await self.app.send_chat_message(
chat_id=message.chat_id,
sender=self.id,
content=response,
@@ -235,7 +235,7 @@ class LlamaAgent(Agent):
is_self = is_digit and str(sender) == str(self.id)
return self.name if is_self else sender
chat_view = self.app.get_chat_view(message.chat_id)
chat_view = await self.app.get_chat_view(message.chat_id)
thread_messages = chat_view.get_thread(message.thread_id).messages
chat_summary = "\n".join(f"[{sender_name(message)}]: {message.content}" for message in thread_messages)
messages = [
@@ -244,7 +244,7 @@ class LlamaAgent(Agent):
]
response = await replicate_llama_completion(self.app, messages)
self.app.send_chat_message(
await self.app.send_chat_message(
chat_id=message.chat_id,
sender=self.id,
content=response,
+3 -3
View File
@@ -144,7 +144,7 @@ class CharacterAgent(Agent):
# TODO: before retrieving the chat context, acquire a lock for it
chat_view = self.app.get_chat_view(message.chat_id)
chat_view = await self.app.get_chat_view(message.chat_id)
current_thread = chat_view.get_thread(message.thread_id)
thread_context = self.app.get_agent_resource(self.id, 'last-thread-context', message.chat_id)
@@ -204,7 +204,7 @@ class CharacterAgent(Agent):
controller_response = await replicate_llama_completion(self.app, llama_messages)
if self.debug:
self.app.send_chat_message(
await self.app.send_chat_message(
chat_id=message.chat_id,
sender=self.id,
content=controller_response,
@@ -216,7 +216,7 @@ class CharacterAgent(Agent):
sections = parse_sections(controller_response)
if 'outwards' in sections:
send_response = self.app.send_chat_message(
send_response = await self.app.send_chat_message(
chat_id=message.chat_id,
sender=self.id,
content=sections['outwards'],
+2 -2
View File
@@ -68,7 +68,7 @@ class ExpertAgent(Agent):
return response
async def received_message(self, message):
chat_view = self.app.get_chat_view(message.chat_id)
chat_view = await self.app.get_chat_view(message.chat_id)
thread = chat_view.get_thread(message.thread_id)
def msg_role(message):
@@ -90,7 +90,7 @@ class ExpertAgent(Agent):
messages=messages
)
self.app.send_chat_message(
await self.app.send_chat_message(
chat_id=message.chat_id,
sender=self.id,
content=response,
+155 -180
View File
@@ -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())
+104 -128
View File
@@ -1,117 +1,100 @@
import asyncio
from datetime import datetime
from sqlalchemy.sql.expression import select
import database as db
# Assuming db is your database.py containing the Tortoise models from the previous step
import database as db
from chat import ChatView, Message, MESSAGE_PAGE_SIZE
class ChatMixin:
def create_chat(self, user_id, name, description, picture, members=[]):
with self.sql_session() as session:
chat = db.Chat()
chat.user_id = user_id
chat.name = name
chat.description = description
chat.picture = picture
async def create_chat(self, user_id, name, description, picture, members=[]):
# Tortoise uses .create() for one-line instantiation and saving
chat = await db.Chat.create(
user_id=user_id,
name=name,
description=description,
picture=picture
)
print(f"Created chat #{chat.id}: {chat.name}")
session.add(chat)
session.commit()
print(f"Created chat #{chat.id}: {chat.name}")
chat_id = chat.id
for agent_id in members:
self.add_agent_to_chat(chat_id, agent_id)
await self.add_agent_to_chat(chat.id, agent_id)
return chat_id
return chat.id
def get_chats(self, user_id):
with self.sql_session() as session:
chats = session.execute(select(db.Chat).where(db.Chat.user_id == user_id)).scalars().all()
return [{
"id": chat.id,
"name": chat.name,
"description": chat.description,
"picture": chat.picture
} for chat in chats]
async def get_chats(self, user_id):
# Direct filtering on the model class
chats = await db.Chat.filter(user_id=user_id).all()
return [{
"id": chat.id,
"name": chat.name,
"description": chat.description,
"picture": chat.picture
} for chat in chats]
def delete_chat(self, chat_id):
with self.sql_session() as session:
session.execute(db.Membership.__table__.delete().where(db.Membership.chat_id == chat_id))
session.execute(db.Chat.__table__.delete().where(db.Chat.id == chat_id))
# TODO: Delete messages from the chat
session.commit()
print(f"Deleted chat #{chat_id} and its memberships.")
async def delete_chat(self, chat_id):
# Filtering and deleting is a one-liner
await db.Membership.filter(chat_id=chat_id).delete()
await db.Chat.filter(id=chat_id).delete()
# TODO: Delete messages from the chat
print(f"Deleted chat #{chat_id} and its memberships.")
def get_chat_messages(self, chat_id, earliest_index=None):
with self.sql_session() as session:
condition = (db.Message.chat_id == chat_id)
if earliest_index is not None:
condition = condition and db.Message.id < earliest_index
async def get_chat_messages(self, chat_id, earliest_index=None):
query = db.Message.filter(chat_id=chat_id)
if earliest_index is not None:
query = query.filter(id__gt=earliest_index)
messages = session.execute(
select(db.Message).where(condition).order_by(db.Message.time.desc())
).scalars().all()
# Order by time descending, then reverse in Python to get chronological order
messages = await query.order_by("-time").all()
messages.reverse()
messages.reverse()
return [{
"id": message.id,
"time": message.time,
"sender": message.sender,
"kind": message.kind,
"content": message.content,
"threadId": message.thread_id,
} for message in messages]
return [{
"id": message.id,
"time": message.time,
"sender": message.sender,
"kind": message.kind,
"content": message.content,
"threadId": message.thread_id,
} for message in messages]
def create_chat_thread(self, chat_id, name=None):
with self.sql_session() as session:
thread = db.Thread()
thread.chat_id = chat_id
thread.name = name
async def create_chat_thread(self, chat_id, name=None):
thread = await db.Thread.create(
chat_id=chat_id,
name=name
)
print(f"Created thread #{thread.id}: {thread.name}")
return thread.id
session.add(thread)
session.commit()
print(f"Created thread #{thread.id}: {thread.name}")
async def add_agent_to_chat(self, chat_id, agent_id):
await db.Membership.create(
chat_id=chat_id,
agent_id=agent_id
)
print(f"Added agent #{agent_id} to chat #{chat_id}")
return thread.id
async def send_chat_message(self, chat_id, content, sender='user', thread_id=None, kind=None, subkind=None, format=None):
now = int(datetime.now().timestamp())
message = await db.Message.create(
chat_id=chat_id,
time=now,
last_edit=now,
sender=sender,
kind=kind or 'message',
subkind=subkind,
format=format,
content=content,
thread_id=thread_id
)
def add_agent_to_chat(self, chat_id, agent_id):
with self.sql_session() as session:
membership = db.Membership()
membership.chat_id = chat_id
membership.agent_id = agent_id
session.add(membership)
session.commit()
print(f"Added agent #{agent_id} to chat #{chat_id}")
# Fetch chat owner using values_list for efficiency
chat = await db.Chat.get(id=chat_id).prefetch_related("user")
chat_owner_id = chat.user_id
self.notify_message_to_user(chat_owner_id, message)
# If an integer is passed in sender, it is interpreted as an agent id
def send_chat_message(self, chat_id, content, sender='user', thread_id=None, kind=None, subkind=None, format=None):
with self.sql_session() as session:
message = db.Message()
message.chat_id = chat_id
message.time = int(datetime.now().timestamp())
message.last_edit = message.time
message.sender = sender
message.kind = kind or 'message'
message.subkind = subkind
message.content = content
message.thread_id = thread_id
session.add(message)
session.commit()
session.refresh(message)
# Notify chat owner
chat_owner = session.execute(select(db.Chat.user_id).where(db.Chat.id == chat_id)).scalar()
self.notify_message_to_user(chat_owner, message)
# Notify participant agents
agent_ids = session.execute(select(db.Membership.agent_id).where(db.Membership.chat_id == chat_id)).scalars().all()
# Fetch participant agents
agent_ids = await db.Membership.filter(chat_id=chat_id).values_list("agent_id", flat=True)
for agent_id in agent_ids:
if agent_id != sender:
self.notify_message_to_agent(agent_id, message)
if str(agent_id) != str(sender):
await self.notify_message_to_agent(agent_id, message)
return {
'id': message.id,
@@ -122,41 +105,35 @@ class ChatMixin:
'content': message.content,
}
# Updates a single chat message.
# Most fields are optional. If they're not passed a value, they wont be updated.
# When kind is set, subkind is always set too.
def update_chat_message(self, message_id, content=None, kind=None, subkind=None, format=None):
with self.sql_session() as session:
message = session.execute(
select(db.Message).where(db.Message.id == message_id)
).scalars().first()
async def update_chat_message(self, message_id, content=None, kind=None, subkind=None, format=None):
message = await db.Message.get_or_none(id=message_id)
if not message:
raise Exception(f"Message not found with id {message_id}")
if not message:
raise Exception(f"Message not found with id {message_id}")
summary = []
if content != None:
message.content = content
summary.append('content')
if kind != None:
message.kind = kind
message.subkind = subkind
summary.append('kind')
summary.append('subkind')
if format != None:
message.format = format
summary.append('format')
#message.last_edit = int(datetime.now().timestamp())
session.commit()
summary = []
if content is not None:
message.content = content
summary.append('content')
if kind is not None:
message.kind = kind
message.subkind = subkind
summary.append('kind')
summary.append('subkind')
if format is not None:
message.format = format
summary.append('format')
if summary:
await message.save() # Saves the specific instance
print(f"Updated message #{message.id} fields: {summary}")
def get_chat_view(self, chat_id):
return ChatView(self, chat_id)
async def get_chat_view(self, chat_id):
chat_view = ChatView(self, chat_id)
await chat_view.init_load()
return chat_view
def notify_message_to_user(self, user_id, message):
# could be moved to user mixin
if user_id in self.user_connections:
user_message = {
'action': 'chat-message',
@@ -171,7 +148,6 @@ class ChatMixin:
for channel in self.user_connections[user_id]:
asyncio.create_task(channel.put(user_message))
def notify_message_to_agent(self, agent_id, message):
# could be moved to agent mixin
control = self.get_agent_control(agent_id)
asyncio.create_task(control.received_message(message))
async def notify_message_to_agent(self, agent_id, message):
control = await self.get_agent_control(agent_id)
asyncio.create_task(control.received_message(message))
+58 -46
View File
@@ -67,57 +67,69 @@ class ChatView:
self.current_thread_id = None
self._earliest_message_id = None
self._load_basic_data()
self._load_next_messages(MESSAGE_PAGE_SIZE)
def _load_basic_data(self):
with self.app.sql_session() as session:
record = session.execute(
select(db.Chat).where(db.Chat.id == self.chat_id)
).scalars().first()
async def init_load(self):
await self._load_basic_data()
await self._load_next_messages(MESSAGE_PAGE_SIZE)
if not record:
raise Exception(f"Chat not found with id {self.chat_id}")
async def _load_basic_data(self):
# Using get_or_none for cleaner existence checks
record = await db.Chat.get_or_none(id=self.chat_id)
if not record:
raise Exception(f"Chat not found with id {self.chat_id}")
self.name = record.name
self.description = record.description
async def _load_next_messages(self, load_count):
# Start building the filter
filters = {"chat_id": self.chat_id}
# Add the 'earliest message' condition if it exists
if self._earliest_message_id is not None:
filters["id__lte"] = self._earliest_message_id
# Tortoise query: filter, order, and limit
records_rev = await db.Message.filter(**filters).order_by("-time").limit(load_count)
# Update the earliest id for the next load operation
if records_rev:
# Note: Since we ordered by -time (descending), index 0 is the newest in this batch
self._earliest_message_id = records_rev[-1].id # The "earliest" is the one with the smallest ID/Time in this set
self.name = record.name
self.description = record.description
if not self.threads:
# current thread is the latest known from the newest record in batch
self.current_thread_id = records_rev[0].thread_id
def _load_next_messages(self, load_count):
with self.app.sql_session() as session:
records_rev = session.execute(
select(db.Message).where(
db.Message.chat_id == self.chat_id,
self._earliest_message_id == None or db.Message.id <= self._earliest_message_id
).order_by(db.Message.time.desc())
).scalars().fetchmany(load_count)
# Go over messages from newest to oldest
for record in records_rev:
if record.thread_id is None:
thread = self.default_thread
else:
thread = self.thread_map.get(record.thread_id)
if not thread:
# Assuming Thread is a local class, not the DB model
thread = Thread(record.thread_id, self.chat_id)
self.thread_map[thread.id] = thread
self.threads.insert(0, thread)
# Create local Message object
message = Message(
record.id,
record.content,
record.sender,
record.time,
record.kind,
record.subkind
)
# Update the earliest id for the next load operation
if records_rev:
self._earliest_message_id = records_rev[0].id
if not self.threads:
# current thread is the latest known
self.current_thread_id = records_rev[-1].thread_id
# Go over messages from newest to oldest
for record in records_rev:
if record.thread_id == None:
thread = self.default_thread
else:
thread = self.thread_map.get(record.thread_id)
if not thread:
thread = Thread(record.thread_id, self.chat_id)
self.thread_map[thread.id] = thread
self.threads.insert(0, thread)
# This is the oldest message so far, add it to the start of the chat
message = Message(record.id, record.content, record.sender, record.time, record.kind, record.subkind)
if message.kind == 'note':
thread.notes.insert(0, message)
elif message.kind == 'memory':
thread.memories.insert(0, message)
else:
thread.messages.insert(0, message)
if message.kind == 'note':
thread.notes.insert(0, message)
elif message.kind == 'memory':
thread.memories.insert(0, message)
else:
thread.messages.insert(0, message)
def get_thread(self, thread_id):
if thread_id == None:
+116 -1
View File
@@ -2,6 +2,8 @@ from typing import Optional
from sqlalchemy import create_engine, Column, Integer, JSON, ForeignKey
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column
from sqlalchemy.sql.expression import select
from tortoise import fields, models, Tortoise
import asyncio
class Base(DeclarativeBase):
pass
@@ -9,8 +11,120 @@ class Base(DeclarativeBase):
def create_sql_engine (data_directory):
sql_engine = create_engine(f"sqlite:///{data_directory}/horde.db")
Base.metadata.create_all(sql_engine)
return sql_engine
async def init_database (data_directory):
db_url=f"sqlite:///{data_directory}/horde.db"
# Setup Tortoise
await Tortoise.init(
db_url=db_url,
modules={'models': ['database']}
)
await Tortoise.generate_schemas()
class User(models.Model):
id = fields.IntField(pk=True)
email = fields.CharField(max_length=255)
username = fields.CharField(max_length=100)
password = fields.CharField(max_length=255)
class Meta:
table = "users"
class Agent(models.Model):
id = fields.IntField(pk=True)
user = fields.ForeignKeyField("models.User", related_name="agents")
agentname = fields.CharField(max_length=100, null=True)
visibility = fields.IntField(default=0)
permissions = fields.TextField(default="")
main_chat = fields.ForeignKeyField("models.Chat", related_name="main_agents",null=True)
name = fields.CharField(max_length=255)
description = fields.TextField(null=True)
picture = fields.TextField(null=True)
ordering = fields.IntField(null=True)
kind = fields.CharField(max_length=50)
data = fields.JSONField(default={})
class Meta:
table = "agents"
class Chat(models.Model):
id = fields.IntField(pk=True)
user = fields.ForeignKeyField("models.User", related_name="chats")
name = fields.CharField(max_length=255)
description = fields.TextField(null=True)
picture = fields.TextField(null=True)
ordering = fields.IntField(null=True)
last_message_preview = fields.TextField(null=True)
last_message_time = fields.BigIntField(null=True)
class Meta:
table = "chats"
class Membership(models.Model):
id = fields.IntField(pk=True)
agent = fields.ForeignKeyField("models.Agent", related_name="memberships")
chat = fields.ForeignKeyField("models.Chat", related_name="memberships")
class Meta:
table = "memberships"
class Message(models.Model):
id = fields.IntField(pk=True)
chat = fields.ForeignKeyField("models.Chat", related_name="messages")
thread = fields.ForeignKeyField("models.Thread", related_name="messages", null=True)
sender = fields.CharField(max_length=100)
time = fields.BigIntField()
last_edit = fields.BigIntField()
kind = fields.CharField(max_length=50)
subkind = fields.CharField(max_length=50, null=True)
format = fields.CharField(max_length=20, null=True)
content = fields.TextField()
class Meta:
table = "messages"
class Thread(models.Model):
id = fields.IntField(pk=True)
chat = fields.ForeignKeyField("models.Chat", related_name="threads")
name = fields.CharField(max_length=255, null=True)
description = fields.TextField(null=True)
class Meta:
table = "threads"
class AgentResource(models.Model):
id = fields.IntField(pk=True)
agent = fields.ForeignKeyField("models.Agent", related_name="resources")
kind = fields.CharField(max_length=50)
key = fields.CharField(max_length=100)
value = fields.JSONField(null=True)
class Meta:
table = "agent_resources"
class MessageDatum(models.Model):
id = fields.IntField(pk=True)
message = fields.ForeignKeyField("models.Message", related_name="data")
user = fields.ForeignKeyField("models.User", related_name="message_data")
value = fields.JSONField(null=True)
class Meta:
table = "message_data"
"""
class User(Base):
__modelname__ = "user"
__tablename__ = 'users'
@@ -144,4 +258,5 @@ class MessageDatum(Base):
id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True)
message_id: Mapped[int] = mapped_column(ForeignKey("messages.id"))
user_id: Mapped[int] = mapped_column(ForeignKey("users.id"))
value = mapped_column(JSON, default=None)
value = mapped_column(JSON, default=None)
"""
+20 -14
View File
@@ -12,6 +12,8 @@ import nest_asyncio
frontend_path = '../../frontend/dist'
def start_server(app):
nest_asyncio.apply()
server = Quart(__name__, static_folder=frontend_path)
server.secret_key = 'horde-secret'
auth_manager = QuartAuth(server, mode="brearer")
@@ -48,7 +50,7 @@ def start_server(app):
if "username" not in data or "password" not in data:
return "invalid login arguments", 400
user_id = app.verify_user_credentials(data['username'], data['password'])
user_id = await app.verify_user_credentials(data['username'], data['password'])
if user_id == None:
return "bad login credentials", 401
@@ -65,11 +67,11 @@ def start_server(app):
email = data['email']
password = data['password']
user_id = app.create_new_user(username, email, password)
user_id = await app.create_new_user(username, email, password)
if user_id is None:
return "existing or incorrect credentials", 400
app.generate_default_user_data(user_id)
await app.generate_default_user_data(user_id)
token = auth_manager.dump_token(user_id)
return {"token": token}
@@ -84,6 +86,9 @@ def start_server(app):
if 'token' in auth_req:
try:
user_id = auth_manager.load_token(auth_req['token'])
if not await app.verify_user_id(user_id):
await websocket.send(json.dumps({'error': 'invalid credentials'}))
return
except Exception as err:
print("Token authentication failed: ", err)
await websocket.send(json.dumps({'error': 'token authentication failed'}))
@@ -133,7 +138,7 @@ def start_server(app):
elif action == "create-agent":
name = body['name']
agent_id = app.create_agent(
agent_id = await app.create_agent(
user_id=user_id,
name=name,
agentname=body.get('agentname'),
@@ -146,7 +151,7 @@ def start_server(app):
elif action == "update-agent":
agent_id = body['agentId']
app.update_agent(
await app.update_agent(
agent_id=agent_id,
name=body.get('name'),
agentname=body.get('agentname'),
@@ -157,7 +162,7 @@ def start_server(app):
await send({'message': "agent updated", 'agentId': agent_id, 'nonce': nonce})
elif action == "get-agent-config-fields":
await send({'fields': app.get_agent_config_fields(body['agentId']), 'nonce': nonce})
await send({'fields': await app.get_agent_config_fields(body['agentId']), 'nonce': nonce})
elif action == "get-agent-kinds":
await send({'message': "agents retrieved", 'kinds': app.get_agent_kinds(), 'nonce': nonce})
@@ -166,10 +171,11 @@ def start_server(app):
await send({'message': "agent kind retrieved", 'info': app.get_agent_kind_info(body['kind']), 'nonce': nonce})
elif action == "get-agents":
await send({'message': "agents retrieved", 'agents': app.get_agents(user_id), 'nonce': nonce})
#await send({'agents': [], 'nonce': nonce})
await send({'message': "agents retrieved", 'agents': await app.get_agents(user_id), 'nonce': nonce})
elif action == "create-chat":
chat_id = app.create_chat(
chat_id = await app.create_chat(
user_id=user_id,
name=body['name'],
description=body.get('description'),
@@ -179,24 +185,24 @@ def start_server(app):
await send({'message': "chat created", 'chatId': chat_id, 'nonce': nonce})
elif action == "get-chats":
chats_data = app.get_chats(user_id)
chats_data = await app.get_chats(user_id)
await send({'message': "chats retrieved", 'chats': chats_data, 'nonce': nonce})
elif action == "delete-chat":
chat_id = app.delete_chat(
chat_id = await app.delete_chat(
chat_id=body['chatId']
)
await send({'message': "chat deleted", 'chatId': chat_id, 'nonce': nonce})
elif action == "get-chat-messages":
messages_data = app.get_chat_messages(
messages_data = await app.get_chat_messages(
body['chatId'],
earliest_index=body.get('earliestIndex', 0)
)
await send({'message': "chat messages retrieved", 'messages': messages_data, 'nonce': nonce})
elif action == "send-chat-message":
message = app.send_chat_message(
message = await app.send_chat_message(
chat_id=body['chatId'],
thread_id=body.get('threadId'),
content=body['content'],
@@ -207,7 +213,7 @@ def start_server(app):
await send({'message': "message sent", 'messageId': message['id'], 'nonce': nonce})
elif action == "update-chat-message":
message = app.update_chat_message(
message = await app.update_chat_message(
message_id=body['messageId'],
content=body['content'],
kind=body.get('kind'),
@@ -217,7 +223,7 @@ def start_server(app):
await send({'message': "message sent", 'messageId': body['messageId'], 'nonce': nonce})
elif action == "create-chat-thread":
thread_id = app.create_chat_thread(
thread_id = await app.create_chat_thread(
chat_id=body['chatId'],
name=body.get('name'),
)