Add models and database persistence for protected phrase extraction
- Introduced dataclasses for phrase candidates, judgments, and matches in `models.py`. - Implemented database operations for candidate and protected phrases in `store.py`, including loading, saving, and deleting phrases. - Enhanced text normalization functions in `text_normalization.py` with detailed docstrings. - Refactored search functionality to utilize new models and methods for detecting protected phrases.
This commit is contained in:
@@ -0,0 +1,528 @@
|
||||
"""Database persistence for candidate and protected phrase rows."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
import re
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from sqlalchemy import delete, func, or_, select
|
||||
from sqlalchemy.dialects.postgresql import insert as pg_insert
|
||||
from sqlalchemy.dialects.sqlite import insert as sqlite_insert
|
||||
|
||||
from python.ebook_search.protected_phrases.extraction import minimum_candidate_raw_count
|
||||
from python.ebook_search.protected_phrases.models import PhraseCandidate, PhraseRecalculationResult
|
||||
from python.ebook_search.protected_phrases.text_normalization import normalize_text
|
||||
from python.orm.richie import (
|
||||
EbookCandidatePhrase,
|
||||
EbookChunk,
|
||||
EbookChunkPhraseMention,
|
||||
EbookPhraseAlias,
|
||||
EbookProtectedPhrase,
|
||||
EbookSource,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Sequence
|
||||
|
||||
from sqlalchemy.dialects.postgresql.dml import Insert as PostgresInsert
|
||||
from sqlalchemy.dialects.sqlite.dml import Insert as SqliteInsert
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from python.ebook_search.config import EbookSearchConfig
|
||||
from python.ebook_search.protected_phrases.models import LLMJudgment
|
||||
from python.orm.richie.base import TableBase
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def dialect_insert(session: Session, table: type[TableBase]) -> PostgresInsert | SqliteInsert:
|
||||
"""Return a dialect-specific INSERT construct that supports ``ON CONFLICT DO UPDATE``.
|
||||
|
||||
Production runs on PostgreSQL while tests run on SQLite; both support upserts with
|
||||
compatible SQLAlchemy constructs, so the correct one is chosen from the bound dialect.
|
||||
|
||||
Args:
|
||||
session (Session): Active database session whose bind selects the dialect.
|
||||
table (type[TableBase]): Mapped table to insert into.
|
||||
|
||||
Returns:
|
||||
PostgresInsert | SqliteInsert: A dialect insert exposing ``on_conflict_do_update``.
|
||||
"""
|
||||
if session.get_bind().dialect.name == "sqlite":
|
||||
return sqlite_insert(table)
|
||||
return pg_insert(table)
|
||||
|
||||
|
||||
|
||||
|
||||
def load_book_text(session: Session, book_id: int) -> str:
|
||||
"""Load a book's indexed chunk text as one string for phrase extraction.
|
||||
|
||||
Args:
|
||||
session (Session): Active database session.
|
||||
book_id (int): Book whose chunk text is loaded.
|
||||
|
||||
Returns:
|
||||
str: The book's chunk text joined into a single string.
|
||||
"""
|
||||
texts = session.scalars(
|
||||
select(EbookChunk.text).where(EbookChunk.source_id == book_id).order_by(EbookChunk.chunk_index)
|
||||
)
|
||||
return "\n\n".join(stripped for text in texts if (stripped := text.strip()))
|
||||
|
||||
|
||||
def load_book_chapter_texts(session: Session, book_id: int) -> list[str]:
|
||||
"""Reconstruct chapter-like text blocks from indexed chunks for phrase extraction.
|
||||
|
||||
Args:
|
||||
session (Session): Active database session.
|
||||
book_id (int): Book whose chunks are grouped into chapters.
|
||||
|
||||
Returns:
|
||||
list[str]: Non-empty chapter-like text blocks in chunk order.
|
||||
"""
|
||||
rows = session.execute(
|
||||
select(EbookChunk.chapter_id, EbookChunk.text)
|
||||
.where(EbookChunk.source_id == book_id)
|
||||
.order_by(EbookChunk.chunk_index)
|
||||
)
|
||||
chapters: list[str] = []
|
||||
current_chapter_id: int | None = None
|
||||
current_parts: list[str] = []
|
||||
have_current = False
|
||||
|
||||
for chapter_id, text in rows:
|
||||
if have_current and chapter_id != current_chapter_id:
|
||||
chapter_text = "\n\n".join(current_parts).strip()
|
||||
if chapter_text:
|
||||
chapters.append(chapter_text)
|
||||
current_parts = []
|
||||
current_chapter_id = chapter_id
|
||||
current_parts.append(str(text))
|
||||
have_current = True
|
||||
|
||||
if current_parts:
|
||||
chapter_text = "\n\n".join(current_parts).strip()
|
||||
if chapter_text:
|
||||
chapters.append(chapter_text)
|
||||
return chapters
|
||||
|
||||
|
||||
def metadata_for_source(source: EbookSource) -> dict[str, object | None]:
|
||||
"""Return phrase extraction metadata for one indexed source.
|
||||
|
||||
Args:
|
||||
source (EbookSource): Indexed source to read metadata from.
|
||||
|
||||
Returns:
|
||||
dict[str, object | None]: Title, author, language, publisher, and identifier values.
|
||||
"""
|
||||
return {
|
||||
"title": source.title,
|
||||
"author": source.author,
|
||||
"language": source.language,
|
||||
"publisher": source.publisher,
|
||||
"identifier": source.identifier,
|
||||
}
|
||||
|
||||
|
||||
def metadata_for_source_id(session: Session, source_id: int) -> dict[str, object | None]:
|
||||
"""Return phrase extraction metadata for one indexed source by id.
|
||||
|
||||
Args:
|
||||
session (Session): Active database session.
|
||||
source_id (int): Id of the indexed source to read metadata from.
|
||||
|
||||
Returns:
|
||||
dict[str, object | None]: Title, author, language, publisher, and identifier values.
|
||||
|
||||
Raises:
|
||||
ValueError: If no source exists with the given id.
|
||||
"""
|
||||
source = session.get(EbookSource, source_id)
|
||||
if source is None:
|
||||
msg = f"No indexed source with id {source_id}"
|
||||
raise ValueError(msg)
|
||||
return metadata_for_source(source)
|
||||
|
||||
|
||||
def count_protected_phrases(session: Session, book_id: int) -> int:
|
||||
"""Count stored protected phrases for one book.
|
||||
|
||||
Args:
|
||||
session (Session): Active database session.
|
||||
book_id (int): Book whose protected phrases are counted.
|
||||
|
||||
Returns:
|
||||
int: Number of protected phrases stored for the book.
|
||||
"""
|
||||
return session.scalars(
|
||||
select(func.count(EbookProtectedPhrase.id)).where(EbookProtectedPhrase.book_id == book_id)
|
||||
).one()
|
||||
|
||||
|
||||
def count_unjudged_candidates(session: Session, book_id: int, config: EbookSearchConfig) -> int:
|
||||
"""Count storable candidate rows for a book that have not yet been judged.
|
||||
|
||||
Args:
|
||||
session (Session): Active database session.
|
||||
book_id (int): Book whose unjudged candidates are counted.
|
||||
config (EbookSearchConfig): Runtime phrase-tuning settings supplying storage thresholds.
|
||||
|
||||
Returns:
|
||||
int: Number of storable, unjudged candidate rows for the book.
|
||||
"""
|
||||
return session.scalars(
|
||||
select(func.count(EbookCandidatePhrase.id)).where(
|
||||
EbookCandidatePhrase.book_id == book_id,
|
||||
EbookCandidatePhrase.llm_judged.is_(False),
|
||||
EbookCandidatePhrase.token_count >= config.phrase_min_tokens,
|
||||
EbookCandidatePhrase.raw_count >= minimum_candidate_raw_count(config),
|
||||
)
|
||||
).one()
|
||||
|
||||
|
||||
def load_candidates_for_judgment(
|
||||
session: Session,
|
||||
book_id: int,
|
||||
judgment_limit: int,
|
||||
config: EbookSearchConfig,
|
||||
) -> Sequence[EbookCandidatePhrase]:
|
||||
"""Load the top unjudged candidate rows for a book.
|
||||
|
||||
Common-word phrases are filtered before storage (``filter_storable_candidates``), so
|
||||
no re-check is needed here.
|
||||
|
||||
Args:
|
||||
session (Session): Active database session.
|
||||
book_id (int): Book whose candidates are loaded.
|
||||
judgment_limit (int): Maximum number of candidate rows to return.
|
||||
config (EbookSearchConfig): Runtime phrase-tuning settings supplying storage thresholds.
|
||||
|
||||
Returns:
|
||||
Sequence[EbookCandidatePhrase]: Top storable, unjudged candidate rows ordered by score.
|
||||
"""
|
||||
return session.scalars(
|
||||
select(EbookCandidatePhrase)
|
||||
.where(
|
||||
EbookCandidatePhrase.book_id == book_id,
|
||||
EbookCandidatePhrase.llm_judged.is_(False),
|
||||
EbookCandidatePhrase.token_count >= config.phrase_min_tokens,
|
||||
EbookCandidatePhrase.raw_count >= minimum_candidate_raw_count(config),
|
||||
)
|
||||
.order_by(
|
||||
EbookCandidatePhrase.candidate_score.desc(),
|
||||
EbookCandidatePhrase.raw_count.desc(),
|
||||
EbookCandidatePhrase.id,
|
||||
)
|
||||
.limit(judgment_limit)
|
||||
).all()
|
||||
|
||||
|
||||
def phrase_candidate_from_row(row: EbookCandidatePhrase) -> PhraseCandidate:
|
||||
"""Recreate an in-memory candidate from a persisted candidate row.
|
||||
|
||||
Args:
|
||||
row (EbookCandidatePhrase): Stored candidate row to convert.
|
||||
|
||||
Returns:
|
||||
PhraseCandidate: An in-memory candidate mirroring the row's fields.
|
||||
"""
|
||||
return PhraseCandidate(
|
||||
phrase_text=row.phrase_text,
|
||||
phrase_norm=row.phrase_norm,
|
||||
token_count=row.token_count,
|
||||
source_raw_ngram=row.source_raw_ngram,
|
||||
source_yake=row.source_yake,
|
||||
source_spacy_ner=row.source_spacy_ner,
|
||||
source_spacy_noun_chunk=row.source_spacy_noun_chunk,
|
||||
source_capitalized=row.source_capitalized,
|
||||
source_metadata=row.source_metadata,
|
||||
spacy_label=row.spacy_label,
|
||||
raw_count=row.raw_count,
|
||||
chapter_count=row.chapter_count,
|
||||
yake_score=row.yake_score,
|
||||
candidate_score=row.candidate_score,
|
||||
sample_contexts=row.sample_contexts or [],
|
||||
)
|
||||
|
||||
|
||||
def save_candidate_to_db(
|
||||
session: Session,
|
||||
book_id: int,
|
||||
series_id: int | None,
|
||||
candidate: PhraseCandidate,
|
||||
*,
|
||||
judgment: LLMJudgment | None,
|
||||
) -> EbookCandidatePhrase:
|
||||
"""Insert or update one candidate phrase row.
|
||||
|
||||
Args:
|
||||
session (Session): Active database session.
|
||||
book_id (int): Book the candidate belongs to.
|
||||
series_id (int | None): Series scope stored on the row.
|
||||
candidate (PhraseCandidate): Candidate whose fields are written to the row.
|
||||
judgment (LLMJudgment | None): Judgment to record, or ``None`` to leave the row unjudged.
|
||||
|
||||
Returns:
|
||||
EbookCandidatePhrase: The inserted or updated candidate row.
|
||||
"""
|
||||
values: dict[str, object] = {
|
||||
"book_id": book_id,
|
||||
"phrase_norm": candidate.phrase_norm,
|
||||
"series_id": series_id,
|
||||
"phrase_text": candidate.phrase_text,
|
||||
"token_count": candidate.token_count,
|
||||
"source_raw_ngram": candidate.source_raw_ngram,
|
||||
"source_yake": candidate.source_yake,
|
||||
"source_spacy_ner": candidate.source_spacy_ner,
|
||||
"source_spacy_noun_chunk": candidate.source_spacy_noun_chunk,
|
||||
"source_capitalized": candidate.source_capitalized,
|
||||
"source_metadata": candidate.source_metadata,
|
||||
"spacy_label": candidate.spacy_label,
|
||||
"raw_count": candidate.raw_count,
|
||||
"chapter_count": candidate.chapter_count,
|
||||
"yake_score": candidate.yake_score,
|
||||
"candidate_score": candidate.candidate_score,
|
||||
"llm_judged": judgment is not None,
|
||||
}
|
||||
if candidate.sample_contexts:
|
||||
values["sample_contexts"] = list(candidate.sample_contexts)
|
||||
if judgment is not None:
|
||||
values.update(
|
||||
llm_keep=judgment.keep,
|
||||
llm_confidence=judgment.confidence,
|
||||
llm_category=judgment.category,
|
||||
llm_reason=judgment.reason,
|
||||
)
|
||||
|
||||
# Preserve an existing judgment when this call is only refreshing candidate fields.
|
||||
skip_update = {"book_id", "phrase_norm"}
|
||||
if judgment is None:
|
||||
skip_update.add("llm_judged")
|
||||
insert_statement = dialect_insert(session, EbookCandidatePhrase).values(**values)
|
||||
statement = insert_statement.on_conflict_do_update(
|
||||
index_elements=["book_id", "phrase_norm"],
|
||||
set_={column: insert_statement.excluded[column] for column in values if column not in skip_update},
|
||||
).returning(EbookCandidatePhrase)
|
||||
return session.scalars(statement, execution_options={"populate_existing": True}).one()
|
||||
|
||||
|
||||
def upsert_protected_phrase(
|
||||
session: Session,
|
||||
book_id: int,
|
||||
series_id: int | None,
|
||||
candidate: PhraseCandidate,
|
||||
judgment: LLMJudgment,
|
||||
source_candidate: EbookCandidatePhrase,
|
||||
) -> EbookProtectedPhrase:
|
||||
"""Insert or update one accepted protected phrase and its aliases.
|
||||
|
||||
Args:
|
||||
session (Session): Active database session.
|
||||
book_id (int): Book the protected phrase belongs to.
|
||||
series_id (int | None): Series scope stored on the phrase.
|
||||
candidate (PhraseCandidate): Candidate the phrase was promoted from.
|
||||
judgment (LLMJudgment): Accepted judgment supplying canonical text, category, and aliases.
|
||||
source_candidate (EbookCandidatePhrase): Candidate row the phrase was promoted from.
|
||||
|
||||
Returns:
|
||||
EbookProtectedPhrase: The inserted or updated protected phrase row.
|
||||
|
||||
Raises:
|
||||
ValueError: If the chosen phrase text normalizes to empty.
|
||||
"""
|
||||
phrase_text = judgment.canonical or candidate.phrase_text
|
||||
phrase_norm = normalize_text(phrase_text)
|
||||
if not phrase_norm:
|
||||
msg = f"Protected phrase normalized to empty text: {phrase_text!r}"
|
||||
raise ValueError(msg)
|
||||
|
||||
values = {
|
||||
"book_id": book_id,
|
||||
"phrase_norm": phrase_norm,
|
||||
"series_id": series_id,
|
||||
"phrase_text": phrase_text,
|
||||
"canonical_id": make_canonical_id(judgment, phrase_norm),
|
||||
"phrase_type": judgment.category,
|
||||
"token_count": len(phrase_norm.split()),
|
||||
"confidence": judgment.confidence,
|
||||
"importance": judgment.importance,
|
||||
"allow_nested": judgment.allow_nested,
|
||||
"suppress_children": judgment.suppress_children,
|
||||
"source_candidate_id": source_candidate.id,
|
||||
}
|
||||
insert_statement = dialect_insert(session, EbookProtectedPhrase).values(**values)
|
||||
statement = insert_statement.on_conflict_do_update(
|
||||
index_elements=["book_id", "phrase_norm"],
|
||||
set_={
|
||||
column: insert_statement.excluded[column] for column in values if column not in {"book_id", "phrase_norm"}
|
||||
},
|
||||
).returning(EbookProtectedPhrase)
|
||||
row = session.scalars(statement, execution_options={"populate_existing": True}).one()
|
||||
|
||||
for alias_text in judgment.aliases:
|
||||
upsert_phrase_alias(session, row, alias_text)
|
||||
return row
|
||||
|
||||
|
||||
def upsert_phrase_alias(session: Session, phrase: EbookProtectedPhrase, alias_text: str) -> EbookPhraseAlias | None:
|
||||
"""Insert or update one protected phrase alias.
|
||||
|
||||
Args:
|
||||
session (Session): Active database session.
|
||||
phrase (EbookProtectedPhrase): Protected phrase the alias points to.
|
||||
alias_text (str): Alias surface form to store.
|
||||
|
||||
Returns:
|
||||
EbookPhraseAlias | None: The alias row, or ``None`` when the alias is empty or equals the phrase.
|
||||
"""
|
||||
alias_norm = normalize_text(alias_text)
|
||||
if not alias_norm or alias_norm == phrase.phrase_norm:
|
||||
return None
|
||||
|
||||
insert_statement = dialect_insert(session, EbookPhraseAlias).values(
|
||||
phrase_id=phrase.id,
|
||||
alias_norm=alias_norm,
|
||||
alias_text=alias_text,
|
||||
confidence=1.0,
|
||||
)
|
||||
statement = insert_statement.on_conflict_do_update(
|
||||
index_elements=["phrase_id", "alias_norm"],
|
||||
set_={
|
||||
"alias_text": insert_statement.excluded.alias_text,
|
||||
"confidence": insert_statement.excluded.confidence,
|
||||
},
|
||||
).returning(EbookPhraseAlias)
|
||||
return session.scalars(statement, execution_options={"populate_existing": True}).one()
|
||||
|
||||
|
||||
def make_canonical_id(judgment: LLMJudgment, phrase_norm: str) -> str:
|
||||
"""Create a deterministic canonical id from a judgment category and phrase.
|
||||
|
||||
Args:
|
||||
judgment (LLMJudgment): Judgment supplying the phrase category.
|
||||
phrase_norm (str): Normalized phrase text to slugify.
|
||||
|
||||
Returns:
|
||||
str: A ``category:slug`` canonical identifier.
|
||||
"""
|
||||
category = slugify_identifier(judgment.category or "phrase")
|
||||
phrase_slug = slugify_identifier(phrase_norm)
|
||||
return f"{category}:{phrase_slug}"
|
||||
|
||||
|
||||
def slugify_identifier(value: str) -> str:
|
||||
"""Normalize text for use inside a canonical id.
|
||||
|
||||
Args:
|
||||
value (str): Text to slugify.
|
||||
|
||||
Returns:
|
||||
str: A lowercase underscore slug, or ``"unknown"`` when empty.
|
||||
"""
|
||||
slug = re.sub(r"[^a-z0-9]+", "_", normalize_text(value).replace("'", ""))
|
||||
return slug.strip("_") or "unknown"
|
||||
|
||||
|
||||
def prune_unstorable_unjudged_candidate_phrases(
|
||||
session: Session,
|
||||
book_id: int,
|
||||
config: EbookSearchConfig,
|
||||
) -> int:
|
||||
"""Delete old unjudged candidate rows that no longer satisfy storage filters.
|
||||
|
||||
Args:
|
||||
session (Session): Active database session.
|
||||
book_id (int): Book whose stale candidates are pruned.
|
||||
config (EbookSearchConfig): Runtime phrase-tuning settings supplying storage thresholds.
|
||||
|
||||
Returns:
|
||||
int: Number of candidate rows deleted.
|
||||
"""
|
||||
deleted = rowcount(
|
||||
session.execute(
|
||||
delete(EbookCandidatePhrase).where(
|
||||
EbookCandidatePhrase.book_id == book_id,
|
||||
EbookCandidatePhrase.llm_judged.is_(False),
|
||||
or_(
|
||||
EbookCandidatePhrase.token_count < config.phrase_min_tokens,
|
||||
EbookCandidatePhrase.raw_count < minimum_candidate_raw_count(config),
|
||||
),
|
||||
)
|
||||
)
|
||||
)
|
||||
if deleted:
|
||||
logger.info(
|
||||
"ebook_candidate_phrase_unstorable_pruned book_id=%s deleted=%s min_tokens=%s min_uses=%s",
|
||||
book_id,
|
||||
deleted,
|
||||
config.phrase_min_tokens,
|
||||
minimum_candidate_raw_count(config),
|
||||
)
|
||||
return deleted
|
||||
|
||||
|
||||
def delete_phrase_data_for_book(session: Session, book_id: int) -> PhraseRecalculationResult:
|
||||
"""Delete all candidate, protected, alias, and mention phrase data for one book.
|
||||
|
||||
Args:
|
||||
session (Session): Active database session.
|
||||
book_id (int): Book whose phrase data is deleted.
|
||||
|
||||
Returns:
|
||||
PhraseRecalculationResult: Deleted-row counts with ``candidate_phrases`` set to 0.
|
||||
"""
|
||||
protected_ids = session.scalars(
|
||||
select(EbookProtectedPhrase.id).where(EbookProtectedPhrase.book_id == book_id)
|
||||
).all()
|
||||
deleted_aliases = 0
|
||||
if protected_ids:
|
||||
deleted_aliases = rowcount(
|
||||
session.execute(delete(EbookPhraseAlias).where(EbookPhraseAlias.phrase_id.in_(protected_ids)))
|
||||
)
|
||||
|
||||
deleted_mentions = rowcount(
|
||||
session.execute(delete(EbookChunkPhraseMention).where(EbookChunkPhraseMention.book_id == book_id))
|
||||
)
|
||||
if protected_ids:
|
||||
deleted_mentions += rowcount(
|
||||
session.execute(delete(EbookChunkPhraseMention).where(EbookChunkPhraseMention.phrase_id.in_(protected_ids)))
|
||||
)
|
||||
|
||||
deleted_protected = rowcount(
|
||||
session.execute(delete(EbookProtectedPhrase).where(EbookProtectedPhrase.book_id == book_id))
|
||||
)
|
||||
deleted_candidates = rowcount(
|
||||
session.execute(delete(EbookCandidatePhrase).where(EbookCandidatePhrase.book_id == book_id))
|
||||
)
|
||||
session.flush()
|
||||
logger.info(
|
||||
"ebook_candidate_phrase_data_deleted book_id=%s candidates=%s protected=%s aliases=%s mentions=%s",
|
||||
book_id,
|
||||
deleted_candidates,
|
||||
deleted_protected,
|
||||
deleted_aliases,
|
||||
deleted_mentions,
|
||||
)
|
||||
return PhraseRecalculationResult(
|
||||
book_id=book_id,
|
||||
deleted_candidates=deleted_candidates,
|
||||
deleted_protected_phrases=deleted_protected,
|
||||
deleted_aliases=deleted_aliases,
|
||||
deleted_mentions=deleted_mentions,
|
||||
candidate_phrases=0,
|
||||
)
|
||||
|
||||
def rowcount(result: object) -> int:
|
||||
"""Return a safe integer rowcount from a SQLAlchemy execution result.
|
||||
|
||||
Args:
|
||||
result (object): SQLAlchemy execution result that may expose ``rowcount``.
|
||||
|
||||
Returns:
|
||||
int: The result's rowcount, or 0 when it is missing or negative.
|
||||
"""
|
||||
count = getattr(result, "rowcount", 0)
|
||||
return int(count if count is not None and count >= 0 else 0)
|
||||
Reference in New Issue
Block a user