diff --git a/backend/requirements.txt b/backend/requirements.txt index d47cc7b..1d63699 100644 --- a/backend/requirements.txt +++ b/backend/requirements.txt @@ -5,4 +5,5 @@ sqlalchemy nest_asyncio replicate qdrant-client[fastembed] -ollama \ No newline at end of file +ollama +tortoise-orm \ No newline at end of file diff --git a/backend/src/agent.py b/backend/src/agent.py index ac377e3..6e8c111 100644 --- a/backend/src/agent.py +++ b/backend/src/agent.py @@ -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, diff --git a/backend/src/agents/character.py b/backend/src/agents/character.py index 6705c59..5706464 100644 --- a/backend/src/agents/character.py +++ b/backend/src/agents/character.py @@ -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'], diff --git a/backend/src/agents/expert.py b/backend/src/agents/expert.py index 18cb2eb..25ba74d 100644 --- a/backend/src/agents/expert.py +++ b/backend/src/agents/expert.py @@ -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, diff --git a/backend/src/app.py b/backend/src/app.py index 33617ac..aa4b9f3 100644 --- a/backend/src/app.py +++ b/backend/src/app.py @@ -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) \ No newline at end of file + + # 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()) \ No newline at end of file diff --git a/backend/src/app_mixins/chat.py b/backend/src/app_mixins/chat.py index 51054c6..63ba66c 100644 --- a/backend/src/app_mixins/chat.py +++ b/backend/src/app_mixins/chat.py @@ -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)) \ No newline at end of file diff --git a/backend/src/chat.py b/backend/src/chat.py index cf1bafd..e226bff 100644 --- a/backend/src/chat.py +++ b/backend/src/chat.py @@ -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: diff --git a/backend/src/database.py b/backend/src/database.py index 7a3e6a5..ec2bfdf 100644 --- a/backend/src/database.py +++ b/backend/src/database.py @@ -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) \ No newline at end of file + value = mapped_column(JSON, default=None) +""" \ No newline at end of file diff --git a/backend/src/server.py b/backend/src/server.py index 3cc3b71..f97f4df 100644 --- a/backend/src/server.py +++ b/backend/src/server.py @@ -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'), ) diff --git a/frontend/src/App.vue b/frontend/src/App.vue index 4679cf8..b7040c8 100644 --- a/frontend/src/App.vue +++ b/frontend/src/App.vue @@ -73,14 +73,14 @@ backend.authResult().then(handleAuthResult);