diff --git a/python/alembic/richie/versions/2026_06_29-adding_phrase_metadata_tables_dddee09eddcc.py b/python/alembic/richie/versions/2026_06_29-adding_phrase_metadata_tables_dddee09eddcc.py new file mode 100644 index 0000000..d5e6902 --- /dev/null +++ b/python/alembic/richie/versions/2026_06_29-adding_phrase_metadata_tables_dddee09eddcc.py @@ -0,0 +1,206 @@ +"""adding Phrase metadata tables. + +Revision ID: dddee09eddcc +Revises: 96d72c748c24 +Create Date: 2026-06-29 00:49:07.344159 + +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import sqlalchemy as sa +from alembic import op +from sqlalchemy.dialects import postgresql + +from python.orm import RichieBase + +if TYPE_CHECKING: + from collections.abc import Sequence + +# revision identifiers, used by Alembic. +revision: str = "dddee09eddcc" +down_revision: str | None = "96d72c748c24" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + +schema = RichieBase.schema_name + + +def upgrade() -> None: + """Upgrade.""" + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + "candidate_phrases", + sa.Column("book_id", sa.Integer(), nullable=False), + sa.Column("series_id", sa.Integer(), nullable=True), + sa.Column("phrase_text", sa.Text(), nullable=False), + sa.Column("phrase_norm", sa.Text(), nullable=False), + sa.Column("token_count", sa.Integer(), nullable=False), + sa.Column("source_raw_ngram", sa.Boolean(), nullable=False), + sa.Column("source_yake", sa.Boolean(), nullable=False), + sa.Column("source_spacy_ner", sa.Boolean(), nullable=False), + sa.Column("source_spacy_noun_chunk", sa.Boolean(), nullable=False), + sa.Column("source_capitalized", sa.Boolean(), nullable=False), + sa.Column("source_metadata", sa.Boolean(), nullable=False), + sa.Column("spacy_label", sa.String(), nullable=True), + sa.Column("raw_count", sa.Integer(), nullable=False), + sa.Column("chapter_count", sa.Integer(), nullable=False), + sa.Column("yake_score", sa.Float(), nullable=True), + sa.Column("candidate_score", sa.Float(), nullable=False), + sa.Column( + "sample_contexts", + sa.JSON().with_variant(postgresql.JSONB(astext_type=sa.Text()), "postgresql"), + nullable=True, + ), + sa.Column("llm_judged", sa.Boolean(), nullable=False), + sa.Column("llm_keep", sa.Boolean(), nullable=True), + sa.Column("llm_confidence", sa.Float(), nullable=True), + sa.Column("llm_category", sa.String(), nullable=True), + sa.Column("llm_reason", sa.Text(), nullable=True), + sa.Column("id", sa.Integer(), nullable=False), + sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False), + sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False), + sa.ForeignKeyConstraint( + ["book_id"], + [f"{schema}.ebook_source.id"], + name=op.f("fk_candidate_phrases_book_id_ebook_source"), + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_candidate_phrases")), + sa.UniqueConstraint("book_id", "phrase_norm", name="uq_candidate_phrases_book_id_phrase_norm"), + schema=schema, + ) + op.create_index( + "candidate_phrases_book_norm_idx", "candidate_phrases", ["book_id", "phrase_norm"], unique=False, schema=schema + ) + op.create_index( + "candidate_phrases_book_score_idx", + "candidate_phrases", + ["book_id", "candidate_score"], + unique=False, + schema=schema, + ) + op.create_table( + "protected_phrases", + sa.Column("book_id", sa.Integer(), nullable=True), + sa.Column("series_id", sa.Integer(), nullable=True), + sa.Column("phrase_text", sa.Text(), nullable=False), + sa.Column("phrase_norm", sa.Text(), nullable=False), + sa.Column("canonical_id", sa.String(), nullable=False), + sa.Column("phrase_type", sa.String(), nullable=True), + sa.Column("token_count", sa.Integer(), nullable=False), + sa.Column("confidence", sa.Float(), nullable=False), + sa.Column("importance", sa.Float(), nullable=False), + sa.Column("allow_nested", sa.Boolean(), nullable=False), + sa.Column("suppress_children", sa.Boolean(), nullable=False), + sa.Column("source_candidate_id", sa.Integer(), nullable=True), + sa.Column("id", sa.Integer(), nullable=False), + sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False), + sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False), + sa.ForeignKeyConstraint( + ["book_id"], + [f"{schema}.ebook_source.id"], + name=op.f("fk_protected_phrases_book_id_ebook_source"), + ondelete="CASCADE", + ), + sa.ForeignKeyConstraint( + ["source_candidate_id"], + [f"{schema}.candidate_phrases.id"], + name=op.f("fk_protected_phrases_source_candidate_id_candidate_phrases"), + ondelete="SET NULL", + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_protected_phrases")), + sa.UniqueConstraint("book_id", "phrase_norm", name="uq_protected_phrases_book_id_phrase_norm"), + schema=schema, + ) + op.create_index( + "protected_phrases_book_norm_idx", "protected_phrases", ["book_id", "phrase_norm"], unique=False, schema=schema + ) + op.create_index("protected_phrases_norm_idx", "protected_phrases", ["phrase_norm"], unique=False, schema=schema) + op.create_index( + "protected_phrases_series_norm_idx", + "protected_phrases", + ["series_id", "phrase_norm"], + unique=False, + schema=schema, + ) + op.create_table( + "chunk_phrase_mentions", + sa.Column("chunk_id", sa.BigInteger(), nullable=False), + sa.Column("phrase_id", sa.Integer(), nullable=False), + sa.Column("book_id", sa.Integer(), nullable=True), + sa.Column("series_id", sa.Integer(), nullable=True), + sa.Column("start_char", sa.Integer(), nullable=False), + sa.Column("end_char", sa.Integer(), nullable=True), + sa.Column("id", sa.Integer(), nullable=False), + sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False), + sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False), + sa.ForeignKeyConstraint( + ["book_id"], + [f"{schema}.ebook_source.id"], + name=op.f("fk_chunk_phrase_mentions_book_id_ebook_source"), + ondelete="CASCADE", + ), + sa.ForeignKeyConstraint( + ["chunk_id"], + [f"{schema}.ebook_chunk.id"], + name=op.f("fk_chunk_phrase_mentions_chunk_id_ebook_chunk"), + ondelete="CASCADE", + ), + sa.ForeignKeyConstraint( + ["phrase_id"], + [f"{schema}.protected_phrases.id"], + name=op.f("fk_chunk_phrase_mentions_phrase_id_protected_phrases"), + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_chunk_phrase_mentions")), + sa.UniqueConstraint("chunk_id", "phrase_id", "start_char", name="uq_chunk_phrase_mentions_chunk_phrase_start"), + schema=schema, + ) + op.create_index( + "chunk_phrase_mentions_chunk_idx", "chunk_phrase_mentions", ["chunk_id"], unique=False, schema=schema + ) + op.create_index( + "chunk_phrase_mentions_phrase_idx", "chunk_phrase_mentions", ["phrase_id"], unique=False, schema=schema + ) + op.create_table( + "phrase_aliases", + sa.Column("phrase_id", sa.Integer(), nullable=False), + sa.Column("alias_text", sa.Text(), nullable=False), + sa.Column("alias_norm", sa.Text(), nullable=False), + sa.Column("confidence", sa.Float(), nullable=False), + sa.Column("id", sa.Integer(), nullable=False), + sa.Column("created", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False), + sa.Column("updated", sa.DateTime(timezone=True), server_default=sa.text("now()"), nullable=False), + sa.ForeignKeyConstraint( + ["phrase_id"], + [f"{schema}.protected_phrases.id"], + name=op.f("fk_phrase_aliases_phrase_id_protected_phrases"), + ondelete="CASCADE", + ), + sa.PrimaryKeyConstraint("id", name=op.f("pk_phrase_aliases")), + sa.UniqueConstraint("phrase_id", "alias_norm", name="uq_phrase_aliases_phrase_id_alias_norm"), + schema=schema, + ) + op.create_index("phrase_aliases_norm_idx", "phrase_aliases", ["alias_norm"], unique=False, schema=schema) + # ### end Alembic commands ### + + +def downgrade() -> None: + """Downgrade.""" + # ### commands auto generated by Alembic - please adjust! ### + op.drop_index("phrase_aliases_norm_idx", table_name="phrase_aliases", schema=schema) + op.drop_table("phrase_aliases", schema=schema) + op.drop_index("chunk_phrase_mentions_phrase_idx", table_name="chunk_phrase_mentions", schema=schema) + op.drop_index("chunk_phrase_mentions_chunk_idx", table_name="chunk_phrase_mentions", schema=schema) + op.drop_table("chunk_phrase_mentions", schema=schema) + op.drop_index("protected_phrases_series_norm_idx", table_name="protected_phrases", schema=schema) + op.drop_index("protected_phrases_norm_idx", table_name="protected_phrases", schema=schema) + op.drop_index("protected_phrases_book_norm_idx", table_name="protected_phrases", schema=schema) + op.drop_table("protected_phrases", schema=schema) + op.drop_index("candidate_phrases_book_score_idx", table_name="candidate_phrases", schema=schema) + op.drop_index("candidate_phrases_book_norm_idx", table_name="candidate_phrases", schema=schema) + op.drop_table("candidate_phrases", schema=schema) + # ### end Alembic commands ### diff --git a/python/orm/richie/__init__.py b/python/orm/richie/__init__.py index a28ce7a..8549c9b 100644 --- a/python/orm/richie/__init__.py +++ b/python/orm/richie/__init__.py @@ -12,12 +12,16 @@ from python.orm.richie.contact import ( RelationshipType, ) from python.orm.richie.ebook import ( + EbookCandidatePhrase, EbookChapter, EbookChunk, EbookChunkEmbedding1024, EbookChunkEmbedding2560, EbookChunkEmbedding4096, + EbookChunkPhraseMention, EbookEmbeddingModel, + EbookPhraseAlias, + EbookProtectedPhrase, EbookSource, ) @@ -28,12 +32,16 @@ __all__ = [ "Contact", "ContactNeed", "ContactRelationship", + "EbookCandidatePhrase", "EbookChapter", "EbookChunk", "EbookChunkEmbedding1024", "EbookChunkEmbedding2560", "EbookChunkEmbedding4096", + "EbookChunkPhraseMention", "EbookEmbeddingModel", + "EbookPhraseAlias", + "EbookProtectedPhrase", "EbookSource", "Need", "RelationshipType", diff --git a/python/orm/richie/ebook.py b/python/orm/richie/ebook.py index 8e32409..e2da0f3 100644 --- a/python/orm/richie/ebook.py +++ b/python/orm/richie/ebook.py @@ -5,11 +5,23 @@ from __future__ import annotations from datetime import datetime from pgvector.sqlalchemy import Vector -from sqlalchemy import BigInteger, Boolean, DateTime, ForeignKey, Index, String, UniqueConstraint +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.""" @@ -94,7 +106,7 @@ class EbookEmbeddingModel(TableBase): name: Mapped[str] = mapped_column(String, unique=True) dimension: Mapped[int] - is_default: Mapped[bool] = mapped_column(Boolean, default=False) + is_default: Mapped[bool] = mapped_column(default=False) class EbookChunkEmbedding1024(TableBaseBig): @@ -136,3 +148,97 @@ class EbookChunkEmbedding4096(TableBaseBig): 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_spacy_ner: Mapped[bool] = mapped_column(default=False) + source_spacy_noun_chunk: Mapped[bool] = mapped_column(default=False) + source_capitalized: Mapped[bool] = mapped_column(default=False) + source_metadata: Mapped[bool] = mapped_column(default=False) + spacy_label: Mapped[str | None] + 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]