feat: Tortoise ORM migrations

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

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_token and replicate.Client(api_token=replicate_token) or None
-            self.replicate_client = replicate.Client(api_token=replicate_api_token)
+        self.openrouter_client = openrouter_token and OpenRouter(api_key=openrouter_token) or None
-        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.agent_classes = [
             agents_mod.ExpertAgent,
@@ -350,10 +343,10 @@ class App(ChatMixin):
 # --- Main Entry Point ---
 
 async def main():
-    if len(sys.argv) > 1:
+    # Get data directory from environment variable
-        data_directory = sys.argv[1]
+    data_directory = os.getenv("HORDE_DATA_DIR")
-    else:
+    if not data_directory:
-        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

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):
+# Creates a Tortoise config dict using files from the passed directory
-    pass
+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):
+# Only used for tortoise tool
-    sql_engine = create_engine(f"sqlite:///{data_directory}/horde.db")
+TORTOISE_ORM=tortoise_config(os.getenv("HORDE_DATA_DIR"))
-    Base.metadata.create_all(sql_engine)
 
-    return sql_engine
+# 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"
+    from tortoise.migrations.api import migrate
-
+    config = tortoise_config(data_directory)
-    # Setup Tortoise
+    await Tortoise.init(config=config)
-    await Tortoise.init(
+    # Apply migrations automatically
-        db_url=db_url,
+    await migrate(config=config)
-        modules={'models': ['database']}
-    )
-    await Tortoise.generate_schemas()
 
 # Checks whether an object is an instance of any database model
 def is_instance (value):
@@ -146,140 +151,3 @@ class Model(models.Model):
     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)
-"""

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'],
+        ),
+    ]

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 $HORDE_DATA_DIR
-    mkdir -p $data_dir
+    python backend/src/app.py
-    python backend/src/app.py $data_dir
     ;;
   "cli")
+    # deprecated lol
     node cli/main.js "$2"
     ;;
   "test")
-    # Setup test data directory
+    export HORDE_DATA_DIR=$(realpath ./test-data)
-    data_dir=$(realpath ./test-data)
+    mkdir -p $HORDE_DATA_DIR
-    mkdir -p $data_dir
+    rm -rf $HORDE_DATA_DIR/*
-    rm -rf $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)
+      export HORDE_DATA_DIR=$(realpath $1)
-      python backend/src/app.py $data_dir
+      python backend/src/app.py
     else
       echo "Error: unknown command or directory '$1'"
     fi