305 lines
12 KiB
Python
305 lines
12 KiB
Python
import datetime
|
|
|
|
from dotenv import load_dotenv
|
|
from nanoid import generate as generate_nanoid
|
|
from pgvector.sqlalchemy import Vector
|
|
from sqlalchemy import (
|
|
BigInteger,
|
|
Boolean,
|
|
CheckConstraint,
|
|
DateTime,
|
|
ForeignKey,
|
|
Identity,
|
|
Index,
|
|
UniqueConstraint,
|
|
text,
|
|
)
|
|
from sqlalchemy.dialects.postgresql import JSONB, TEXT
|
|
from sqlalchemy.orm import Mapped, mapped_column, relationship
|
|
from sqlalchemy.sql import func
|
|
|
|
from .db import Base
|
|
|
|
load_dotenv()
|
|
|
|
|
|
class App(Base):
|
|
__tablename__ = "apps"
|
|
id: Mapped[int] = mapped_column(
|
|
BigInteger, Identity(), primary_key=True, index=True, autoincrement=True
|
|
)
|
|
public_id: Mapped[str] = mapped_column(
|
|
TEXT, index=True, unique=True, default=generate_nanoid
|
|
)
|
|
name: Mapped[str] = mapped_column(TEXT, index=True, unique=True)
|
|
users = relationship("User", back_populates="app")
|
|
created_at: Mapped[datetime.datetime] = mapped_column(
|
|
DateTime(timezone=True), index=True, default=func.now()
|
|
)
|
|
h_metadata: Mapped[dict] = mapped_column("metadata", JSONB, default={})
|
|
|
|
__table_args__ = (
|
|
CheckConstraint("length(public_id) = 21", name="public_id_length"),
|
|
CheckConstraint("length(name) <= 512", name="name_length"),
|
|
CheckConstraint("public_id ~ '^[A-Za-z0-9_-]+$'", name="public_id_format"),
|
|
)
|
|
|
|
|
|
class User(Base):
|
|
__tablename__ = "users"
|
|
id: Mapped[int] = mapped_column(
|
|
BigInteger, Identity(), primary_key=True, index=True, autoincrement=True
|
|
)
|
|
public_id: Mapped[str] = mapped_column(
|
|
TEXT, index=True, unique=True, default=generate_nanoid
|
|
)
|
|
name: Mapped[str] = mapped_column(TEXT, index=True)
|
|
h_metadata: Mapped[dict] = mapped_column("metadata", JSONB, default={})
|
|
created_at: Mapped[datetime.datetime] = mapped_column(
|
|
DateTime(timezone=True), index=True, default=func.now()
|
|
)
|
|
app_id: Mapped[str] = mapped_column(ForeignKey("apps.public_id"), index=True)
|
|
app = relationship("App", back_populates="users")
|
|
sessions = relationship("Session", back_populates="user")
|
|
collections = relationship("Collection", back_populates="user")
|
|
metamessages = relationship("Metamessage", back_populates="user")
|
|
|
|
__table_args__ = (
|
|
UniqueConstraint("name", "app_id", name="unique_name_app_user"),
|
|
CheckConstraint("length(public_id) = 21", name="public_id_length"),
|
|
CheckConstraint("length(name) <= 512", name="name_length"),
|
|
CheckConstraint("public_id ~ '^[A-Za-z0-9_-]+$'", name="public_id_format"),
|
|
Index("idx_users_app_lookup", "app_id", "public_id"),
|
|
)
|
|
|
|
def __repr__(self) -> str:
|
|
return f"User(id={self.id}, app_id={self.app_id}, public_id={self.public_id} created_at={self.created_at}, h_metadata={self.h_metadata})"
|
|
|
|
|
|
class Session(Base):
|
|
__tablename__ = "sessions"
|
|
id: Mapped[int] = mapped_column(
|
|
BigInteger, Identity(), primary_key=True, index=True, autoincrement=True
|
|
)
|
|
public_id: Mapped[str] = mapped_column(
|
|
TEXT, index=True, unique=True, default=generate_nanoid
|
|
)
|
|
is_active: Mapped[bool] = mapped_column(default=True)
|
|
h_metadata: Mapped[dict] = mapped_column("metadata", JSONB, default={})
|
|
created_at: Mapped[datetime.datetime] = mapped_column(
|
|
DateTime(timezone=True), index=True, default=func.now()
|
|
)
|
|
messages = relationship("Message", back_populates="session")
|
|
metamessages = relationship("Metamessage", back_populates="session")
|
|
user_id: Mapped[str] = mapped_column(ForeignKey("users.public_id"), index=True)
|
|
app_id: Mapped[str] = mapped_column(ForeignKey("apps.public_id"), index=True)
|
|
user = relationship("User", back_populates="sessions")
|
|
|
|
__table_args__ = (
|
|
CheckConstraint("length(public_id) = 21", name="public_id_length"),
|
|
CheckConstraint("public_id ~ '^[A-Za-z0-9_-]+$'", name="public_id_format"),
|
|
Index("idx_sessions_user_lookup", "user_id", "public_id"),
|
|
)
|
|
|
|
def __repr__(self) -> str:
|
|
return f"Session(id={self.id}, user_id={self.user_id}, is_active={self.is_active}, created_at={self.created_at}, h_metadata={self.h_metadata})"
|
|
|
|
|
|
class Message(Base):
|
|
__tablename__ = "messages"
|
|
id: Mapped[int] = mapped_column(
|
|
BigInteger, Identity(), primary_key=True, index=True, autoincrement=True
|
|
)
|
|
public_id: Mapped[str] = mapped_column(
|
|
TEXT, index=True, unique=True, default=generate_nanoid
|
|
)
|
|
session_id: Mapped[str] = mapped_column(
|
|
ForeignKey("sessions.public_id"), index=True
|
|
)
|
|
is_user: Mapped[bool]
|
|
content: Mapped[str] = mapped_column(TEXT)
|
|
h_metadata: Mapped[dict] = mapped_column("metadata", JSONB, default={})
|
|
|
|
created_at: Mapped[datetime.datetime] = mapped_column(
|
|
DateTime(timezone=True), index=True, default=func.now()
|
|
)
|
|
session = relationship("Session", back_populates="messages")
|
|
metamessages = relationship("Metamessage", back_populates="message")
|
|
user_id: Mapped[str] = mapped_column(ForeignKey("users.public_id"), index=True)
|
|
app_id: Mapped[str] = mapped_column(ForeignKey("apps.public_id"), index=True)
|
|
|
|
__table_args__ = (
|
|
CheckConstraint("length(public_id) = 21", name="public_id_length"),
|
|
CheckConstraint("public_id ~ '^[A-Za-z0-9_-]+$'", name="public_id_format"),
|
|
CheckConstraint("length(content) <= 65535", name="content_length"),
|
|
Index(
|
|
"idx_messages_session_lookup",
|
|
"session_id",
|
|
"id",
|
|
postgresql_include=["public_id", "is_user", "created_at"],
|
|
),
|
|
)
|
|
|
|
def __repr__(self) -> str:
|
|
return f"Message(id={self.id}, session_id={self.session_id}, is_user={self.is_user}, content={self.content[10:]})"
|
|
|
|
|
|
class Metamessage(Base):
|
|
__tablename__ = "metamessages"
|
|
id: Mapped[int] = mapped_column(
|
|
BigInteger, Identity(), primary_key=True, index=True, autoincrement=True
|
|
)
|
|
public_id: Mapped[str] = mapped_column(
|
|
TEXT, index=True, unique=True, default=generate_nanoid
|
|
)
|
|
label: Mapped[str] = mapped_column(TEXT, index=True)
|
|
content: Mapped[str] = mapped_column(TEXT)
|
|
|
|
# Foreign keys - message_id is now optional
|
|
user_id: Mapped[str] = mapped_column(ForeignKey("users.public_id"), index=True)
|
|
app_id: Mapped[str] = mapped_column(ForeignKey("apps.public_id"), index=True)
|
|
session_id: Mapped[str | None] = mapped_column(
|
|
ForeignKey("sessions.public_id"), index=True, nullable=True
|
|
)
|
|
message_id: Mapped[str | None] = mapped_column(
|
|
ForeignKey("messages.public_id"), index=True, nullable=True
|
|
)
|
|
|
|
# Relationships
|
|
user = relationship("User", back_populates="metamessages")
|
|
session = relationship("Session", back_populates="metamessages")
|
|
message = relationship("Message", back_populates="metamessages")
|
|
|
|
created_at: Mapped[datetime.datetime] = mapped_column(
|
|
DateTime(timezone=True), index=True, default=func.now()
|
|
)
|
|
h_metadata: Mapped[dict] = mapped_column("metadata", JSONB, default={})
|
|
|
|
__table_args__ = (
|
|
CheckConstraint("length(public_id) = 21", name="public_id_length"),
|
|
CheckConstraint("public_id ~ '^[A-Za-z0-9_-]+$'", name="public_id_format"),
|
|
CheckConstraint("length(content) <= 65535", name="content_length"),
|
|
CheckConstraint("length(label) <= 512", name="label_length"),
|
|
# Added constraints to ensure consistency
|
|
CheckConstraint(
|
|
"(message_id IS NULL) OR (session_id IS NOT NULL)",
|
|
name="message_requires_session",
|
|
),
|
|
# Keep existing index
|
|
Index(
|
|
"idx_metamessages_lookup",
|
|
"label",
|
|
text("id DESC"),
|
|
postgresql_include=["public_id", "message_id", "created_at"],
|
|
),
|
|
# Indices for user, session, and message lookups
|
|
Index(
|
|
"idx_metamessages_user_lookup",
|
|
"user_id",
|
|
"label",
|
|
text("id DESC"),
|
|
),
|
|
Index(
|
|
"idx_metamessages_session_lookup",
|
|
"session_id",
|
|
"label",
|
|
text("id DESC"),
|
|
),
|
|
Index(
|
|
"idx_metamessages_message_lookup",
|
|
"message_id",
|
|
"label",
|
|
text("id DESC"),
|
|
),
|
|
)
|
|
|
|
def __repr__(self) -> str:
|
|
return f"Metamessages(id={self.id}, user_id={self.user_id}, session_id={self.session_id}, message_id={self.message_id}, label={self.label})"
|
|
|
|
|
|
class Collection(Base):
|
|
__tablename__ = "collections"
|
|
|
|
id: Mapped[int] = mapped_column(
|
|
BigInteger, Identity(), primary_key=True, index=True, autoincrement=True
|
|
)
|
|
public_id: Mapped[str] = mapped_column(
|
|
TEXT, index=True, unique=True, default=generate_nanoid
|
|
)
|
|
name: Mapped[str] = mapped_column(TEXT, index=True)
|
|
created_at: Mapped[datetime.datetime] = mapped_column(
|
|
DateTime(timezone=True), index=True, default=func.now()
|
|
)
|
|
h_metadata: Mapped[dict] = mapped_column("metadata", JSONB, default={})
|
|
documents = relationship(
|
|
"Document", back_populates="collection", cascade="all, delete, delete-orphan"
|
|
)
|
|
user = relationship("User", back_populates="collections")
|
|
user_id: Mapped[str] = mapped_column(
|
|
TEXT, ForeignKey("users.public_id"), index=True
|
|
)
|
|
app_id: Mapped[str] = mapped_column(ForeignKey("apps.public_id"), index=True)
|
|
|
|
__table_args__ = (
|
|
UniqueConstraint("name", "user_id", name="unique_name_collection_user"),
|
|
CheckConstraint("length(public_id) = 21", name="public_id_length"),
|
|
CheckConstraint("public_id ~ '^[A-Za-z0-9_-]+$'", name="public_id_format"),
|
|
CheckConstraint("length(name) <= 512", name="name_length"),
|
|
)
|
|
|
|
|
|
class Document(Base):
|
|
__tablename__ = "documents"
|
|
id: Mapped[int] = mapped_column(
|
|
BigInteger, Identity(), primary_key=True, index=True, autoincrement=True
|
|
)
|
|
public_id: Mapped[str] = mapped_column(
|
|
TEXT, index=True, unique=True, default=generate_nanoid
|
|
)
|
|
h_metadata: Mapped[dict] = mapped_column("metadata", JSONB, default={})
|
|
content: Mapped[str] = mapped_column(TEXT)
|
|
embedding = mapped_column(Vector(1536))
|
|
created_at: Mapped[datetime.datetime] = mapped_column(
|
|
DateTime(timezone=True), index=True, default=func.now()
|
|
)
|
|
|
|
collection_id: Mapped[str] = mapped_column(
|
|
TEXT, ForeignKey("collections.public_id"), index=True
|
|
)
|
|
user_id: Mapped[str] = mapped_column(ForeignKey("users.public_id"), index=True)
|
|
app_id: Mapped[str] = mapped_column(ForeignKey("apps.public_id"), index=True)
|
|
collection = relationship("Collection", back_populates="documents")
|
|
|
|
__table_args__ = (
|
|
CheckConstraint("length(public_id) = 21", name="public_id_length"),
|
|
CheckConstraint("length(content) <= 65535", name="content_length"),
|
|
CheckConstraint("public_id ~ '^[A-Za-z0-9_-]+$'", name="public_id_format"),
|
|
# HNSW index on embedding column
|
|
Index(
|
|
"idx_documents_embedding_hnsw",
|
|
"embedding",
|
|
postgresql_using="hnsw", # HNSW index type
|
|
postgresql_with={"m": 16, "ef_construction": 64}, # HNSW parameters
|
|
postgresql_ops={"embedding": "vector_cosine_ops"}, # Cosine distance operator
|
|
),
|
|
)
|
|
|
|
|
|
class QueueItem(Base):
|
|
__tablename__ = "queue"
|
|
id: Mapped[int] = mapped_column(
|
|
BigInteger, Identity(), primary_key=True, autoincrement=True
|
|
)
|
|
session_id: Mapped[int] = mapped_column(ForeignKey("sessions.id"), index=True)
|
|
payload: Mapped[dict] = mapped_column(JSONB, nullable=False)
|
|
processed: Mapped[bool] = mapped_column(Boolean, default=False)
|
|
|
|
|
|
class ActiveQueueSession(Base):
|
|
__tablename__ = "active_queue_sessions"
|
|
|
|
session_id: Mapped[int] = mapped_column(BigInteger, primary_key=True, index=True)
|
|
last_updated: Mapped[datetime.datetime] = mapped_column(
|
|
DateTime(timezone=True), default=func.now(), onupdate=func.now()
|
|
)
|