Files
dotfiles/python/ebook_search/protected_phrases/store.py
T
Richie 6bf77299e8
treefmt / nix fmt (pull_request) Failing after 5s
pytest / pytest (pull_request) Successful in 28s
test ebook search / test-ebook-search (pull_request) Failing after 35s
build_systems / build-brain (pull_request) Successful in 49s
build_systems / build-bob (pull_request) Successful in 49s
build_systems / build-rhapsody-in-green (pull_request) Successful in 1m0s
build_systems / build-jeeves (pull_request) Successful in 2m20s
Refactor logging statements to use f-strings for improved readability and consistency across the codebase. This change enhances the clarity of log messages by directly embedding variable values, making it easier to trace and debug application behavior.
2026-07-12 19:34:19 -04:00

685 lines
25 KiB
Python

"""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 (
CorpusPhraseStats,
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.ext.asyncio import AsyncSession
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: AsyncSession, 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 (AsyncSession): 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)
async def load_book_text(session: AsyncSession, book_id: int) -> str:
"""Load a book's indexed chunk text as one string for phrase extraction.
Args:
session (AsyncSession): 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 = await 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()))
async def load_book_chapter_texts(session: AsyncSession, book_id: int) -> list[str]:
"""Reconstruct chapter-like text blocks from indexed chunks for phrase extraction.
Args:
session (AsyncSession): 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 = await 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,
}
async def metadata_for_source_id(session: AsyncSession, source_id: int) -> dict[str, object | None]:
"""Return phrase extraction metadata for one indexed source by id.
Args:
session (AsyncSession): 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 = await 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)
async def count_protected_phrases(session: AsyncSession, book_id: int) -> int:
"""Count stored protected phrases for one book.
Args:
session (AsyncSession): Active database session.
book_id (int): Book whose protected phrases are counted.
Returns:
int: Number of protected phrases stored for the book.
"""
return (
await session.scalars(
select(func.count(EbookProtectedPhrase.id)).where(EbookProtectedPhrase.book_id == book_id)
)
).one()
async def count_unjudged_candidates(session: AsyncSession, book_id: int, config: EbookSearchConfig) -> int:
"""Count storable candidate rows for a book that have not yet been judged.
Args:
session (AsyncSession): 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 (
await 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()
async def corpus_phrase_stats(session: AsyncSession) -> CorpusPhraseStats:
"""Summarize candidate and protected phrase coverage across the whole corpus.
Args:
session (AsyncSession): Active database session.
Returns:
CorpusPhraseStats: Corpus-wide phrase counts and per-book coverage counts.
"""
total_books = (await session.scalars(select(func.count(EbookSource.id)))).one()
candidate_phrases, judged_candidates, books_with_candidates, books_with_unjudged = (
await session.execute(
select(
func.count(EbookCandidatePhrase.id),
func.count(EbookCandidatePhrase.id).filter(EbookCandidatePhrase.llm_judged.is_(True)),
func.count(func.distinct(EbookCandidatePhrase.book_id)),
func.count(func.distinct(EbookCandidatePhrase.book_id)).filter(
EbookCandidatePhrase.llm_judged.is_(False)
),
)
)
).one()
protected_phrases = (await session.scalars(select(func.count(EbookProtectedPhrase.id)))).one()
return CorpusPhraseStats(
total_books=total_books,
books_with_candidates=books_with_candidates,
books_fully_judged=books_with_candidates - books_with_unjudged,
candidate_phrases=candidate_phrases,
judged_candidates=judged_candidates,
unjudged_candidates=candidate_phrases - judged_candidates,
protected_phrases=protected_phrases,
)
async def book_ids_pending_first_judgment(session: AsyncSession) -> list[int]:
"""Return books that have candidate phrases but no judged candidates yet.
Args:
session (AsyncSession): Active database session.
Returns:
list[int]: Book ids with candidates where judging has never run, ordered by id.
"""
judged_books = select(EbookCandidatePhrase.book_id).where(EbookCandidatePhrase.llm_judged.is_(True)).distinct()
return list(
(
await session.scalars(
select(EbookCandidatePhrase.book_id)
.where(EbookCandidatePhrase.book_id.not_in(judged_books))
.distinct()
.order_by(EbookCandidatePhrase.book_id)
)
).all()
)
async def load_candidates_for_judgment(
session: AsyncSession,
book_id: int,
config: EbookSearchConfig,
) -> Sequence[EbookCandidatePhrase]:
"""Load every storable unjudged candidate row for a book.
Rows may have been stored before the current junk filters and score weights existed, so
callers re-check :func:`is_junk_phrase` and rescore before selecting what to judge.
Args:
session (AsyncSession): Active database session.
book_id (int): Book whose candidates are loaded.
config (EbookSearchConfig): Runtime phrase-tuning settings supplying storage thresholds.
Returns:
Sequence[EbookCandidatePhrase]: Storable, unjudged candidate rows ordered by stored score.
"""
query = (
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,
)
)
return (await session.scalars(query)).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_capitalized=row.source_capitalized,
source_metadata=row.source_metadata,
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 candidate_row_values(
book_id: int,
series_id: int | None,
candidate: PhraseCandidate,
*,
judgment: LLMJudgment | None,
) -> dict[str, object]:
"""Build the column values for one candidate phrase upsert.
Args:
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:
dict[str, object]: Column values keyed by column name.
"""
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_capitalized": candidate.source_capitalized,
"source_metadata": candidate.source_metadata,
"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,
)
return values
async def save_candidate_to_db(
session: AsyncSession,
book_id: int,
series_id: int | None,
candidate: PhraseCandidate,
*,
judgment: LLMJudgment | None,
) -> EbookCandidatePhrase:
"""Insert or update one candidate phrase row.
Args:
session (AsyncSession): 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 = candidate_row_values(book_id, series_id, candidate, judgment=judgment)
# 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 (await session.scalars(statement, execution_options={"populate_existing": True})).one()
BULK_CANDIDATE_UPSERT_CHUNK = 1000
async def bulk_upsert_unjudged_candidates(
session: AsyncSession,
book_id: int,
series_id: int | None,
candidates: Sequence[PhraseCandidate],
) -> int:
"""Insert or update many freshly extracted candidate rows in chunked multi-row upserts.
Saving one row per statement costs one database round trip per candidate, which dominated
generation time for full books, so candidates are written ``BULK_CANDIDATE_UPSERT_CHUNK``
rows per statement instead. Existing judgments and sample contexts are never overwritten:
fresh extractions carry no contexts, and ``llm_judged`` plus the ``llm_*`` columns are left
out of the conflict update. Candidates must have unique ``phrase_norm`` values, as produced
by extraction, since one multi-row upsert cannot touch the same row twice.
Args:
session (Session): Active database session.
book_id (int): Book the candidates belong to.
series_id (int | None): Series scope stored on the rows.
candidates (Sequence[PhraseCandidate]): Freshly extracted candidates to persist.
Returns:
int: Number of candidate rows written.
"""
values = [
candidate_row_values(book_id, series_id, candidate, judgment=None)
for candidate in candidates
if not candidate.sample_contexts
]
if len(values) != len(candidates):
msg = "bulk_upsert_unjudged_candidates only accepts freshly extracted candidates without sample contexts"
raise ValueError(msg)
skip_update = {"book_id", "phrase_norm", "llm_judged"}
for chunk_start in range(0, len(values), BULK_CANDIDATE_UPSERT_CHUNK):
chunk = values[chunk_start : chunk_start + BULK_CANDIDATE_UPSERT_CHUNK]
insert_statement = dialect_insert(session, EbookCandidatePhrase).values(chunk)
statement = insert_statement.on_conflict_do_update(
index_elements=["book_id", "phrase_norm"],
set_={column: insert_statement.excluded[column] for column in chunk[0] if column not in skip_update},
)
await session.execute(statement)
return len(values)
def new_candidate_row(book_id: int, series_id: int | None, candidate: PhraseCandidate) -> EbookCandidatePhrase:
"""Build a fresh unjudged candidate row without checking for an existing one.
Unlike :func:`save_candidate_to_db`, this does no lookup, so it is only safe when the caller
guarantees there is no existing row for ``(book_id, candidate.phrase_norm)`` — for example
right after :func:`delete_phrase_data_for_book` has cleared the book.
Args:
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.
Returns:
EbookCandidatePhrase: A new, unattached candidate row.
"""
row = EbookCandidatePhrase(book_id=book_id, phrase_norm=candidate.phrase_norm)
row.llm_judged = False
row.series_id = series_id
row.phrase_text = candidate.phrase_text
row.token_count = candidate.token_count
row.source_raw_ngram = candidate.source_raw_ngram
row.source_yake = candidate.source_yake
row.source_capitalized = candidate.source_capitalized
row.source_metadata = candidate.source_metadata
row.raw_count = candidate.raw_count
row.chapter_count = candidate.chapter_count
row.yake_score = candidate.yake_score
row.candidate_score = candidate.candidate_score
if candidate.sample_contexts:
row.sample_contexts = list(candidate.sample_contexts)
return row
async def upsert_protected_phrase(
session: AsyncSession,
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 (AsyncSession): 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 = (await session.scalars(statement, execution_options={"populate_existing": True})).one()
for alias_text in judgment.aliases:
await upsert_phrase_alias(session, row, alias_text)
return row
async def upsert_phrase_alias(
session: AsyncSession,
phrase: EbookProtectedPhrase,
alias_text: str,
) -> EbookPhraseAlias | None:
"""Insert or update one protected phrase alias.
Args:
session (AsyncSession): 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 (await 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"
async def prune_unstorable_unjudged_candidate_phrases(
session: AsyncSession,
book_id: int,
config: EbookSearchConfig,
) -> int:
"""Delete old unjudged candidate rows that no longer satisfy storage filters.
Args:
session (AsyncSession): 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(
await 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(
f"ebook_candidate_phrase_unstorable_pruned {book_id=} {deleted=} {config.phrase_min_tokens=} "
f"min_uses={minimum_candidate_raw_count(config)}"
)
return deleted
async def delete_phrase_data_for_book(session: AsyncSession, book_id: int) -> PhraseRecalculationResult:
"""Delete all candidate, protected, alias, and mention phrase data for one book.
Args:
session (AsyncSession): 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 = (
await session.scalars(select(EbookProtectedPhrase.id).where(EbookProtectedPhrase.book_id == book_id))
).all()
deleted_aliases = 0
if protected_ids:
deleted_aliases = rowcount(
await session.execute(delete(EbookPhraseAlias).where(EbookPhraseAlias.phrase_id.in_(protected_ids)))
)
deleted_mentions = rowcount(
await session.execute(delete(EbookChunkPhraseMention).where(EbookChunkPhraseMention.book_id == book_id))
)
if protected_ids:
deleted_mentions += rowcount(
await session.execute(
delete(EbookChunkPhraseMention).where(EbookChunkPhraseMention.phrase_id.in_(protected_ids))
)
)
deleted_protected = rowcount(
await session.execute(delete(EbookProtectedPhrase).where(EbookProtectedPhrase.book_id == book_id))
)
deleted_candidates = rowcount(
await session.execute(delete(EbookCandidatePhrase).where(EbookCandidatePhrase.book_id == book_id))
)
await session.flush()
logger.info(
f"ebook_candidate_phrase_data_deleted {book_id=} {deleted_candidates=} {deleted_protected=} {deleted_aliases=} "
f"{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)