"""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)