242 lines
8.9 KiB
Python
242 lines
8.9 KiB
Python
"""EPUB search models."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from datetime import datetime
|
|
|
|
from pgvector.sqlalchemy import Vector
|
|
from sqlalchemy import (
|
|
JSON,
|
|
BigInteger,
|
|
DateTime,
|
|
ForeignKey,
|
|
Index,
|
|
String,
|
|
Text,
|
|
UniqueConstraint,
|
|
)
|
|
from sqlalchemy.dialects.postgresql import JSONB
|
|
from sqlalchemy.orm import Mapped, mapped_column, relationship
|
|
|
|
from python.orm.richie.base import TableBase, TableBaseBig
|
|
|
|
JSON_DOCUMENT = JSON().with_variant(JSONB, "postgresql")
|
|
|
|
|
|
class EbookSource(TableBase):
|
|
"""One indexed EPUB file."""
|
|
|
|
__tablename__ = "ebook_source"
|
|
__table_args__ = (
|
|
UniqueConstraint("file_path"),
|
|
UniqueConstraint("file_sha256"),
|
|
)
|
|
|
|
title: Mapped[str]
|
|
author: Mapped[str | None]
|
|
language: Mapped[str | None]
|
|
publisher: Mapped[str | None]
|
|
identifier: Mapped[str | None]
|
|
file_path: Mapped[str]
|
|
file_sha256: Mapped[str] = mapped_column(String(64))
|
|
file_mtime: Mapped[datetime] = mapped_column(DateTime(timezone=True))
|
|
file_size: Mapped[int] = mapped_column(BigInteger)
|
|
|
|
chapters: Mapped[list[EbookChapter]] = relationship(
|
|
"EbookChapter",
|
|
back_populates="source",
|
|
cascade="all, delete-orphan",
|
|
passive_deletes=True,
|
|
)
|
|
chunks: Mapped[list[EbookChunk]] = relationship(
|
|
"EbookChunk",
|
|
back_populates="source",
|
|
cascade="all, delete-orphan",
|
|
passive_deletes=True,
|
|
)
|
|
|
|
|
|
class EbookChapter(TableBase):
|
|
"""A chapter or spine document inside an EPUB."""
|
|
|
|
__tablename__ = "ebook_chapter"
|
|
__table_args__ = (UniqueConstraint("source_id", "spine_index"),)
|
|
|
|
source_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_source.id", ondelete="CASCADE"))
|
|
spine_index: Mapped[int]
|
|
title: Mapped[str | None]
|
|
href: Mapped[str | None]
|
|
|
|
source: Mapped[EbookSource] = relationship("EbookSource", back_populates="chapters")
|
|
chunks: Mapped[list[EbookChunk]] = relationship(
|
|
"EbookChunk",
|
|
back_populates="chapter",
|
|
cascade="all, delete-orphan",
|
|
passive_deletes=True,
|
|
)
|
|
|
|
|
|
class EbookChunk(TableBaseBig):
|
|
"""A searchable text chunk."""
|
|
|
|
__tablename__ = "ebook_chunk"
|
|
__table_args__ = (
|
|
UniqueConstraint("source_id", "chunk_index", name="uq_ebook_chunk_source_id_chunk_index"),
|
|
UniqueConstraint("source_id", "content_sha256", name="uq_ebook_chunk_source_id_content_sha256"),
|
|
)
|
|
|
|
source_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_source.id", ondelete="CASCADE"))
|
|
chapter_id: Mapped[int | None] = mapped_column(ForeignKey("main.ebook_chapter.id", ondelete="SET NULL"))
|
|
chunk_index: Mapped[int]
|
|
text: Mapped[str]
|
|
token_start: Mapped[int]
|
|
token_count: Mapped[int]
|
|
page_label: Mapped[str | None]
|
|
content_sha256: Mapped[str] = mapped_column(String(64))
|
|
search_text: Mapped[str]
|
|
|
|
source: Mapped[EbookSource] = relationship("EbookSource", back_populates="chunks")
|
|
chapter: Mapped[EbookChapter | None] = relationship("EbookChapter", back_populates="chunks")
|
|
|
|
|
|
class EbookEmbeddingModel(TableBase):
|
|
"""A supported embedding model."""
|
|
|
|
__tablename__ = "ebook_embedding_model"
|
|
|
|
name: Mapped[str] = mapped_column(String, unique=True)
|
|
dimension: Mapped[int]
|
|
is_default: Mapped[bool] = mapped_column(default=False)
|
|
|
|
|
|
class EbookChunkEmbedding1024(TableBaseBig):
|
|
"""1024-dimensional chunk embedding."""
|
|
|
|
__tablename__ = "ebook_chunk_embedding_1024"
|
|
__table_args__ = (
|
|
UniqueConstraint("chunk_id", "model_id"),
|
|
Index(
|
|
"ix_ebook_chunk_embedding_1024_embedding_cosine",
|
|
"embedding",
|
|
postgresql_using="hnsw",
|
|
postgresql_ops={"embedding": "vector_cosine_ops"},
|
|
),
|
|
)
|
|
|
|
chunk_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_chunk.id", ondelete="CASCADE"))
|
|
model_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_embedding_model.id", ondelete="CASCADE"))
|
|
embedding: Mapped[list[float]] = mapped_column(Vector(1024))
|
|
|
|
|
|
class EbookChunkEmbedding2560(TableBaseBig):
|
|
"""2560-dimensional chunk embedding."""
|
|
|
|
__tablename__ = "ebook_chunk_embedding_2560"
|
|
__table_args__ = (UniqueConstraint("chunk_id", "model_id"),)
|
|
|
|
chunk_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_chunk.id", ondelete="CASCADE"))
|
|
model_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_embedding_model.id", ondelete="CASCADE"))
|
|
embedding: Mapped[list[float]] = mapped_column(Vector(2560))
|
|
|
|
|
|
class EbookChunkEmbedding4096(TableBaseBig):
|
|
"""4096-dimensional chunk embedding."""
|
|
|
|
__tablename__ = "ebook_chunk_embedding_4096"
|
|
__table_args__ = (UniqueConstraint("chunk_id", "model_id"),)
|
|
|
|
chunk_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_chunk.id", ondelete="CASCADE"))
|
|
model_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_embedding_model.id", ondelete="CASCADE"))
|
|
embedding: Mapped[list[float]] = mapped_column(Vector(4096))
|
|
|
|
|
|
class EbookCandidatePhrase(TableBase):
|
|
"""A high-recall phrase candidate extracted from one book."""
|
|
|
|
__tablename__ = "candidate_phrases"
|
|
__table_args__ = (
|
|
UniqueConstraint("book_id", "phrase_norm", name="uq_candidate_phrases_book_id_phrase_norm"),
|
|
Index("candidate_phrases_book_score_idx", "book_id", "candidate_score"),
|
|
Index("candidate_phrases_book_norm_idx", "book_id", "phrase_norm"),
|
|
)
|
|
|
|
book_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_source.id", ondelete="CASCADE"))
|
|
series_id: Mapped[int | None]
|
|
phrase_text: Mapped[str] = mapped_column(Text)
|
|
phrase_norm: Mapped[str] = mapped_column(Text)
|
|
token_count: Mapped[int]
|
|
source_raw_ngram: Mapped[bool] = mapped_column(default=False)
|
|
source_yake: Mapped[bool] = mapped_column(default=False)
|
|
source_capitalized: Mapped[bool] = mapped_column(default=False)
|
|
source_metadata: Mapped[bool] = mapped_column(default=False)
|
|
raw_count: Mapped[int] = mapped_column(default=0)
|
|
chapter_count: Mapped[int] = mapped_column(default=0)
|
|
yake_score: Mapped[float | None]
|
|
candidate_score: Mapped[float] = mapped_column(default=0.0)
|
|
sample_contexts: Mapped[list[str] | None] = mapped_column(JSON_DOCUMENT)
|
|
llm_judged: Mapped[bool] = mapped_column(default=False)
|
|
llm_keep: Mapped[bool | None]
|
|
llm_confidence: Mapped[float | None]
|
|
llm_category: Mapped[str | None]
|
|
llm_reason: Mapped[str | None] = mapped_column(Text)
|
|
|
|
|
|
class EbookProtectedPhrase(TableBase):
|
|
"""A phrase accepted by the LLM judge for protected query matching."""
|
|
|
|
__tablename__ = "protected_phrases"
|
|
__table_args__ = (
|
|
UniqueConstraint("book_id", "phrase_norm", name="uq_protected_phrases_book_id_phrase_norm"),
|
|
Index("protected_phrases_norm_idx", "phrase_norm"),
|
|
Index("protected_phrases_book_norm_idx", "book_id", "phrase_norm"),
|
|
Index("protected_phrases_series_norm_idx", "series_id", "phrase_norm"),
|
|
)
|
|
|
|
book_id: Mapped[int | None] = mapped_column(ForeignKey("main.ebook_source.id", ondelete="CASCADE"))
|
|
series_id: Mapped[int | None]
|
|
phrase_text: Mapped[str] = mapped_column(Text)
|
|
phrase_norm: Mapped[str] = mapped_column(Text)
|
|
canonical_id: Mapped[str]
|
|
phrase_type: Mapped[str | None]
|
|
token_count: Mapped[int]
|
|
confidence: Mapped[float]
|
|
importance: Mapped[float] = mapped_column(default=0.5)
|
|
allow_nested: Mapped[bool] = mapped_column(default=False)
|
|
suppress_children: Mapped[bool] = mapped_column(default=True)
|
|
source_candidate_id: Mapped[int | None] = mapped_column(
|
|
ForeignKey("main.candidate_phrases.id", ondelete="SET NULL")
|
|
)
|
|
|
|
|
|
class EbookPhraseAlias(TableBase):
|
|
"""A normalized alias that maps to a protected phrase."""
|
|
|
|
__tablename__ = "phrase_aliases"
|
|
__table_args__ = (
|
|
UniqueConstraint("phrase_id", "alias_norm", name="uq_phrase_aliases_phrase_id_alias_norm"),
|
|
Index("phrase_aliases_norm_idx", "alias_norm"),
|
|
)
|
|
|
|
phrase_id: Mapped[int] = mapped_column(ForeignKey("main.protected_phrases.id", ondelete="CASCADE"))
|
|
alias_text: Mapped[str] = mapped_column(Text)
|
|
alias_norm: Mapped[str] = mapped_column(Text)
|
|
confidence: Mapped[float] = mapped_column(default=1.0)
|
|
|
|
|
|
class EbookChunkPhraseMention(TableBase):
|
|
"""A precomputed occurrence of a protected phrase inside one chunk."""
|
|
|
|
__tablename__ = "chunk_phrase_mentions"
|
|
__table_args__ = (
|
|
UniqueConstraint("chunk_id", "phrase_id", "start_char", name="uq_chunk_phrase_mentions_chunk_phrase_start"),
|
|
Index("chunk_phrase_mentions_phrase_idx", "phrase_id"),
|
|
Index("chunk_phrase_mentions_chunk_idx", "chunk_id"),
|
|
)
|
|
|
|
chunk_id: Mapped[int] = mapped_column(ForeignKey("main.ebook_chunk.id", ondelete="CASCADE"))
|
|
phrase_id: Mapped[int] = mapped_column(ForeignKey("main.protected_phrases.id", ondelete="CASCADE"))
|
|
book_id: Mapped[int | None] = mapped_column(ForeignKey("main.ebook_source.id", ondelete="CASCADE"))
|
|
series_id: Mapped[int | None]
|
|
start_char: Mapped[int]
|
|
end_char: Mapped[int | None]
|