from datetime import datetime from enum import StrEnum, auto from typing import TypedDict from uuid import UUID, uuid4 from sqlalchemy import TIMESTAMP, Column from sqlalchemy.dialects.postgresql import JSONB from sqlalchemy.orm import Mapped from sqlalchemy.sql.functions import now from sqlmodel import Field, Relationship, SQLModel, col, func class DialogueEntry(TypedDict): speaker: str text: str start_time: float end_time: float # Create factory functions for columns to avoid reusing column objects def created_datetime_column(): return Column(TIMESTAMP(timezone=True), nullable=False, server_default=now(), default=None) def updated_datetime_column(): return Column(TIMESTAMP(timezone=True), nullable=False, server_default=now(), default=None) class BaseTableMixin(SQLModel): # Note, we can't add created/updated_datetime Columns here, as each table needs its own instance of these Columns model_config = { # noqa: RUF012 "from_attributes": True, } id: UUID = Field( default_factory=uuid4, primary_key=True, sa_column_kwargs={"server_default": func.gen_random_uuid()} ) class JobStatus(StrEnum): AWAITING_START = auto() IN_PROGRESS = auto() COMPLETED = auto() FAILED = auto() class ContentSource(StrEnum): MANUAL_EDIT = auto() AI_EDIT = auto() INITIAL_GENERATION = auto() class MinuteVersion(BaseTableMixin, table=True): __tablename__ = "minute_version" created_datetime: datetime = Field(sa_column=created_datetime_column(), default=None) updated_datetime: datetime = Field(sa_column=updated_datetime_column(), default=None) minute_id: UUID = Field(foreign_key="minute.id", ondelete="CASCADE") minute: Mapped["Minute"] = Relationship(back_populates="minute_versions") hallucinations: list["Hallucination"] = Relationship(back_populates="minute_version", cascade_delete=True) html_content: str = Field(default="", sa_column_kwargs={"server_default": ""}) status: JobStatus = Field( default=JobStatus.AWAITING_START, sa_column_kwargs={"server_default": JobStatus.AWAITING_START.name} ) error: str | None = None ai_edit_instructions: str | None = Field( default=None, description="If the content source is an AI edit, store the instruction here" ) content_source: ContentSource = Field( default=ContentSource.INITIAL_GENERATION, sa_column_kwargs={"server_default": ContentSource.INITIAL_GENERATION.name}, ) class Minute(BaseTableMixin, table=True): __tablename__ = "minute" created_datetime: datetime = Field(sa_column=created_datetime_column(), default=None) updated_datetime: datetime = Field(sa_column=updated_datetime_column(), default=None) transcription_id: UUID = Field(foreign_key="transcription.id", ondelete="CASCADE") transcription: Mapped["Transcription"] = Relationship(back_populates="minutes") template_name: str = Field(default="General") user_template_id: UUID | None = Field( foreign_key="user_template.id", nullable=True, ondelete="SET NULL", default=None ) user_template: "UserTemplate" = Relationship(back_populates="minutes") agenda: str | None = Field(nullable=True, default=None) minute_versions: Mapped[list["MinuteVersion"]] = Relationship( back_populates="minute", cascade_delete=True, sa_relationship_kwargs={"order_by": col(MinuteVersion.created_datetime).desc()}, ) class HallucinationType(StrEnum): FACTUAL_FABRICATION = auto() NONSENSICAL = auto() CONTRADICTION = auto() MISLEADING = auto() OTHER = auto() class Hallucination(BaseTableMixin, table=True): __tablename__ = "hallucination" created_datetime: datetime = Field(sa_column=created_datetime_column(), default=None) updated_datetime: datetime = Field(sa_column=updated_datetime_column(), default=None) minute_version_id: UUID = Field(foreign_key="minute_version.id", ondelete="CASCADE") minute_version: MinuteVersion = Relationship(back_populates="hallucinations") hallucination_type: HallucinationType = Field(description="Type of hallucination", default=HallucinationType.OTHER) hallucination_text: str | None = Field(description="Text of hallucination", default=None) hallucination_reason: str | None = Field(description="Reason for hallucination", default=None) # Main models with table=True for DB tables class User(BaseTableMixin, table=True): __tablename__ = "user" created_datetime: datetime = Field(sa_column=created_datetime_column(), default=None) updated_datetime: datetime = Field(sa_column=updated_datetime_column(), default=None) email: str = Field(index=True) data_retention_days: int | None = Field(default=30) transcriptions: list["Transcription"] = Relationship(back_populates="user") class Recording(BaseTableMixin, table=True): __tablename__ = "recording" created_datetime: datetime = Field(sa_column=created_datetime_column(), default=None) user_id: UUID = Field(foreign_key="user.id", nullable=False) s3_file_key: str transcription_id: UUID | None = Field(default=None, foreign_key="transcription.id", ondelete="SET NULL") transcription: "Transcription" = Relationship(back_populates="recordings") class Chat(BaseTableMixin, table=True): __tablename__ = "chat" created_datetime: datetime = Field(sa_column=created_datetime_column(), default=None) updated_datetime: datetime = Field(sa_column=updated_datetime_column(), default=None) transcription_id: UUID = Field(foreign_key="transcription.id", ondelete="CASCADE") transcription: Mapped["Transcription"] = Relationship(back_populates="chat") user_content: str = Field(default=None) assistant_content: str | None = Field(default=None) status: JobStatus = Field( default=JobStatus.AWAITING_START, sa_column_kwargs={"server_default": JobStatus.AWAITING_START.name} ) error: str | None = Field(default=None) class Transcription(BaseTableMixin, table=True): __tablename__ = "transcription" created_datetime: datetime = Field(sa_column=created_datetime_column(), default=None) updated_datetime: datetime = Field(sa_column=updated_datetime_column(), default=None) title: str | None = Field(default=None) dialogue_entries: list[DialogueEntry] | None = Field(default=None, sa_column=Column(JSONB)) status: JobStatus = Field( default=JobStatus.AWAITING_START, sa_column_kwargs={"server_default": JobStatus.AWAITING_START.name} ) error: str | None = Field(default=None) user: User | None = Relationship(back_populates="transcriptions") user_id: UUID | None = Field(default=None, foreign_key="user.id") minutes: list[Minute] = Relationship( back_populates="transcription", cascade_delete=True, sa_relationship_kwargs={"order_by": col(Minute.created_datetime).desc()}, ) # Kept old minute versions so we can migrate them legacy_minute_versions: list[dict] | None = Field(sa_column=Column(name="minute_versions", type_=JSONB), default=[]) recordings: Mapped[list[Recording]] = Relationship( back_populates="transcription", sa_relationship_kwargs={"order_by": col(Recording.created_datetime).desc()}, ) chat: list[Chat] = Relationship( back_populates="transcription", cascade_delete=True, sa_relationship_kwargs={"order_by": col(Chat.created_datetime).desc()}, ) class TemplateType(StrEnum): DOCUMENT = auto() FORM = auto() class TemplateQuestion(BaseTableMixin, table=True): __tablename__ = "template_question" position: int title: str description: str user_template_id: UUID = Field(foreign_key="user_template.id", ondelete="CASCADE") user_template: "UserTemplate" = Relationship(back_populates="questions") class UserTemplate(BaseTableMixin, table=True): __tablename__ = "user_template" created_datetime: datetime = Field(sa_column=created_datetime_column(), default=None) updated_datetime: datetime = Field(sa_column=updated_datetime_column(), default=None) name: str content: str description: str = "" type: TemplateType = TemplateType.DOCUMENT user_id: UUID | None = Field(default=None, foreign_key="user.id") minutes: list[Minute] = Relationship(back_populates="user_template") questions: list[TemplateQuestion] = Relationship( back_populates="user_template", passive_deletes="all", sa_relationship_kwargs={"order_by": TemplateQuestion.position}, )