From 25b5fbd22a8dfc9ede0adffa597340316dc440bc Mon Sep 17 00:00:00 2001 From: Shivering Ocelot Date: Sat, 30 May 2026 10:55:31 -0500 Subject: [PATCH] feat: Tortoise ORM migrations --- backend/requirements.txt | 2 +- backend/src/app.py | 26 ++-- backend/src/database.py | 182 ++++--------------------- backend/src/migrations/0001_initial.py | 129 ++++++++++++++++++ backend/src/migrations/__init__.py | 0 run.sh | 29 ++-- 6 files changed, 184 insertions(+), 184 deletions(-) create mode 100644 backend/src/migrations/0001_initial.py create mode 100644 backend/src/migrations/__init__.py diff --git a/backend/requirements.txt b/backend/requirements.txt index eb9bf90..68c51f4 100644 --- a/backend/requirements.txt +++ b/backend/requirements.txt @@ -6,6 +6,6 @@ nest_asyncio replicate qdrant-client[fastembed] ollama -tortoise-orm +tortoise-orm>=1.0.0 llama-cpp-python openrouter \ No newline at end of file diff --git a/backend/src/app.py b/backend/src/app.py index 53ff0b1..f452eb8 100644 --- a/backend/src/app.py +++ b/backend/src/app.py @@ -22,7 +22,7 @@ from chat import ChatView, Message, MESSAGE_PAGE_SIZE from app_mixins import ChatMixin class App(ChatMixin): - def __init__(self, data_directory, replicate_api_token=None): + def __init__(self, data_directory, replicate_token=None, openrouter_token=None): self.data_directory = data_directory self.user_connections = {} self.remote_connections = {} @@ -34,16 +34,9 @@ class App(ChatMixin): # Clients self.qdrant_client = QdrantClient(path=f"{self.data_directory}/qdrant_data") - if replicate_api_token: - self.replicate_client = replicate.Client(api_token=replicate_api_token) - else: - self.replicate_client = None - - openrouter_api_key = os.getenv("OPENROUTER_API_KEY") - if openrouter_api_key: - self.openrouter_client = OpenRouter(api_key=openrouter_api_key) - else: - self.openrouter_client = None + self.replicate_client = replicate_token and replicate.Client(api_token=replicate_token) or None + self.openrouter_client = openrouter_token and OpenRouter(api_key=openrouter_token) or None + self.agent_classes = [ agents_mod.ExpertAgent, @@ -350,10 +343,10 @@ class App(ChatMixin): # --- Main Entry Point --- async def main(): - if len(sys.argv) > 1: - data_directory = sys.argv[1] - else: - data_directory = "." + # Get data directory from environment variable + data_directory = os.getenv("HORDE_DATA_DIR") + if not data_directory: + raise ValueError("HORDE_DATA_DIR environment variable must be set") # 1. Initialize Database First await db.init_database(data_directory) @@ -362,7 +355,8 @@ async def main(): # 2. Initialize App horde_app = App( data_directory=data_directory, - replicate_api_token=os.getenv('REPLICATE_API_TOKEN'), + replicate_token=os.getenv('REPLICATE_API_TOKEN'), + openrouter_token=os.getenv("OPENROUTER_API_KEY") ) # 3. Start Server diff --git a/backend/src/database.py b/backend/src/database.py index c7d0800..48520ed 100644 --- a/backend/src/database.py +++ b/backend/src/database.py @@ -1,33 +1,38 @@ -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 +import os # These types don't represent all objects in Horde. # Agents and Chats can be created for just a session, in which case, they # are not saved on the database, as aren't all their associated values of # other classes. -class Base(DeclarativeBase): - pass +# Creates a Tortoise config dict using files from the passed directory +def tortoise_config(data_directory): + if data_directory is None: raise Exception("Missing data directory") + db_url = f"sqlite:///{data_directory}/horde.db" + return { + "connections": {"default": db_url}, + "apps": { + "models": { + "models": ["database"], + "default_connection": "default", + "migrations": "migrations", + }, + }, + } -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 +# Only used for tortoise tool +TORTOISE_ORM=tortoise_config(os.getenv("HORDE_DATA_DIR")) + +# Considering creating a migration script "makemigration.py" or similar, that will load the config and make migration 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() + from tortoise.migrations.api import migrate + config = tortoise_config(data_directory) + await Tortoise.init(config=config) + # Apply migrations automatically + await migrate(config=config) # Checks whether an object is an instance of any database model def is_instance (value): @@ -145,141 +150,4 @@ class Model(models.Model): id = fields.IntField(pk=True) name = fields.CharField(max_length=255) kind = fields.CharField(max_length=50) - value = fields.JSONField(null=True) - -""" -class User(Base): - __modelname__ = "user" - __tablename__ = 'users' - - id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True) - - email: Mapped[str] - username: Mapped[str] - # Stores the plain password. This must be changed eventually to store a salted hash instead. - password: Mapped[str] - -class Agent(Base): - __modelname__ = "agent" - __tablename__ = 'agents' - - id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True) - user_id: Mapped[int] = mapped_column(ForeignKey("users.id")) - - agentname: Mapped[Optional[str]] - visibility: Mapped[int] = mapped_column(default=0) - # Comma separated list of permission symbols - permissions: Mapped[str] = mapped_column(default="") - main_chat_id: Mapped[int] = mapped_column(ForeignKey("chats.id")) - - name: Mapped[str] - description: Mapped[Optional[str]] - picture: Mapped[Optional[str]] - ordering: Mapped[Optional[int]] - - # Agents can be any of a set of agent classes - kind: Mapped[str] - # Stores the agent data specific to its class - data = mapped_column(JSON, default={}) - -class Chat(Base): - __modelname__ = 'chat' - __tablename__ = 'chats' - - id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True) - user_id: Mapped[int] = mapped_column(ForeignKey("users.id")) - - name: Mapped[str] - description: Mapped[Optional[str]] - picture: Mapped[Optional[str]] - ordering: Mapped[Optional[int]] - - # For indexing purposes, some information is taken from the last chat message. - # Preview is the (possibly cropped) content, time is the timestamp of the message. - last_message_preview: Mapped[Optional[str]] - last_message_time: Mapped[Optional[int]] - - # Chat Controller defines when to forward messages to agents and users, and what messages each can see - # Controllers are: default (all messages to all participants), mention (notify agents by specific mentions), active (llm agent selecting next agent to respond), arena - #controller_kind: Mapped[str] - #controller_data = mapped_column(JSON, default={}) - -# Stores membership of agents in chats. If a record exists linking an agent and a chat, -# it means the agent participates in that chat. -class Membership(Base): - __modelname__ = 'membership' - __tablename__ = 'memberships' - - id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True) - agent_id: Mapped[int] = mapped_column(ForeignKey("agents.id")) - chat_id: Mapped[int] = mapped_column(ForeignKey("chats.id")) - -class Message(Base): - __modelname__ = 'message' - __tablename__ = 'messages' - - id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True) - chat_id: Mapped[int] = mapped_column(ForeignKey("chats.id")) - thread_id: Mapped[Optional[int]] = mapped_column(ForeignKey("threads.id")) - - # Identifier of the sender: user / system / user-$ID / agent-$ID - sender: Mapped[str] - # Timestamp when this message was sent - time: Mapped[int] - # Timestamp of this message's last edit - last_edit: Mapped[int] - - # Predefined set of message types: - # message, log, memory, file, note - kind: Mapped[str] - # Each kind has a set of subkinds, and a default for empty values - # message: (empty) - # log: debug, info - # memory: (depends on the agent) - # file: (empty) - # note: public, private, pinned - subkind: Mapped[Optional[str]] - # For the kinds message, log and note, a format can be specified, - # empty defaults to 'text' and can also be 'markdown', 'csv' - format: Mapped[Optional[str]] - content: Mapped[str] - -# Represents a temporary topic or stage in a conversation -class Thread(Base): - __modelname__ = 'thread' - __tablename__ = 'threads' - - id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True) - chat_id: Mapped[int] = mapped_column(ForeignKey("chats.id")) - - # Optional name and descriptions of the thread - name: Mapped[Optional[str]] - description: Mapped[Optional[str]] - - # Threads can form trees where threads continue each other - #previous_thread_id: Mapped[Optional[int]] = mapped_column(ForeignKey("threads.id")) - -# A persistent object or value associated with an agent -class AgentResource(Base): - __modelname__ = 'agent_resource' - __tablename__ = 'agent_resources' - - id: Mapped[int] = mapped_column(primary_key=True, autoincrement=True) - agent_id: Mapped[int] = mapped_column(ForeignKey("agents.id")) - - # Kind is the class of the resource - # Key is the specific identifier within the class - # Multiple resources can share key - kind: Mapped[str] - key: Mapped[str] - value = mapped_column(JSON, default=None) - -class MessageDatum(Base): - __modelname__ = 'message_datum' - __tablename__ = 'message_data' - - 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 = fields.JSONField(null=True) \ No newline at end of file diff --git a/backend/src/migrations/0001_initial.py b/backend/src/migrations/0001_initial.py new file mode 100644 index 0000000..5a96d21 --- /dev/null +++ b/backend/src/migrations/0001_initial.py @@ -0,0 +1,129 @@ +from tortoise import migrations +from tortoise.migrations import operations as ops +from orjson import loads +from tortoise.fields.base import OnDelete +from tortoise.fields.data import JSON_DUMPS +from tortoise import fields + +class Migration(migrations.Migration): + initial = True + + operations = [ + ops.CreateModel( + name='Model', + fields=[ + ('id', fields.IntField(generated=True, primary_key=True, unique=True, db_index=True)), + ('name', fields.CharField(max_length=255)), + ('kind', fields.CharField(max_length=50)), + ('value', fields.JSONField(null=True, encoder=JSON_DUMPS, decoder=loads)), + ], + options={'table': 'model', 'app': 'models', 'pk_attr': 'id'}, + bases=['Model'], + ), + ops.CreateModel( + name='User', + fields=[ + ('id', fields.IntField(generated=True, primary_key=True, unique=True, db_index=True)), + ('email', fields.CharField(max_length=255)), + ('username', fields.CharField(max_length=100)), + ('password', fields.CharField(max_length=255)), + ], + options={'table': 'users', 'app': 'models', 'pk_attr': 'id'}, + bases=['Model'], + ), + ops.CreateModel( + name='Chat', + fields=[ + ('id', fields.IntField(generated=True, primary_key=True, unique=True, db_index=True)), + ('user', fields.ForeignKeyField('models.User', source_field='user_id', db_constraint=True, to_field='id', related_name='chats', on_delete=OnDelete.CASCADE)), + ('name', fields.CharField(max_length=255)), + ('description', fields.TextField(null=True, unique=False)), + ('picture', fields.TextField(null=True, unique=False)), + ('ordering', fields.IntField(null=True)), + ('last_message_preview', fields.TextField(null=True, unique=False)), + ('last_message_time', fields.BigIntField(null=True)), + ], + options={'table': 'chats', 'app': 'models', 'pk_attr': 'id'}, + bases=['Model'], + ), + ops.CreateModel( + name='Agent', + fields=[ + ('id', fields.IntField(generated=True, primary_key=True, unique=True, db_index=True)), + ('user', fields.ForeignKeyField('models.User', source_field='user_id', db_constraint=True, to_field='id', related_name='agents', on_delete=OnDelete.CASCADE)), + ('agentname', fields.CharField(null=True, max_length=100)), + ('visibility', fields.IntField(default=0)), + ('permissions', fields.TextField(default='', unique=False)), + ('main_chat', fields.ForeignKeyField('models.Chat', source_field='main_chat_id', null=True, db_constraint=True, to_field='id', related_name='main_agents', on_delete=OnDelete.CASCADE)), + ('name', fields.CharField(max_length=255)), + ('description', fields.TextField(null=True, unique=False)), + ('picture', fields.TextField(null=True, unique=False)), + ('ordering', fields.IntField(null=True)), + ('kind', fields.CharField(max_length=50)), + ('data', fields.JSONField(default={}, encoder=JSON_DUMPS, decoder=loads)), + ], + options={'table': 'agents', 'app': 'models', 'pk_attr': 'id'}, + bases=['Model'], + ), + ops.CreateModel( + name='AgentResource', + fields=[ + ('id', fields.IntField(generated=True, primary_key=True, unique=True, db_index=True)), + ('agent', fields.ForeignKeyField('models.Agent', source_field='agent_id', db_constraint=True, to_field='id', related_name='resources', on_delete=OnDelete.CASCADE)), + ('kind', fields.CharField(max_length=50)), + ('key', fields.CharField(max_length=100)), + ('value', fields.JSONField(null=True, encoder=JSON_DUMPS, decoder=loads)), + ], + options={'table': 'agent_resources', 'app': 'models', 'pk_attr': 'id'}, + bases=['Model'], + ), + ops.CreateModel( + name='Membership', + fields=[ + ('id', fields.IntField(generated=True, primary_key=True, unique=True, db_index=True)), + ('agent', fields.ForeignKeyField('models.Agent', source_field='agent_id', db_constraint=True, to_field='id', related_name='memberships', on_delete=OnDelete.CASCADE)), + ('chat', fields.ForeignKeyField('models.Chat', source_field='chat_id', db_constraint=True, to_field='id', related_name='memberships', on_delete=OnDelete.CASCADE)), + ], + options={'table': 'memberships', 'app': 'models', 'pk_attr': 'id'}, + bases=['Model'], + ), + ops.CreateModel( + name='Thread', + fields=[ + ('id', fields.IntField(generated=True, primary_key=True, unique=True, db_index=True)), + ('chat', fields.ForeignKeyField('models.Chat', source_field='chat_id', db_constraint=True, to_field='id', related_name='threads', on_delete=OnDelete.CASCADE)), + ('name', fields.CharField(null=True, max_length=255)), + ('description', fields.TextField(null=True, unique=False)), + ], + options={'table': 'threads', 'app': 'models', 'pk_attr': 'id'}, + bases=['Model'], + ), + ops.CreateModel( + name='Message', + fields=[ + ('id', fields.IntField(generated=True, primary_key=True, unique=True, db_index=True)), + ('chat', fields.ForeignKeyField('models.Chat', source_field='chat_id', db_constraint=True, to_field='id', related_name='messages', on_delete=OnDelete.CASCADE)), + ('thread', fields.ForeignKeyField('models.Thread', source_field='thread_id', null=True, db_constraint=True, to_field='id', related_name='messages', on_delete=OnDelete.CASCADE)), + ('sender', fields.CharField(max_length=100)), + ('time', fields.BigIntField()), + ('last_edit', fields.BigIntField()), + ('kind', fields.CharField(max_length=50)), + ('subkind', fields.CharField(null=True, max_length=50)), + ('format', fields.CharField(null=True, max_length=20)), + ('content', fields.TextField(unique=False)), + ], + options={'table': 'messages', 'app': 'models', 'pk_attr': 'id'}, + bases=['Model'], + ), + ops.CreateModel( + name='MessageDatum', + fields=[ + ('id', fields.IntField(generated=True, primary_key=True, unique=True, db_index=True)), + ('message', fields.ForeignKeyField('models.Message', source_field='message_id', db_constraint=True, to_field='id', related_name='data', on_delete=OnDelete.CASCADE)), + ('user', fields.ForeignKeyField('models.User', source_field='user_id', db_constraint=True, to_field='id', related_name='message_data', on_delete=OnDelete.CASCADE)), + ('value', fields.JSONField(null=True, encoder=JSON_DUMPS, decoder=loads)), + ], + options={'table': 'message_data', 'app': 'models', 'pk_attr': 'id'}, + bases=['Model'], + ), + ] diff --git a/backend/src/migrations/__init__.py b/backend/src/migrations/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/run.sh b/run.sh index f4f45f4..755371a 100644 --- a/run.sh +++ b/run.sh @@ -10,23 +10,24 @@ else echo "Write a .env file with the variables: REPLICATE_API_TOKEN" fi +export HORDE_DATA_DIR=$(realpath ./local-data) + case "$1" in "") - data_dir=$(realpath ./local-data) - mkdir -p $data_dir - python backend/src/app.py $data_dir + mkdir -p $HORDE_DATA_DIR + python backend/src/app.py ;; "cli") + # deprecated lol node cli/main.js "$2" ;; "test") - # Setup test data directory - data_dir=$(realpath ./test-data) - mkdir -p $data_dir - rm -rf $data_dir/* + export HORDE_DATA_DIR=$(realpath ./test-data) + mkdir -p $HORDE_DATA_DIR + rm -rf $HORDE_DATA_DIR/* # Start backend in the background - python backend/src/app.py $data_dir & + python backend/src/app.py & BACKEND_PID=$! echo "Started backend: $BACKEND_PID" sleep 5 @@ -47,10 +48,18 @@ case "$1" in flutter build web done ;; + + # Create migrations based on the current schema in database.py + "migration") + cd backend/src + tortoise -c database.TORTOISE_ORM makemigrations + ;; + + # Execute with custom home directory if it exists *) if [ -d "$1" ]; then - data_dir=$(realpath $1) - python backend/src/app.py $data_dir + export HORDE_DATA_DIR=$(realpath $1) + python backend/src/app.py else echo "Error: unknown command or directory '$1'" fi