93 lines
3.4 KiB
Python
93 lines
3.4 KiB
Python
from datetime import datetime
|
|
from uuid import uuid4
|
|
|
|
from sqlalchemy import (
|
|
CheckConstraint,
|
|
DateTime,
|
|
ForeignKey,
|
|
Index,
|
|
Integer,
|
|
String,
|
|
Text,
|
|
UniqueConstraint,
|
|
)
|
|
from sqlalchemy.orm import Mapped, mapped_column
|
|
|
|
from kbqa.database import Base, utc_now
|
|
|
|
|
|
def new_uuid() -> str:
|
|
return str(uuid4())
|
|
|
|
|
|
class Document(Base):
|
|
__tablename__ = "documents"
|
|
__table_args__ = (
|
|
CheckConstraint(
|
|
"status IN ('pending','indexing','indexed','failed','deleting')",
|
|
name="ck_documents_status",
|
|
),
|
|
Index("ix_documents_status_updated_at", "status", "updated_at"),
|
|
)
|
|
|
|
id: Mapped[str] = mapped_column(String(36), primary_key=True, default=new_uuid)
|
|
filename: Mapped[str] = mapped_column(Text, nullable=False)
|
|
stored_path: Mapped[str] = mapped_column(Text, nullable=False)
|
|
file_type: Mapped[str] = mapped_column(String(16), nullable=False)
|
|
mime_type: Mapped[str] = mapped_column(String(128), nullable=False)
|
|
file_size: Mapped[int] = mapped_column(Integer, nullable=False)
|
|
sha256: Mapped[str] = mapped_column(String(64), nullable=False)
|
|
status: Mapped[str] = mapped_column(String(16), nullable=False, default="pending")
|
|
chunk_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
|
|
error_code: Mapped[str | None] = mapped_column(String(64))
|
|
error_message: Mapped[str | None] = mapped_column(Text)
|
|
created_at: Mapped[datetime] = mapped_column(
|
|
DateTime(timezone=True), nullable=False, default=utc_now
|
|
)
|
|
updated_at: Mapped[datetime] = mapped_column(
|
|
DateTime(timezone=True), nullable=False, default=utc_now, onupdate=utc_now
|
|
)
|
|
|
|
|
|
class IndexJob(Base):
|
|
__tablename__ = "index_jobs"
|
|
__table_args__ = (
|
|
CheckConstraint(
|
|
"status IN ('queued','running','succeeded','failed')",
|
|
name="ck_index_jobs_status",
|
|
),
|
|
)
|
|
|
|
id: Mapped[str] = mapped_column(String(36), primary_key=True, default=new_uuid)
|
|
document_id: Mapped[str] = mapped_column(
|
|
String(36), ForeignKey("documents.id", ondelete="CASCADE"), nullable=False, unique=True
|
|
)
|
|
status: Mapped[str] = mapped_column(String(16), nullable=False, default="queued")
|
|
attempt_count: Mapped[int] = mapped_column(Integer, nullable=False, default=0)
|
|
last_error: Mapped[str | None] = mapped_column(Text)
|
|
created_at: Mapped[datetime] = mapped_column(
|
|
DateTime(timezone=True), nullable=False, default=utc_now
|
|
)
|
|
started_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))
|
|
finished_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))
|
|
|
|
|
|
class Chunk(Base):
|
|
__tablename__ = "chunks"
|
|
__table_args__ = (
|
|
UniqueConstraint("document_id", "order_index", name="uq_chunks_document_order"),
|
|
Index("ix_chunks_document_id", "document_id"),
|
|
)
|
|
|
|
id: Mapped[str] = mapped_column(String(36), primary_key=True)
|
|
document_id: Mapped[str] = mapped_column(
|
|
String(36), ForeignKey("documents.id", ondelete="CASCADE"), nullable=False
|
|
)
|
|
order_index: Mapped[int] = mapped_column(Integer, nullable=False)
|
|
content: Mapped[str] = mapped_column(Text, nullable=False)
|
|
page_number: Mapped[int | None] = mapped_column(Integer)
|
|
char_count: Mapped[int] = mapped_column(Integer, nullable=False)
|
|
created_at: Mapped[datetime] = mapped_column(
|
|
DateTime(timezone=True), nullable=False, default=utc_now
|
|
)
|