Spaces:
Sleeping
Sleeping
File size: 6,253 Bytes
aacd162 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 | from __future__ import annotations
from datetime import datetime
from sqlalchemy import DateTime, Float, ForeignKey, Integer, String, Text, func, JSON
from sqlalchemy.orm import Mapped, mapped_column, relationship
from data.db import Base
class User(Base):
__tablename__ = "users"
id: Mapped[int] = mapped_column(Integer, primary_key=True, index=True)
email: Mapped[str] = mapped_column(String(255), nullable=False, unique=True, index=True)
display_name: Mapped[str | None] = mapped_column(String(255))
avatar_url: Mapped[str | None] = mapped_column(String(1024))
created_at: Mapped[datetime] = mapped_column(
DateTime(timezone=True), server_default=func.now(), nullable=False
)
notebooks: Mapped[list[Notebook]] = relationship(
"Notebook", back_populates="owner", cascade="all, delete-orphan"
)
class Notebook(Base):
__tablename__ = "notebooks"
id: Mapped[int] = mapped_column(Integer, primary_key=True, index=True)
owner_user_id: Mapped[int] = mapped_column(
ForeignKey("users.id", ondelete="CASCADE"), nullable=False, index=True
)
title: Mapped[str] = mapped_column(String(255), nullable=False)
created_at: Mapped[datetime] = mapped_column(
DateTime(timezone=True), server_default=func.now(), nullable=False
)
updated_at: Mapped[datetime] = mapped_column(
DateTime(timezone=True), server_default=func.now(), onupdate=func.now(), nullable=False
)
owner: Mapped[User] = relationship("User", back_populates="notebooks")
sources: Mapped[list[Source]] = relationship(
"Source", back_populates="notebook", cascade="all, delete-orphan"
)
chat_threads: Mapped[list[ChatThread]] = relationship(
"ChatThread", back_populates="notebook", cascade="all, delete-orphan"
)
artifacts: Mapped[list[Artifact]] = relationship(
"Artifact", back_populates="notebook", cascade="all, delete-orphan"
)
class Source(Base):
__tablename__ = "sources"
id: Mapped[int] = mapped_column(Integer, primary_key=True, index=True)
notebook_id: Mapped[int] = mapped_column(
ForeignKey("notebooks.id", ondelete="CASCADE"), nullable=False, index=True
)
type: Mapped[str] = mapped_column(String(50), nullable=False)
title: Mapped[str | None] = mapped_column(String(255))
original_name: Mapped[str | None] = mapped_column(String(1024))
url: Mapped[str | None] = mapped_column(String(2048))
storage_path: Mapped[str | None] = mapped_column(String(1024))
status: Mapped[str] = mapped_column(String(50), nullable=False, default="pending")
ingested_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))
notebook: Mapped[Notebook] = relationship("Notebook", back_populates="sources")
citations: Mapped[list[MessageCitation]] = relationship(
"MessageCitation", back_populates="source", cascade="all, delete-orphan"
)
class ChatThread(Base):
__tablename__ = "chat_threads"
id: Mapped[int] = mapped_column(Integer, primary_key=True, index=True)
notebook_id: Mapped[int] = mapped_column(
ForeignKey("notebooks.id", ondelete="CASCADE"), nullable=False, index=True
)
title: Mapped[str | None] = mapped_column(String(255))
created_at: Mapped[datetime] = mapped_column(
DateTime(timezone=True), server_default=func.now(), nullable=False
)
notebook: Mapped[Notebook] = relationship("Notebook", back_populates="chat_threads")
messages: Mapped[list[Message]] = relationship(
"Message", back_populates="thread", cascade="all, delete-orphan"
)
class Message(Base):
__tablename__ = "messages"
id: Mapped[int] = mapped_column(Integer, primary_key=True, index=True)
thread_id: Mapped[int] = mapped_column(
ForeignKey("chat_threads.id", ondelete="CASCADE"), nullable=False, index=True
)
role: Mapped[str] = mapped_column(String(20), nullable=False)
content: Mapped[str] = mapped_column(Text, nullable=False)
created_at: Mapped[datetime] = mapped_column(
DateTime(timezone=True), server_default=func.now(), nullable=False
)
thread: Mapped[ChatThread] = relationship("ChatThread", back_populates="messages")
citations: Mapped[list[MessageCitation]] = relationship(
"MessageCitation", back_populates="message", cascade="all, delete-orphan"
)
class MessageCitation(Base):
__tablename__ = "message_citations"
id: Mapped[int] = mapped_column(Integer, primary_key=True, index=True)
message_id: Mapped[int] = mapped_column(
ForeignKey("messages.id", ondelete="CASCADE"), nullable=False, index=True
)
source_id: Mapped[int] = mapped_column(
ForeignKey("sources.id", ondelete="CASCADE"), nullable=False, index=True
)
chunk_ref: Mapped[str | None] = mapped_column(String(255))
quote: Mapped[str | None] = mapped_column(Text)
score: Mapped[float | None] = mapped_column(Float)
message: Mapped[Message] = relationship("Message", back_populates="citations")
source: Mapped[Source] = relationship("Source", back_populates="citations")
class Artifact(Base):
"""Generated artifacts: quizzes, podcasts, reports."""
__tablename__ = "artifacts"
id: Mapped[int] = mapped_column(Integer, primary_key=True, index=True)
notebook_id: Mapped[int] = mapped_column(
ForeignKey("notebooks.id", ondelete="CASCADE"), nullable=False, index=True
)
type: Mapped[str] = mapped_column(String(50), nullable=False) # 'quiz', 'podcast', 'report'
title: Mapped[str | None] = mapped_column(String(255))
status: Mapped[str] = mapped_column(String(50), nullable=False, default="pending")
file_path: Mapped[str | None] = mapped_column(String(1024))
artifact_metadata: Mapped[dict | None] = mapped_column("metadata", JSON)
content: Mapped[str | None] = mapped_column(Text)
error_message: Mapped[str | None] = mapped_column(Text)
created_at: Mapped[datetime] = mapped_column(
DateTime(timezone=True), server_default=func.now(), nullable=False
)
generated_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))
notebook: Mapped[Notebook] = relationship("Notebook", back_populates="artifacts") |