130 lines
4.4 KiB
Python
130 lines
4.4 KiB
Python
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
|
|
|
|
class Base(DeclarativeBase):
|
|
pass
|
|
|
|
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
|
|
|
|
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]]
|
|
|
|
# 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"))
|
|
|
|
# Timestamp of this message
|
|
time: Mapped[int]
|
|
# Identifier of the sender: user / system / user-$ID / agent-$ID
|
|
sender: Mapped[str]
|
|
|
|
# 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]]
|
|
|
|
# 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)
|
|
|