feat: Tortoise ORM migrations

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)
-"""
+    value = fields.JSONField(null=True)