hub/backend/src/database.py

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.

# 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",
            },
        },
    }

# 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):
    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):
    import temporary as tmp
    if hasattr(value, 'id'):
        if tmp.is_tmp_id(value.id):
            return False

    if isinstance(value, models.Model):
        return True

    print("NOT A DB INSTANCE??? " + repr(value))
    return False

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="")
    # DEPRECATED: Use AgentResources instead
    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 UserResource(models.Model):
#     id = fields.IntField(pk=True)
#     user = fields.ForeignKeyField("models.User", related_name="resources")

#     kind = fields.CharField(max_length=50)
#     key = fields.CharField(max_length=100)
#     value = fields.JSONField(null=True)

#     class Meta:
#         table = "user_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 Model(models.Model):
    id = fields.IntField(pk=True)
    #user = fields.ForeignKeyField("models.User", null=True)
    name = fields.CharField(max_length=255)
    kind = fields.CharField(max_length=50)
    value = fields.JSONField(null=True)