Richie
4861f58f27
- 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.
529 lines
19 KiB
Python
529 lines
19 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 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)
|