feat: Tortoise ORM migrations
This commit is contained in:
Shivering Ocelot committed
1 parent
50cd677515
commit
25b5fbd22a
6 files changed
+184
-184
No files matched your search
+10
-16
@@ -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
|
||||
|
||||
+25
-157
@@ -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)
|
||||
@@ -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'],
|
||||
),
|
||||
]
|
||||
Whitespace-only changes.
Reference in new issue
Block a user