Tortoise based database models

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)
+    value = mapped_column(JSON, default=None)
+"""