From 4861f58f274c23bf55aa028a28c682eed0f0aa83 Mon Sep 17 00:00:00 2001 From: Richie Cahill Date: Sat, 4 Jul 2026 23:37:27 -0400 Subject: [PATCH] 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. --- python/ebook_search/api/routes/admin.py | 37 +- python/ebook_search/api/routes/page.py | 2 +- python/ebook_search/ingest.py | 2 +- .../protected_phrases/extraction.py | 731 +++++++ .../protected_phrases/generate_ngrams.py | 274 +++ .../protected_phrases/judge_ngrams.py | 481 ++++ python/ebook_search/protected_phrases/lib.py | 1933 ----------------- .../protected_phrases/matching.py | 487 +++++ .../ebook_search/protected_phrases/models.py | 262 +++ .../ebook_search/protected_phrases/store.py | 528 +++++ .../protected_phrases/text_normalization.py | 40 +- python/ebook_search/search.py | 4 +- 12 files changed, 2803 insertions(+), 1978 deletions(-) create mode 100644 python/ebook_search/protected_phrases/extraction.py create mode 100644 python/ebook_search/protected_phrases/generate_ngrams.py create mode 100644 python/ebook_search/protected_phrases/judge_ngrams.py delete mode 100644 python/ebook_search/protected_phrases/lib.py create mode 100644 python/ebook_search/protected_phrases/matching.py create mode 100644 python/ebook_search/protected_phrases/models.py create mode 100644 python/ebook_search/protected_phrases/store.py diff --git a/python/ebook_search/api/routes/admin.py b/python/ebook_search/api/routes/admin.py index b43e6fe..737ec74 100644 --- a/python/ebook_search/api/routes/admin.py +++ b/python/ebook_search/api/routes/admin.py @@ -14,11 +14,8 @@ from python.ebook_search.api.dependencies import ( from python.ebook_search.api.web import templates from python.ebook_search.embeddings import embed_missing_chunks, embedding_model_stats from python.ebook_search.ingest import ingest_configured_paths -from python.ebook_search.protected_phrases.lib import ( - build_missing_protected_phrases, - generate_candidate_phrases_for_books, - judge_candidate_phrases_for_books, -) +from python.ebook_search.protected_phrases.generate_ngrams import generate_candidate_phrases_for_books +from python.ebook_search.protected_phrases.judge_ngrams import judge_candidate_phrases_for_books from python.fastapi_tools import DbSession # noqa: TC001 FastAPI resolves this annotated dependency at runtime logger = logging.getLogger(__name__) @@ -50,36 +47,6 @@ def scan_library(request: Request, config: AppConfig, session: DbSession) -> HTM return templates.TemplateResponse(request, "partials/admin_status.html", {"message": f"Indexed {count} EPUBs"}) -@router.post("/build-phrases", response_class=HTMLResponse) -def build_phrases(request: Request, config: AppConfig, session: DbSession) -> HTMLResponse: - """Build protected phrases for indexed books that are missing them.""" - try: - result = build_missing_protected_phrases(session, config) - session.commit() - except Exception as error: - session.rollback() - logger.exception("ebook_admin_build_phrases_failed") - return templates.TemplateResponse(request, "partials/error.html", {"message": str(error)}, status_code=500) - - logger.info( - "ebook_admin_build_phrases_complete books_seen=%s books_built=%s protected=%s mentions=%s", - result.books_seen, - result.books_built, - result.protected_phrases, - result.phrase_mentions, - ) - return templates.TemplateResponse( - request, - "partials/admin_status.html", - { - "message": ( - f"Built phrases for {result.books_built} of {result.books_seen} books; " - f"{result.protected_phrases} protected phrases, {result.phrase_mentions} mentions" - ) - }, - ) - - @router.post("/generate-ngrams", response_class=HTMLResponse) def generate_ngrams(request: Request, config: AppConfig, session: DbSession) -> HTMLResponse: """Generate candidate n-grams for indexed books without LLM judging.""" diff --git a/python/ebook_search/api/routes/page.py b/python/ebook_search/api/routes/page.py index 0ba67af..4608f39 100644 --- a/python/ebook_search/api/routes/page.py +++ b/python/ebook_search/api/routes/page.py @@ -12,7 +12,7 @@ from python.ebook_search.api.dependencies import ( AppConfig, # noqa: TC001 FastAPI resolves this annotated dependency at runtime ) from python.ebook_search.api.web import templates -from python.ebook_search.protected_phrases.lib import recalculate_candidate_phrases_for_book +from python.ebook_search.protected_phrases.generate_ngrams import recalculate_candidate_phrases_for_book from python.fastapi_tools import DbSession # noqa: TC001 FastAPI resolves this annotated dependency at runtime from python.orm.richie import EbookCandidatePhrase, EbookProtectedPhrase, EbookSource diff --git a/python/ebook_search/ingest.py b/python/ebook_search/ingest.py index 4111f69..b5f9254 100644 --- a/python/ebook_search/ingest.py +++ b/python/ebook_search/ingest.py @@ -13,7 +13,7 @@ import tiktoken from sqlalchemy import or_, select from python.ebook_search.epub_parse import parse_epub -from python.ebook_search.protected_phrases.lib import index_chunk_phrase_mentions_for_book +from python.ebook_search.protected_phrases.matching import index_chunk_phrase_mentions_for_book from python.orm.richie import EbookChapter, EbookChunk, EbookSource logger = logging.getLogger(__name__) diff --git a/python/ebook_search/protected_phrases/extraction.py b/python/ebook_search/protected_phrases/extraction.py new file mode 100644 index 0000000..93c0668 --- /dev/null +++ b/python/ebook_search/protected_phrases/extraction.py @@ -0,0 +1,731 @@ +"""Candidate phrase extraction and scoring for protected phrases.""" + +from __future__ import annotations + +import logging +import re +from collections import defaultdict +from time import perf_counter +from typing import TYPE_CHECKING, Protocol + +from yake import KeywordExtractor + +from python.ebook_search.protected_phrases.config import ( + get_bad_ends, + get_bad_starts, + get_ignored_phrases, + get_most_common_words, +) +from python.ebook_search.protected_phrases.models import PhraseCandidate +from python.ebook_search.protected_phrases.text_normalization import tokenize, tokenize_with_offsets + +if TYPE_CHECKING: + from collections.abc import Iterable, Mapping, Sequence + + from python.ebook_search.config import EbookSearchConfig + +logger = logging.getLogger(__name__) + + +BAD_START_SCORE_PENALTY = 3.0 +BAD_END_SCORE_PENALTY = 3.0 +SOURCE_FIELDS = ( + "source_raw_ngram", + "source_yake", + "source_spacy_ner", + "source_spacy_noun_chunk", + "source_capitalized", + "source_metadata", +) +CAPITALIZED_PHRASE_RE = re.compile(r"\b(?:[A-Z][a-zA-Z']+)(?:\s+(?:of|the|and|in|on|for|[A-Z][a-zA-Z']+)){0,6}") + + +class SpacySpan(Protocol): + """Small protocol for the spaCy span attributes used by this module.""" + + text: str + + +class SpacyEntity(SpacySpan, Protocol): + """Small protocol for the spaCy entity attributes used by this module.""" + + label_: str + + +class SpacyDoc(Protocol): + """Small protocol for the spaCy doc attributes used by this module.""" + + ents: Iterable[SpacyEntity] + noun_chunks: Iterable[SpacySpan] + + +class SpacyLanguage(Protocol): + """Small protocol for a callable spaCy language pipeline.""" + + def __call__(self, text: str) -> SpacyDoc: + """Parse text into a spaCy-like doc.""" + + +class YakeExtractor(Protocol): + """Small protocol for the YAKE extractor used by this module.""" + + def extract_keywords(self, text: str) -> Iterable[tuple[str, float]]: + """Return YAKE keyword tuples.""" + + +class YakeExtractorFactory(Protocol): + """Callable constructor protocol for YAKE keyword extractors.""" + + def __call__(self, *, lan: str, n: int, dedupLim: float, top: int) -> YakeExtractor: # noqa: N803 + """Create a YAKE keyword extractor. + + Args: + lan (str): Language code passed to YAKE. + n (int): Maximum n-gram size to extract. + dedupLim (float): Deduplication similarity threshold. + top (int): Maximum number of keyphrases to return. + + Returns: + YakeExtractor: The constructed keyword extractor. + """ + + +def strip_leading_articles(phrase_norm: str) -> str: + """Remove one leading English article from a normalized phrase. + + Args: + phrase_norm (str): Normalized phrase text to strip. + + Returns: + str: The phrase with a single leading ``the``, ``a``, or ``an`` removed. + """ + tokens_ = phrase_norm.split() + if tokens_ and tokens_[0] in {"the", "a", "an"}: + tokens_ = tokens_[1:] + return " ".join(tokens_) + + +def normalize_candidate_phrase( + phrase_text: str, + config: EbookSearchConfig, + *, + min_tokens: int | None = None, + max_tokens: int | None = None, + strip_leading_article: bool = False, +) -> tuple[str, str, int] | None: + """Normalize a candidate phrase and validate token bounds. + + Args: + phrase_text (str): Raw phrase text to normalize. + config (EbookSearchConfig): Runtime phrase-tuning settings. + min_tokens (int | None): Minimum token count override; defaults to ``config.phrase_min_tokens``. + max_tokens (int | None): Maximum token count override; defaults to ``config.phrase_max_tokens``. + strip_leading_article (bool): Whether to drop a single leading English article. + + Returns: + tuple[str, str, int] | None: Display text, normalized phrase, and token count, or ``None`` + when the phrase falls outside the token bounds or is ignored. + """ + normalized_tokens = tokenize_with_offsets(phrase_text) + start = 0 + if strip_leading_article and normalized_tokens and normalized_tokens[0].text in {"the", "a", "an"}: + start = 1 + + selected_tokens = normalized_tokens[start:] + min_count = config.phrase_min_tokens if min_tokens is None else min_tokens + max_count = config.phrase_max_tokens if max_tokens is None else max_tokens + if len(selected_tokens) < min_count or len(selected_tokens) > max_count: + return None + + phrase_norm = " ".join(token.text for token in selected_tokens) + if phrase_norm in get_ignored_phrases(): + return None + + display_text = phrase_text[selected_tokens[0].start_char : selected_tokens[-1].end_char].strip() + return display_text or phrase_norm, phrase_norm, len(selected_tokens) + + +def extract_raw_ngrams(text: str, config: EbookSearchConfig) -> dict[str, PhraseCandidate]: + """Extract raw normalized n-grams as high-recall candidates. + + Args: + text (str): Book text to slide n-gram windows over. + config (EbookSearchConfig): Runtime phrase-tuning settings. + + Returns: + dict[str, PhraseCandidate]: Candidates keyed by normalized phrase, with raw counts. + """ + tokens_ = tokenize(text) + out: dict[str, PhraseCandidate] = {} + for ngram_size in range(config.phrase_min_tokens, config.phrase_max_tokens + 1): + for start in range(len(tokens_) - ngram_size + 1): + normalized = normalize_candidate_phrase(" ".join(tokens_[start : start + ngram_size]), config) + if normalized is None: + continue + phrase_text, phrase_norm, token_count = normalized + item = out.setdefault( + phrase_norm, + PhraseCandidate( + phrase_text=phrase_text, + phrase_norm=phrase_norm, + token_count=token_count, + source_raw_ngram=True, + ), + ) + item.raw_count += 1 + return out + + + + +def extract_yake_candidates( + book_text: str, + config: EbookSearchConfig, + top_k: int = 1000, +) -> dict[str, PhraseCandidate]: + """Extract YAKE keyphrases when the optional YAKE package is installed. + + Args: + book_text (str): Full book text to extract keyphrases from. + config (EbookSearchConfig): Runtime phrase-tuning settings. + top_k (int): Maximum number of YAKE keyphrases to request. + + Returns: + dict[str, PhraseCandidate]: Candidates keyed by normalized phrase, with YAKE scores. + """ + extractor = KeywordExtractor(lan="en", n=config.phrase_max_tokens, dedupLim=0.85, top=top_k) + out: dict[str, PhraseCandidate] = {} + for phrase_text, yake_score in extractor.extract_keywords(book_text): + normalized = normalize_candidate_phrase(phrase_text, config) + if normalized is None: + continue + display_text, phrase_norm, token_count = normalized + out[phrase_norm] = PhraseCandidate( + phrase_text=display_text, + phrase_norm=phrase_norm, + token_count=token_count, + source_yake=True, + yake_score=float(yake_score), + ) + return out + + +def extract_spacy_candidates( + book_text: str, + nlp: SpacyLanguage, + config: EbookSearchConfig, +) -> dict[str, PhraseCandidate]: + """Extract spaCy named entities and noun chunks from one text block. + + Args: + book_text (str): Text block to parse with spaCy. + nlp (SpacyLanguage): Callable spaCy language pipeline. + config (EbookSearchConfig): Runtime phrase-tuning settings. + + Returns: + dict[str, PhraseCandidate]: Candidates keyed by normalized phrase from entities and noun chunks. + """ + out: dict[str, PhraseCandidate] = {} + doc = nlp(book_text) + + for ent in doc.ents: + normalized = normalize_candidate_phrase( + ent.text, + config, + max_tokens=config.phrase_max_entity_tokens, + ) + if normalized is None: + continue + phrase_text, phrase_norm, token_count = normalized + out[phrase_norm] = PhraseCandidate( + phrase_text=phrase_text, + phrase_norm=phrase_norm, + token_count=token_count, + source_spacy_ner=True, + spacy_label=ent.label_, + ) + + for chunk in doc.noun_chunks: + normalized = normalize_candidate_phrase(chunk.text, config, strip_leading_article=True) + if normalized is None: + continue + phrase_text, phrase_norm, token_count = normalized + out[phrase_norm] = PhraseCandidate( + phrase_text=phrase_text, + phrase_norm=phrase_norm, + token_count=token_count, + source_spacy_noun_chunk=True, + ) + return out + + +def extract_capitalized_phrases(original_text: str, config: EbookSearchConfig) -> dict[str, PhraseCandidate]: + """Extract capitalized phrase runs that often carry fictional terms. + + Args: + original_text (str): Original-case book text to scan for capitalized runs. + config (EbookSearchConfig): Runtime phrase-tuning settings. + + Returns: + dict[str, PhraseCandidate]: Candidates keyed by normalized phrase from capitalized runs. + """ + out: dict[str, PhraseCandidate] = {} + for match in CAPITALIZED_PHRASE_RE.finditer(original_text): + phrase_text = match.group(0).strip() + normalized = normalize_candidate_phrase( + phrase_text, + config, + max_tokens=config.phrase_max_entity_tokens, + ) + if normalized is None: + continue + display_text, phrase_norm, token_count = normalized + out[phrase_norm] = PhraseCandidate( + phrase_text=display_text, + phrase_norm=phrase_norm, + token_count=token_count, + source_capitalized=True, + ) + return out + + +def extract_metadata_candidates( + metadata: Mapping[str, object] | None, + config: EbookSearchConfig, +) -> dict[str, PhraseCandidate]: + """Extract phrases from book metadata values such as title, author, and series. + + Args: + metadata (Mapping[str, object] | None): Book metadata values, or ``None`` when unavailable. + config (EbookSearchConfig): Runtime phrase-tuning settings. + + Returns: + dict[str, PhraseCandidate]: Candidates keyed by normalized phrase from metadata values. + """ + if metadata is None: + return {} + + out: dict[str, PhraseCandidate] = {} + for value in metadata.values(): + if value is None: + continue + phrase_text = str(value).strip() + normalized = normalize_candidate_phrase( + phrase_text, + config, + max_tokens=config.phrase_max_entity_tokens, + ) + if normalized is None: + continue + display_text, phrase_norm, token_count = normalized + out[phrase_norm] = PhraseCandidate( + phrase_text=display_text, + phrase_norm=phrase_norm, + token_count=token_count, + source_metadata=True, + ) + return out + + +def merge_candidate_sources(*sources: Mapping[str, PhraseCandidate]) -> dict[str, PhraseCandidate]: + """Merge candidate dictionaries by normalized phrase. + + Args: + *sources (Mapping[str, PhraseCandidate]): Candidate maps to combine, keyed by normalized phrase. + + Returns: + dict[str, PhraseCandidate]: One merged candidate per normalized phrase. + """ + merged: dict[str, PhraseCandidate] = {} + for source in sources: + for phrase_norm, item in source.items(): + existing = merged.setdefault( + phrase_norm, + PhraseCandidate( + phrase_text=item.phrase_text, + phrase_norm=phrase_norm, + token_count=item.token_count, + ), + ) + merge_candidate(existing, item) + return merged + + +def merge_candidate(existing: PhraseCandidate, item: PhraseCandidate) -> None: + """Merge one candidate into an existing candidate object. + + Args: + existing (PhraseCandidate): Candidate mutated in place to absorb ``item``. + item (PhraseCandidate): Candidate whose sources, counts, and scores are merged in. + """ + existing.source_raw_ngram = existing.source_raw_ngram or item.source_raw_ngram + existing.source_yake = existing.source_yake or item.source_yake + existing.source_spacy_ner = existing.source_spacy_ner or item.source_spacy_ner + existing.source_spacy_noun_chunk = existing.source_spacy_noun_chunk or item.source_spacy_noun_chunk + existing.source_capitalized = existing.source_capitalized or item.source_capitalized + existing.source_metadata = existing.source_metadata or item.source_metadata + existing.raw_count += item.raw_count + if item.yake_score is not None: + existing.yake_score = item.yake_score + if item.spacy_label: + existing.spacy_label = item.spacy_label + + +def enrich_with_frequency_and_chapter_counts( + candidates: Mapping[str, PhraseCandidate], + chapters: Sequence[str], +) -> dict[str, PhraseCandidate]: + """Add raw occurrence and chapter-spread counts to candidates. + + Args: + candidates (Mapping[str, PhraseCandidate]): Candidates to enrich, keyed by normalized phrase. + chapters (Sequence[str]): Chapter-like text blocks used to count occurrences and spread. + + Returns: + dict[str, PhraseCandidate]: Candidates with updated ``raw_count`` and ``chapter_count`` values. + """ + if not candidates: + return {} + + candidate_sets_by_size: dict[int, set[str]] = defaultdict(set) + for phrase_norm, candidate in candidates.items(): + candidate_sets_by_size[candidate.token_count].add(phrase_norm) + + total_counts: defaultdict[str, int] = defaultdict(int) + chapter_counts: defaultdict[str, int] = defaultdict(int) + for chapter in chapters: + seen_in_chapter: set[str] = set() + chapter_tokens = tokenize(chapter) + for ngram_size, candidate_norms in candidate_sets_by_size.items(): + for start in range(len(chapter_tokens) - ngram_size + 1): + phrase_norm = " ".join(chapter_tokens[start : start + ngram_size]) + if phrase_norm not in candidate_norms: + continue + total_counts[phrase_norm] += 1 + seen_in_chapter.add(phrase_norm) + for phrase_norm in seen_in_chapter: + chapter_counts[phrase_norm] += 1 + + enriched = dict(candidates) + for phrase_norm, candidate in enriched.items(): + candidate.raw_count = max(candidate.raw_count, total_counts[phrase_norm]) + candidate.chapter_count = chapter_counts[phrase_norm] + return enriched + + +def filter_storable_candidates( + candidates: Mapping[str, PhraseCandidate], + config: EbookSearchConfig, +) -> tuple[dict[str, PhraseCandidate], int, int, int]: + """Remove candidates that should not be persisted. + + Args: + candidates (Mapping[str, PhraseCandidate]): Candidates to filter, keyed by normalized phrase. + config (EbookSearchConfig): Runtime phrase-tuning settings. + + Returns: + tuple[dict[str, PhraseCandidate], int, int, int]: The storable candidates followed by the counts + dropped for being too short, too rare, and too common. + """ + min_raw_count = minimum_candidate_raw_count(config) + filtered: dict[str, PhraseCandidate] = {} + too_short = 0 + too_rare = 0 + too_common = 0 + for phrase_norm, candidate in candidates.items(): + if candidate.token_count < config.phrase_min_tokens: + too_short += 1 + continue + if candidate.raw_count < min_raw_count: + too_rare += 1 + continue + if is_most_common_word_phrase(phrase_norm): + too_common += 1 + continue + filtered[phrase_norm] = candidate + return filtered, too_short, too_rare, too_common + + +def minimum_candidate_raw_count(config: EbookSearchConfig) -> int: + """Return the minimum occurrence count required before storing a candidate. + + Args: + config (EbookSearchConfig): Runtime phrase-tuning settings. + + Returns: + int: The minimum raw occurrence count, never less than 1. + """ + return max(config.phrase_raw_ngram_min_count, 1) + + +def is_most_common_word_phrase(phrase_norm: str) -> bool: + """Return whether every token in a normalized phrase is a common word. + + Args: + phrase_norm (str): Normalized phrase text to inspect. + + Returns: + bool: True when the phrase is non-empty and every token is a common word. + """ + tokens_ = phrase_norm.split() + common_words = get_most_common_words() + return bool(tokens_) and all(token in common_words for token in tokens_) + + +def score_candidate(candidate: PhraseCandidate, config: EbookSearchConfig) -> float: + """Score a phrase candidate before LLM judging. + + Args: + candidate (PhraseCandidate): Candidate to score. + config (EbookSearchConfig): Runtime phrase-tuning settings. + + Returns: + float: Combined score from sources, frequency, and length, less any penalties. + """ + score = source_score(candidate) + frequency_score(candidate, config) + token_count_score(candidate, config) + if candidate.phrase_norm in get_ignored_phrases(): + score -= 100.0 + if has_bad_start(candidate.phrase_norm): + score -= BAD_START_SCORE_PENALTY + if has_bad_end(candidate.phrase_norm): + score -= BAD_END_SCORE_PENALTY + return score + + +def has_bad_start(phrase_norm: str) -> bool: + """Return whether a normalized phrase starts with a bad starting token. + + Args: + phrase_norm (str): Normalized phrase text to inspect. + + Returns: + bool: True when the first token is a known bad starting token. + """ + tokens_ = phrase_norm.split() + return bool(tokens_ and tokens_[0] in get_bad_starts()) + + +def has_bad_end(phrase_norm: str) -> bool: + """Return whether a normalized phrase ends with a bad ending token. + + Args: + phrase_norm (str): Normalized phrase text to inspect. + + Returns: + bool: True when the last token is a known bad ending token. + """ + tokens_ = phrase_norm.split() + return bool(tokens_ and tokens_[-1] in get_bad_ends()) + + +def source_score(candidate: PhraseCandidate) -> float: + """Return the score contribution from extraction sources. + + Args: + candidate (PhraseCandidate): Candidate whose enabled sources are weighted. + + Returns: + float: Summed weight of the candidate's enabled extraction sources. + """ + return sum( + weight + for enabled, weight in ( + (candidate.source_yake, 2.0), + (candidate.source_spacy_ner, 2.5), + (candidate.source_spacy_noun_chunk, 1.5), + (candidate.source_capitalized, 2.0), + (candidate.source_raw_ngram, 0.5), + ) + if enabled + ) + + +def frequency_score(candidate: PhraseCandidate, config: EbookSearchConfig) -> float: + """Return the score contribution from frequency and chapter spread. + + Args: + candidate (PhraseCandidate): Candidate whose counts are scored. + config (EbookSearchConfig): Runtime phrase-tuning settings holding score thresholds. + + Returns: + float: Summed weight for each frequency and chapter-spread threshold the candidate meets. + """ + return sum( + weight + for count, threshold, weight in ( + (candidate.raw_count, config.phrase_raw_count_score_threshold, 1.0), + (candidate.raw_count, config.phrase_raw_count_high_score_threshold, 1.0), + (candidate.chapter_count, config.phrase_chapter_count_score_threshold, 1.0), + (candidate.chapter_count, config.phrase_chapter_count_high_score_threshold, 1.0), + ) + if count >= threshold + ) + + +def token_count_score(candidate: PhraseCandidate, config: EbookSearchConfig) -> float: + """Return the score contribution from phrase length. + + Args: + candidate (PhraseCandidate): Candidate whose token count is scored. + config (EbookSearchConfig): Runtime phrase-tuning settings holding the max token bound. + + Returns: + float: Length-based score contribution, which may be negative for over- or under-length phrases. + """ + if candidate.token_count == 1: + return -0.5 + if candidate.token_count in {2, 3, 4}: + return 0.5 + if candidate.token_count > config.phrase_max_tokens: + return -1.0 + return 0.0 + + +def get_sample_contexts(normalized_book_text: str, phrase_norm: str, max_contexts: int = 5) -> list[str]: + """Return normalized context snippets containing a candidate phrase. + + ``normalized_book_text`` is expected to already be ``normalize_text``-ed by the caller + so the whole book is not re-normalized for every phrase. + + Args: + normalized_book_text (str): Whole book text, already normalized, to search. + phrase_norm (str): Normalized phrase to find contexts around. + max_contexts (int): Maximum number of context snippets to return. + + Returns: + list[str]: Up to ``max_contexts`` normalized snippets surrounding the phrase. + """ + contexts: list[str] = [] + start = 0 + while len(contexts) < max_contexts: + index = normalized_book_text.find(phrase_norm, start) + if index == -1: + break + left = max(0, index - 300) + right = min(len(normalized_book_text), index + len(phrase_norm) + 300) + contexts.append(normalized_book_text[left:right]) + start = index + len(phrase_norm) + return contexts + + +def candidate_source_names(candidate: PhraseCandidate) -> list[str]: + """Return enabled source names for an extracted candidate. + + Args: + candidate (PhraseCandidate): Candidate whose enabled sources are listed. + + Returns: + list[str]: Names of the extraction sources that produced the candidate. + """ + names: list[str] = [] + if candidate.source_raw_ngram: + names.append("raw_ngram") + if candidate.source_yake: + names.append("yake") + if candidate.source_spacy_ner: + names.append("spacy_ner") + if candidate.source_spacy_noun_chunk: + names.append("spacy_noun_chunk") + if candidate.source_capitalized: + names.append("capitalized") + if candidate.source_metadata: + names.append("metadata") + return names + + +def extract_phrase_candidates_for_book( + book_text: str, + chapters: Sequence[str], + config: EbookSearchConfig, + *, + nlp: SpacyLanguage | None = None, + metadata: Mapping[str, object] | None = None, +) -> list[PhraseCandidate]: + """Extract, score, and limit phrase candidates for one book. + + Args: + book_text (str): Full book text used for most extraction sources. + chapters (Sequence[str]): Chapter-like text blocks used for spaCy and frequency counts. + config (EbookSearchConfig): Runtime phrase-tuning settings. + nlp (SpacyLanguage | None): Optional spaCy pipeline for entity and noun-chunk sources. + metadata (Mapping[str, object] | None): Optional book metadata used as a candidate source. + + Returns: + list[PhraseCandidate]: Scored candidates sorted best-first and capped per book. + """ + started_at = perf_counter() + logger.info( + "ebook_phrase_candidate_extract_start chapters=%s chars=%s min_tokens=%s max_tokens=%s max_candidates=%s", + len(chapters), + len(book_text), + config.phrase_min_tokens, + config.phrase_max_tokens, + config.protected_phrase_max_candidates_per_book, + ) + raw_started_at = perf_counter() + raw = extract_raw_ngrams(book_text, config) + logger.info( + "ebook_phrase_candidate_extract_raw_complete candidates=%s duration_ms=%.1f", + len(raw), + (perf_counter() - raw_started_at) * 1000, + ) + yake_started_at = perf_counter() + yake_candidates = extract_yake_candidates(book_text, config) + logger.info( + "ebook_phrase_candidate_extract_yake_complete candidates=%s duration_ms=%.1f", + len(yake_candidates), + (perf_counter() - yake_started_at) * 1000, + ) + spacy_candidates: dict[str, PhraseCandidate] = {} + if nlp is not None: + spacy_started_at = perf_counter() + for chapter in chapters: + spacy_candidates = merge_candidate_sources(spacy_candidates, extract_spacy_candidates(chapter, nlp, config)) + logger.info( + "ebook_phrase_candidate_extract_spacy_complete candidates=%s duration_ms=%.1f", + len(spacy_candidates), + (perf_counter() - spacy_started_at) * 1000, + ) + capitalized_started_at = perf_counter() + capitalized = extract_capitalized_phrases(book_text, config) + logger.info( + "ebook_phrase_candidate_extract_capitalized_complete candidates=%s duration_ms=%.1f", + len(capitalized), + (perf_counter() - capitalized_started_at) * 1000, + ) + metadata_candidates = extract_metadata_candidates(metadata, config) + + candidates = merge_candidate_sources(raw, yake_candidates, spacy_candidates, capitalized, metadata_candidates) + enriched_started_at = perf_counter() + candidates = enrich_with_frequency_and_chapter_counts(candidates, chapters) + pre_filter_count = len(candidates) + candidates, filtered_too_short, filtered_too_rare, filtered_too_common = filter_storable_candidates( + candidates, + config, + ) + for candidate in candidates.values(): + candidate.candidate_score = score_candidate(candidate, config) + + limited = sorted(candidates.values(), key=lambda item: item.candidate_score, reverse=True)[ + : config.protected_phrase_max_candidates_per_book + ] + logger.info( + "ebook_phrase_candidate_extract_complete raw=%s yake=%s spacy=%s capitalized=%s metadata=%s " + "merged=%s filtered_too_short=%s filtered_too_rare=%s filtered_too_common=%s min_uses=%s " + "storable=%s limited=%s enrich_score_ms=%.1f duration_ms=%.1f", + len(raw), + len(yake_candidates), + len(spacy_candidates), + len(capitalized), + len(metadata_candidates), + pre_filter_count, + filtered_too_short, + filtered_too_rare, + filtered_too_common, + minimum_candidate_raw_count(config), + len(candidates), + len(limited), + (perf_counter() - enriched_started_at) * 1000, + (perf_counter() - started_at) * 1000, + ) + return limited diff --git a/python/ebook_search/protected_phrases/generate_ngrams.py b/python/ebook_search/protected_phrases/generate_ngrams.py new file mode 100644 index 0000000..c3a6499 --- /dev/null +++ b/python/ebook_search/protected_phrases/generate_ngrams.py @@ -0,0 +1,274 @@ +"""Book-level orchestration for candidate n-gram generation and recalculation.""" + +from __future__ import annotations + +import logging +from time import perf_counter +from typing import TYPE_CHECKING + +from sqlalchemy import select + +from python.ebook_search.protected_phrases.extraction import extract_phrase_candidates_for_book +from python.ebook_search.protected_phrases.models import ( + BookCandidateResult, + PhraseCandidateGenerationResult, + PhraseRecalculationResult, +) +from python.ebook_search.protected_phrases.store import ( + delete_phrase_data_for_book, + load_book_chapter_texts, + metadata_for_source, + prune_unstorable_unjudged_candidate_phrases, + save_candidate_to_db, +) +from python.orm.richie import EbookSource + +if TYPE_CHECKING: + from collections.abc import Mapping, Sequence + + from sqlalchemy.orm import Session + + from python.ebook_search.config import EbookSearchConfig + from python.ebook_search.protected_phrases.extraction import SpacyLanguage + from python.orm.richie import EbookCandidatePhrase + +logger = logging.getLogger(__name__) + + +def generate_candidate_phrases_for_books( + session: Session, + config: EbookSearchConfig, + *, + nlp: SpacyLanguage | None = None, +) -> PhraseCandidateGenerationResult: + """Create or refresh candidate phrases for indexed books without calling the LLM judge. + + Args: + session (Session): Active database session. + config (EbookSearchConfig): Runtime phrase-tuning settings. + nlp (SpacyLanguage | None): Optional spaCy pipeline for entity and noun-chunk sources. + + Returns: + PhraseCandidateGenerationResult: Per-corpus counts of books seen, built, and candidates stored. + """ + sources = session.scalars(select(EbookSource).order_by(EbookSource.id)).all() + books_seen = len(sources) + logger.info( + "ebook_candidate_phrase_generation_start books_seen=%s min_tokens=%s max_tokens=%s max_candidates_per_book=%s", + books_seen, + config.phrase_min_tokens, + config.phrase_max_tokens, + config.protected_phrase_max_candidates_per_book, + ) + + outcomes = [generate_candidates_for_source(session, source, config, nlp=nlp) for source in sources] + + result = PhraseCandidateGenerationResult( + books_seen=books_seen, + books_built=sum(1 for outcome in outcomes if outcome.built), + candidate_phrases=sum(outcome.candidates for outcome in outcomes), + ) + logger.info( + "ebook_candidate_phrase_generation_complete books_seen=%s books_built=%s candidate_total=%s", + result.books_seen, + result.books_built, + result.candidate_phrases, + ) + return result + + +def generate_candidates_for_source( + session: Session, + source: EbookSource, + config: EbookSearchConfig, + *, + nlp: SpacyLanguage | None = None, +) -> BookCandidateResult: + """Generate and store candidate phrases for one book, managing its own transaction. + + Commits on success; rolls back and re-raises on error so callers stop the backfill. + + Args: + session (Session): Active database session. + source (EbookSource): Indexed book to generate candidates for. + config (EbookSearchConfig): Runtime phrase-tuning settings. + nlp (SpacyLanguage | None): Optional spaCy pipeline for entity and noun-chunk sources. + + Returns: + BookCandidateResult: Candidate count and whether the book was committed. + """ + book_started_at = perf_counter() + logger.info( + "ebook_candidate_phrase_generation_book_start source_id=%s title=%r", + source.id, + source.title, + ) + try: + chapters = load_book_chapter_texts(session, source.id) + if not chapters: + logger.warning("ebook_candidate_phrase_generation_book_empty source_id=%s", source.id) + return BookCandidateResult() + book_text = "\n\n".join(chapters) + logger.info( + "ebook_candidate_phrase_generation_book_loaded source_id=%s chapters=%s chars=%s", + source.id, + len(chapters), + len(book_text), + ) + + candidates = generate_candidate_phrases_for_book( + session, + source.id, + series_id=None, + book_text=book_text, + chapters=chapters, + config=config, + nlp=nlp, + metadata=metadata_for_source(source), + ) + session.commit() + except Exception: + session.rollback() + logger.exception("ebook_candidate_phrase_generation_book_failed source_id=%s", source.id) + raise + + logger.info( + "ebook_candidate_phrase_generation_book_committed source_id=%s candidates=%s duration_ms=%.1f", + source.id, + len(candidates), + (perf_counter() - book_started_at) * 1000, + ) + return BookCandidateResult(candidates=len(candidates), built=True) + + +def recalculate_candidate_phrases_for_book( + session: Session, + source: EbookSource, + config: EbookSearchConfig, + *, + nlp: SpacyLanguage | None = None, +) -> PhraseRecalculationResult: + """Remove all book phrase data, regenerate candidates, and commit the completed book. + + Args: + session (Session): Active database session. + source (EbookSource): Indexed book to recalculate. + config (EbookSearchConfig): Runtime phrase-tuning settings. + nlp (SpacyLanguage | None): Optional spaCy pipeline for entity and noun-chunk sources. + + Returns: + PhraseRecalculationResult: Deleted-row counts and the number of candidates regenerated. + """ + started_at = perf_counter() + logger.info( + "ebook_candidate_phrase_recalculation_start source_id=%s title=%r", + source.id, + source.title, + ) + try: + deleted = delete_phrase_data_for_book(session, source.id) + chapters = load_book_chapter_texts(session, source.id) + if not chapters: + logger.warning("ebook_candidate_phrase_recalculation_book_empty source_id=%s", source.id) + session.commit() + return PhraseRecalculationResult( + book_id=source.id, + deleted_candidates=deleted.deleted_candidates, + deleted_protected_phrases=deleted.deleted_protected_phrases, + deleted_aliases=deleted.deleted_aliases, + deleted_mentions=deleted.deleted_mentions, + candidate_phrases=0, + ) + + candidates = generate_candidate_phrases_for_book( + session, + source.id, + series_id=None, + book_text="\n\n".join(chapters), + chapters=chapters, + config=config, + nlp=nlp, + metadata=metadata_for_source(source), + ) + session.commit() + except Exception: + session.rollback() + logger.exception("ebook_candidate_phrase_recalculation_failed source_id=%s", source.id) + raise + + result = PhraseRecalculationResult( + book_id=source.id, + deleted_candidates=deleted.deleted_candidates, + deleted_protected_phrases=deleted.deleted_protected_phrases, + deleted_aliases=deleted.deleted_aliases, + deleted_mentions=deleted.deleted_mentions, + candidate_phrases=len(candidates), + ) + logger.info( + "ebook_candidate_phrase_recalculation_complete source_id=%s deleted_candidates=%s " + "deleted_protected=%s deleted_aliases=%s deleted_mentions=%s candidates=%s duration_ms=%.1f", + source.id, + result.deleted_candidates, + result.deleted_protected_phrases, + result.deleted_aliases, + result.deleted_mentions, + result.candidate_phrases, + (perf_counter() - started_at) * 1000, + ) + return result + + +def generate_candidate_phrases_for_book( + session: Session, + book_id: int, + series_id: int | None, + book_text: str, + chapters: Sequence[str], + config: EbookSearchConfig, + *, + nlp: SpacyLanguage | None = None, + metadata: Mapping[str, object] | None = None, +) -> list[EbookCandidatePhrase]: + """Extract and store candidate phrases for one book without LLM judging. + + Args: + session (Session): Active database session. + book_id (int): Book the candidates belong to. + series_id (int | None): Series scope for the stored candidates. + book_text (str): Full book text used for extraction. + chapters (Sequence[str]): Chapter-like text blocks used for frequency counts. + config (EbookSearchConfig): Runtime phrase-tuning settings. + nlp (SpacyLanguage | None): Optional spaCy pipeline for entity and noun-chunk sources. + metadata (Mapping[str, object] | None): Optional book metadata used as a candidate source. + + Returns: + list[EbookCandidatePhrase]: The stored candidate phrase rows. + """ + started_at = perf_counter() + limited_candidates = extract_phrase_candidates_for_book( + book_text, + chapters, + config, + nlp=nlp, + metadata=metadata, + ) + save_started_at = perf_counter() + pruned_count = prune_unstorable_unjudged_candidate_phrases(session, book_id, config) + logger.info( + "ebook_candidate_phrase_save_start book_id=%s candidates=%s pruned_unstorable=%s", + book_id, + len(limited_candidates), + pruned_count, + ) + rows = [ + save_candidate_to_db(session, book_id, series_id, candidate, judgment=None) for candidate in limited_candidates + ] + session.flush() + logger.info( + "ebook_candidate_phrase_generation_complete book_id=%s candidates=%s save_ms=%.1f duration_ms=%.1f", + book_id, + len(rows), + (perf_counter() - save_started_at) * 1000, + (perf_counter() - started_at) * 1000, + ) + return rows diff --git a/python/ebook_search/protected_phrases/judge_ngrams.py b/python/ebook_search/protected_phrases/judge_ngrams.py new file mode 100644 index 0000000..1be6e17 --- /dev/null +++ b/python/ebook_search/protected_phrases/judge_ngrams.py @@ -0,0 +1,481 @@ +"""Book-level orchestration for LLM judging and promotion of candidate phrases.""" + +from __future__ import annotations + +import json +import logging +import re +from time import perf_counter +from typing import TYPE_CHECKING + +from sqlalchemy import select + +from python.ebook_search.llm_interface import request_chat_completion +from python.ebook_search.protected_phrases.extraction import ( + candidate_source_names, + get_sample_contexts, + is_most_common_word_phrase, + minimum_candidate_raw_count, +) +from python.ebook_search.protected_phrases.matching import index_chunk_phrase_mentions_for_book +from python.ebook_search.protected_phrases.models import BookJudgmentResult, LLMJudgment, PhraseJudgmentBackfillResult +from python.ebook_search.protected_phrases.store import ( + count_protected_phrases, + count_unjudged_candidates, + load_book_text, + load_candidates_for_judgment, + phrase_candidate_from_row, + save_candidate_to_db, + upsert_protected_phrase, +) +from python.ebook_search.protected_phrases.text_normalization import normalize_text +from python.orm.richie import EbookSource + +if TYPE_CHECKING: + from sqlalchemy.orm import Session + + from python.ebook_search.config import EbookSearchConfig + from python.ebook_search.protected_phrases.models import PhraseCandidate + from python.orm.richie import EbookCandidatePhrase, EbookProtectedPhrase + +JSON_OBJECT_RE = re.compile(r"\{.*\}", re.DOTALL) + + + +logger = logging.getLogger(__name__) + + +def judge_candidate_phrases_for_books( + session: Session, + config: EbookSearchConfig, +) -> PhraseJudgmentBackfillResult: + """Judge stored candidate phrases and promote accepted phrases for indexed books. + + Args: + session (Session): Active database session. + config (EbookSearchConfig): Runtime phrase-tuning settings. + + Returns: + PhraseJudgmentBackfillResult: Per-corpus counts of books judged, failures, candidates, + protected phrases, and mentions. + """ + source_ids = session.scalars(select(EbookSource.id).order_by(EbookSource.id)).all() + books_seen = len(source_ids) + logger.info( + "ebook_candidate_phrase_judgment_start books_seen=%s llm_candidates_per_book=%s " + "target_protected_per_book=%s confidence_threshold=%.2f min_tokens=%s min_uses=%s", + books_seen, + config.protected_phrase_llm_candidates_per_book, + config.phrase_target_protected_per_book, + config.protected_phrase_confidence_threshold, + config.phrase_min_tokens, + minimum_candidate_raw_count(config), + ) + + outcomes = [judge_book_for_backfill(session, source_id, config) for source_id in source_ids] + + result = PhraseJudgmentBackfillResult( + books_seen=books_seen, + books_judged=sum(1 for outcome in outcomes if outcome.committed), + books_failed=sum(1 for outcome in outcomes if outcome.failed), + candidates_judged=sum(outcome.judged for outcome in outcomes), + protected_phrases=sum(outcome.protected for outcome in outcomes), + phrase_mentions=sum(outcome.mentions for outcome in outcomes), + ) + logger.info( + "ebook_candidate_phrase_judgment_complete books_seen=%s books_judged=%s books_failed=%s " + "candidates_judged=%s protected=%s mentions=%s", + result.books_seen, + result.books_judged, + result.books_failed, + result.candidates_judged, + result.protected_phrases, + result.phrase_mentions, + ) + return result + + +def judge_book_for_backfill( + session: Session, + source_id: int, + config: EbookSearchConfig, +) -> BookJudgmentResult: + """Judge one book's candidates and return its outcome. + + Args: + session (Session): Active database session. + source_id (int): Book to judge candidates for. + config (EbookSearchConfig): Runtime phrase-tuning settings. + + Returns: + BookJudgmentResult: The book's judgment outcome, or an empty result when nothing was unjudged. + """ + unjudged_count = count_unjudged_candidates(session, source_id, config) + if not unjudged_count: + logger.info( + "ebook_candidate_phrase_judgment_book_skip_no_unjudged source_id=%s", + source_id, + ) + return BookJudgmentResult() + logger.info( + "ebook_candidate_phrase_judgment_book_start source_id=%s unjudged=%s", + source_id, + unjudged_count, + ) + return run_book_judgment(session, source_id, unjudged_count, config) + + +def run_book_judgment( + session: Session, + source_id: int, + unjudged_count: int, + config: EbookSearchConfig, +) -> BookJudgmentResult: + """Judge and index one book's candidates, managing its own transaction. + + Commits on success and rolls back on error, returning a :class:`BookJudgmentResult` + that describes the outcome rather than raising it to the caller. + + Args: + session (Session): Active database session. + source_id (int): Book to judge candidates for. + unjudged_count (int): Storable unjudged candidates counted before judging, for logging. + config (EbookSearchConfig): Runtime phrase-tuning settings. + + Returns: + BookJudgmentResult: Judged, protected, and mention counts with commit and failure flags. + """ + book_started_at = perf_counter() + try: + book_text = load_book_text(session, source_id) + if not book_text: + logger.warning("ebook_candidate_phrase_judgment_book_empty source_id=%s", source_id) + return BookJudgmentResult() + + judged, protected = judge_candidate_phrases_for_book( + session, + source_id, + series_id=None, + normalized_book_text=normalize_text(book_text), + config=config, + ) + if judged == 0: + logger.info( + "ebook_candidate_phrase_judgment_book_skip_no_judgments source_id=%s unjudged=%s", + source_id, + unjudged_count, + ) + return BookJudgmentResult() + + mentions = index_chunk_phrase_mentions_for_book(session, source_id, config) if protected else 0 + session.commit() + except Exception: + session.rollback() + logger.exception("ebook_candidate_phrase_judgment_book_failed source_id=%s", source_id) + return BookJudgmentResult(failed=True) + + logger.info( + "ebook_candidate_phrase_judgment_book_committed source_id=%s judged=%s protected=%s mentions=%s " + "duration_ms=%.1f", + source_id, + judged, + len(protected), + mentions, + (perf_counter() - book_started_at) * 1000, + ) + return BookJudgmentResult(judged=judged, protected=len(protected), mentions=mentions, committed=True) + + +def judge_candidate_phrases_for_book( + session: Session, + book_id: int, + series_id: int | None, + normalized_book_text: str, + config: EbookSearchConfig, +) -> tuple[int, list[EbookProtectedPhrase]]: + """Judge unjudged candidate phrase rows for one book. + + Args: + session (Session): Active database session. + book_id (int): Book whose candidates are judged. + series_id (int | None): Series scope for promoted protected phrases. + normalized_book_text (str): Whole book text, already normalized, for context lookups. + config (EbookSearchConfig): Runtime phrase-tuning settings. + + Returns: + tuple[int, list[EbookProtectedPhrase]]: Number of candidates judged and the promoted phrases. + """ + judgment_limit = config.protected_phrase_llm_candidates_per_book + if judgment_limit <= 0: + logger.info("ebook_candidate_phrase_judgment_skipped_llm_limit_zero book_id=%s", book_id) + return 0, [] + + existing_protected = count_protected_phrases(session, book_id) + target_remaining: int | None = None + if config.phrase_target_protected_per_book > 0: + target_remaining = max(config.phrase_target_protected_per_book - existing_protected, 0) + if target_remaining == 0: + logger.info( + "ebook_candidate_phrase_judgment_skipped_target_met book_id=%s existing_protected=%s target=%s", + book_id, + existing_protected, + config.phrase_target_protected_per_book, + ) + return 0, [] + + rows = load_candidates_for_judgment(session, book_id, judgment_limit, config) + logger.info( + "ebook_candidate_phrase_judgment_candidates_loaded book_id=%s candidates=%s existing_protected=%s " + "target_remaining=%s judgment_limit=%s", + book_id, + len(rows), + existing_protected, + target_remaining, + judgment_limit, + ) + + judged_count = 0 + protected: list[EbookProtectedPhrase] = [] + for row_number, row in enumerate(rows, start=1): + candidate, judgment, candidate_row = judge_candidate_row( + session, book_id, series_id, normalized_book_text, row, row_number, len(rows), config + ) + judged_count += 1 + if not should_protect_judged_candidate(row, candidate, judgment, book_id, config): + continue + protected.append(upsert_protected_phrase(session, book_id, series_id, candidate, judgment, candidate_row)) + if target_remaining is not None and len(protected) >= target_remaining: + break + + session.flush() + return judged_count, protected + + +def judge_candidate_row( + session: Session, + book_id: int, + series_id: int | None, + normalized_book_text: str, + row: EbookCandidatePhrase, + row_number: int, + total_rows: int, + config: EbookSearchConfig, +) -> tuple[PhraseCandidate, LLMJudgment, EbookCandidatePhrase]: + """Run and persist the LLM judgment for a single candidate row. + + Args: + session (Session): Active database session. + book_id (int): Book the candidate belongs to. + series_id (int | None): Series scope for the saved candidate. + normalized_book_text (str): Whole book text, already normalized, for context lookups. + row (EbookCandidatePhrase): Stored candidate row to judge. + row_number (int): 1-based position of the row in the batch, for logging. + total_rows (int): Total rows in the batch, for logging. + config (EbookSearchConfig): Runtime phrase-tuning settings. + + Returns: + tuple[PhraseCandidate, LLMJudgment, EbookCandidatePhrase]: The candidate, its judgment, + and the persisted candidate row. + """ + row_started_at = perf_counter() + candidate = phrase_candidate_from_row(row) + candidate.sample_contexts = row.sample_contexts or get_sample_contexts(normalized_book_text, candidate.phrase_norm) + logger.info( + "ebook_candidate_phrase_judgment_candidate_start book_id=%s candidate_id=%s row_number=%s rows=%s " + "phrase=%r score=%.3f raw_count=%s chapter_count=%s", + book_id, + row.id, + row_number, + total_rows, + candidate.phrase_norm, + candidate.candidate_score, + candidate.raw_count, + candidate.chapter_count, + ) + judgment = judge_candidate_with_llm(candidate, config) + candidate_row = save_candidate_to_db(session, book_id, series_id, candidate, judgment=judgment) + logger.info( + "ebook_candidate_phrase_judgment_candidate_complete book_id=%s candidate_id=%s phrase=%r keep=%s " + "confidence=%.3f importance=%.3f category=%r duration_ms=%.1f", + book_id, + row.id, + candidate.phrase_norm, + judgment.keep, + judgment.confidence, + judgment.importance, + judgment.category, + (perf_counter() - row_started_at) * 1000, + ) + return candidate, judgment, candidate_row + + +def should_protect_judged_candidate( + row: EbookCandidatePhrase, + candidate: PhraseCandidate, + judgment: LLMJudgment, + book_id: int, + config: EbookSearchConfig, +) -> bool: + """Report whether a judged candidate qualifies to become a protected phrase. + + Args: + row (EbookCandidatePhrase): Stored candidate row the judgment came from. + candidate (PhraseCandidate): In-memory candidate that was judged. + judgment (LLMJudgment): Judge decision for the candidate. + book_id (int): Book the candidate belongs to, for logging. + config (EbookSearchConfig): Runtime phrase-tuning settings. + + Returns: + bool: True when the judged candidate should be promoted to a protected phrase. + """ + if not judgment.keep or judgment.confidence < config.protected_phrase_confidence_threshold: + return False + accepted_norm = normalize_text(judgment.canonical or candidate.phrase_text) + accepted_token_count = len(accepted_norm.split()) + if accepted_token_count < config.phrase_min_tokens: + logger.info( + "ebook_candidate_phrase_judgment_candidate_skip_short_canonical book_id=%s candidate_id=%s " + "phrase=%r canonical=%r token_count=%s min_tokens=%s", + book_id, + row.id, + candidate.phrase_norm, + accepted_norm, + accepted_token_count, + config.phrase_min_tokens, + ) + return False + if is_most_common_word_phrase(accepted_norm): + logger.info( + "ebook_candidate_phrase_judgment_candidate_skip_common_canonical book_id=%s candidate_id=%s " + "phrase=%r canonical=%r", + book_id, + row.id, + candidate.phrase_norm, + accepted_norm, + ) + return False + return True + + +"""LLM judging of extracted candidate phrases.""" + + +def judge_candidate_with_llm(candidate: PhraseCandidate, config: EbookSearchConfig) -> LLMJudgment: + """Ask the configured chat model to judge one pre-extracted candidate. + + Args: + candidate (PhraseCandidate): Candidate to send to the LLM judge. + config (EbookSearchConfig): Runtime phrase-tuning settings and chat configuration. + + Returns: + LLMJudgment: The parsed structured judgment for the candidate. + """ + payload = { + "phrase": candidate.phrase_norm, + "token_count": candidate.token_count, + "sources": candidate_source_names(candidate), + "raw_count": candidate.raw_count, + "chapter_count": candidate.chapter_count, + "contexts": candidate.sample_contexts, + } + messages = [ + { + "role": "system", + "content": ( + "Judge whether a candidate phrase from a book should be protected for RAG retrieval. " + "Do not extract new phrases. Reject common grammar fragments, ordinary nonspecific phrases, " + "unstable fragments, and phrases kept only because they are frequent. Keep people, places, " + "organizations, factions, events, technologies, fictional conditions, magic systems, formal titles, " + "named concepts, and recurring world-specific terms. Return only a JSON object with keys: keep, " + "canonical, category, aliases, confidence, importance, allow_nested, suppress_children, reason." + ), + }, + {"role": "user", "content": json.dumps(payload, ensure_ascii=True)}, + ] + return parse_llm_judgment(request_chat_completion(config, messages), config) + + +def parse_llm_judgment(content: str, config: EbookSearchConfig) -> LLMJudgment: + """Parse and validate an LLM phrase-judge response. + + Args: + content (str): Raw model response text. + config (EbookSearchConfig): Runtime phrase-tuning settings supplying nesting defaults. + + Returns: + LLMJudgment: The parsed and validated judgment. + + Raises: + TypeError: If the decoded JSON body is not an object. + """ + body = json.loads(extract_json_object(content)) + if not isinstance(body, dict): + msg = "LLM phrase judge response is not a JSON object" + raise TypeError(msg) + + aliases = body.get("aliases", ()) + if not isinstance(aliases, list | tuple): + aliases = () + return LLMJudgment( + keep=bool(body.get("keep", False)), + canonical=optional_text(body.get("canonical")), + category=optional_text(body.get("category")), + aliases=tuple(str(alias) for alias in aliases if isinstance(alias, str) and alias.strip()), + confidence=clamped_float(body.get("confidence"), default=0.0), + importance=clamped_float(body.get("importance"), default=0.5), + allow_nested=bool(body.get("allow_nested", config.phrase_default_allow_nested)), + suppress_children=bool(body.get("suppress_children", config.phrase_default_suppress_children)), + reason=optional_text(body.get("reason")), + ) + + +def extract_json_object(content: str) -> str: + """Extract a JSON object from plain or fenced model output. + + Args: + content (str): Raw model response text. + + Returns: + str: The substring spanning the first JSON object. + + Raises: + ValueError: If no JSON object is found in the response. + """ + stripped = content.strip() + if stripped.startswith("{") and stripped.endswith("}"): + return stripped + match = JSON_OBJECT_RE.search(stripped) + if match is None: + msg = "LLM phrase judge response did not contain a JSON object" + raise ValueError(msg) + return match.group(0) + + +def optional_text(value: object) -> str | None: + """Return stripped text for a nullable JSON value. + + Args: + value (object): Decoded JSON value that may or may not be a string. + + Returns: + str | None: The stripped string, or ``None`` when it is not a non-empty string. + """ + if not isinstance(value, str): + return None + stripped = value.strip() + return stripped or None + + +def clamped_float(value: object, *, default: float) -> float: + """Coerce a JSON number into the 0.0 to 1.0 range. + + Args: + value (object): Decoded JSON value that may or may not be a number. + default (float): Fallback returned when ``value`` is not numeric. + + Returns: + float: The value clamped to ``[0.0, 1.0]``, or ``default`` when non-numeric. + """ + if not isinstance(value, int | float): + return default + return min(max(float(value), 0.0), 1.0) diff --git a/python/ebook_search/protected_phrases/lib.py b/python/ebook_search/protected_phrases/lib.py deleted file mode 100644 index 2a093f3..0000000 --- a/python/ebook_search/protected_phrases/lib.py +++ /dev/null @@ -1,1933 +0,0 @@ -"""Protected phrase extraction, storage, and runtime matching.""" - -from __future__ import annotations - -import importlib -import json -import logging -import re -from collections import defaultdict -from dataclasses import dataclass -from time import perf_counter -from typing import TYPE_CHECKING, Protocol - -from sqlalchemy import and_, delete, func, or_, select - -from python.ebook_search.llm_interface import request_chat_completion -from python.ebook_search.protected_phrases.config import ( - get_bad_ends, - get_bad_starts, - get_ignored_phrases, - get_most_common_words, -) -from python.ebook_search.protected_phrases.text_normalization import ( - NormalizedToken, - normalize_text, - tokenize, - tokenize_with_offsets, -) -from python.orm.richie import ( - EbookCandidatePhrase, - EbookChunk, - EbookChunkPhraseMention, - EbookPhraseAlias, - EbookProtectedPhrase, - EbookSource, -) - -if TYPE_CHECKING: - from collections.abc import Iterable, Iterator, Mapping, Sequence - from types import ModuleType - - from sqlalchemy.orm import Session - - from python.ebook_search.config import EbookSearchConfig - -logger = logging.getLogger(__name__) - - -BAD_START_SCORE_PENALTY = 3.0 -BAD_END_SCORE_PENALTY = 3.0 -SOURCE_FIELDS = ( - "source_raw_ngram", - "source_yake", - "source_spacy_ner", - "source_spacy_noun_chunk", - "source_capitalized", - "source_metadata", -) -JSON_OBJECT_RE = re.compile(r"\{.*\}", re.DOTALL) -CAPITALIZED_PHRASE_RE = re.compile(r"\b(?:[A-Z][a-zA-Z']+)(?:\s+(?:of|the|and|in|on|for|[A-Z][a-zA-Z']+)){0,6}") - - -class SpacySpan(Protocol): - """Small protocol for the spaCy span attributes used by this module.""" - - text: str - - -class SpacyEntity(SpacySpan, Protocol): - """Small protocol for the spaCy entity attributes used by this module.""" - - label_: str - - -class SpacyDoc(Protocol): - """Small protocol for the spaCy doc attributes used by this module.""" - - ents: Iterable[SpacyEntity] - noun_chunks: Iterable[SpacySpan] - - -class SpacyLanguage(Protocol): - """Small protocol for a callable spaCy language pipeline.""" - - def __call__(self, text: str) -> SpacyDoc: - """Parse text into a spaCy-like doc.""" - - -class YakeExtractor(Protocol): - """Small protocol for the YAKE extractor used by this module.""" - - def extract_keywords(self, text: str) -> Iterable[tuple[str, float]]: - """Return YAKE keyword tuples.""" - - -class YakeExtractorFactory(Protocol): - """Callable constructor protocol for YAKE keyword extractors.""" - - def __call__(self, *, lan: str, n: int, dedupLim: float, top: int) -> YakeExtractor: # noqa: N803 - """Create a YAKE keyword extractor.""" - - -@dataclass(slots=True) -class PhraseCandidate: - """A phrase candidate with merged extraction-source metadata.""" - - phrase_text: str - phrase_norm: str - token_count: int - source_raw_ngram: bool = False - source_yake: bool = False - source_spacy_ner: bool = False - source_spacy_noun_chunk: bool = False - source_capitalized: bool = False - source_metadata: bool = False - spacy_label: str | None = None - raw_count: int = 0 - chapter_count: int = 0 - yake_score: float | None = None - candidate_score: float = 0.0 - sample_contexts: tuple[str, ...] = () - - -@dataclass(frozen=True, slots=True) -class LLMJudgment: - """A structured phrase judgment returned by the LLM judge.""" - - keep: bool - canonical: str | None - category: str | None - aliases: tuple[str, ...] - confidence: float - importance: float = 0.5 - allow_nested: bool = False - suppress_children: bool = True - reason: str | None = None - - -@dataclass(frozen=True, slots=True) -class PhraseLookup: - """In-memory lookup maps used for constant-time phrase-window checks.""" - - norm_to_phrase_ids: Mapping[str, tuple[int, ...]] - alias_to_phrase_ids: Mapping[str, tuple[int, ...]] - min_tokens: int - max_tokens: int - - -@dataclass(frozen=True, slots=True) -class PhraseMatch: - """An unhydrated query or chunk phrase match.""" - - phrase_norm: str - start_token: int - end_token: int - token_count: int - phrase_id: int | None = None - start_char: int | None = None - end_char: int | None = None - - -@dataclass(frozen=True, slots=True) -class HydratedPhraseMatch: - """A phrase match with protected-phrase metadata attached.""" - - phrase_id: int - matched_norm: str - phrase_text: str - phrase_norm: str - canonical_id: str - phrase_type: str | None - token_count: int - confidence: float - importance: float - allow_nested: bool - suppress_children: bool - start_token: int - end_token: int - start_char: int | None = None - end_char: int | None = None - book_id: int | None = None - series_id: int | None = None - - -@dataclass(frozen=True, slots=True) -class ProtectedPhraseBackfillResult: - """Summary of protected phrase generation for existing books.""" - - books_seen: int - books_built: int - protected_phrases: int - phrase_mentions: int - - -@dataclass(frozen=True, slots=True) -class PhraseCandidateGenerationResult: - """Summary of candidate phrase extraction for indexed books.""" - - books_seen: int - books_built: int - candidate_phrases: int - - -@dataclass(frozen=True, slots=True) -class PhraseJudgmentBackfillResult: - """Summary of LLM judging for stored candidate phrases.""" - - books_seen: int - books_judged: int - books_failed: int - candidates_judged: int - protected_phrases: int - phrase_mentions: int - - -@dataclass(frozen=True, slots=True) -class PhraseRecalculationResult: - """Summary of phrase cleanup and candidate regeneration for one book.""" - - book_id: int - deleted_candidates: int - deleted_protected_phrases: int - deleted_aliases: int - deleted_mentions: int - candidate_phrases: int - - -def strip_leading_articles(phrase_norm: str) -> str: - """Remove one leading English article from a normalized phrase.""" - tokens_ = phrase_norm.split() - if tokens_ and tokens_[0] in {"the", "a", "an"}: - tokens_ = tokens_[1:] - return " ".join(tokens_) - - -def normalize_candidate_phrase( - phrase_text: str, - config: EbookSearchConfig, - *, - min_tokens: int | None = None, - max_tokens: int | None = None, - strip_leading_article: bool = False, -) -> tuple[str, str, int] | None: - """Normalize a candidate phrase and validate token bounds.""" - normalized_tokens = tokenize_with_offsets(phrase_text) - start = 0 - if strip_leading_article and normalized_tokens and normalized_tokens[0].text in {"the", "a", "an"}: - start = 1 - - selected_tokens = normalized_tokens[start:] - min_count = config.phrase_min_tokens if min_tokens is None else min_tokens - max_count = config.phrase_max_tokens if max_tokens is None else max_tokens - if len(selected_tokens) < min_count or len(selected_tokens) > max_count: - return None - - phrase_norm = " ".join(token.text for token in selected_tokens) - if phrase_norm in get_ignored_phrases(): - return None - - display_text = phrase_text[selected_tokens[0].start_char : selected_tokens[-1].end_char].strip() - return display_text or phrase_norm, phrase_norm, len(selected_tokens) - - -def extract_raw_ngrams(text: str, config: EbookSearchConfig) -> dict[str, PhraseCandidate]: - """Extract raw normalized n-grams as high-recall candidates.""" - tokens_ = tokenize(text) - out: dict[str, PhraseCandidate] = {} - for ngram_size in range(config.phrase_min_tokens, config.phrase_max_tokens + 1): - for start in range(len(tokens_) - ngram_size + 1): - normalized = normalize_candidate_phrase(" ".join(tokens_[start : start + ngram_size]), config) - if normalized is None: - continue - phrase_text, phrase_norm, token_count = normalized - item = out.setdefault( - phrase_norm, - PhraseCandidate( - phrase_text=phrase_text, - phrase_norm=phrase_norm, - token_count=token_count, - source_raw_ngram=True, - ), - ) - item.raw_count += 1 - return out - - -def extract_yake_candidates( - book_text: str, - config: EbookSearchConfig, - top_k: int = 1000, -) -> dict[str, PhraseCandidate]: - """Extract YAKE keyphrases when the optional YAKE package is installed.""" - yake_module = optional_module("yake") - if yake_module is None: - logger.warning("ebook_phrase_yake_unavailable_skipping") - return {} - - keyword_extractor = yake_module.KeywordExtractor - extractor = keyword_extractor(lan="en", n=config.phrase_max_tokens, dedupLim=0.85, top=top_k) - out: dict[str, PhraseCandidate] = {} - for phrase_text, yake_score in extractor.extract_keywords(book_text): - normalized = normalize_candidate_phrase(phrase_text, config) - if normalized is None: - continue - display_text, phrase_norm, token_count = normalized - out[phrase_norm] = PhraseCandidate( - phrase_text=display_text, - phrase_norm=phrase_norm, - token_count=token_count, - source_yake=True, - yake_score=float(yake_score), - ) - return out - - -def optional_module(module_name: str) -> ModuleType | None: - """Import an optional module and return None when it is not installed.""" - try: - return importlib.import_module(module_name) - except ImportError: - return None - - -def extract_spacy_candidates( - book_text: str, - nlp: SpacyLanguage, - config: EbookSearchConfig, -) -> dict[str, PhraseCandidate]: - """Extract spaCy named entities and noun chunks from one text block.""" - out: dict[str, PhraseCandidate] = {} - doc = nlp(book_text) - - for ent in doc.ents: - normalized = normalize_candidate_phrase( - ent.text, - config, - max_tokens=config.phrase_max_entity_tokens, - ) - if normalized is None: - continue - phrase_text, phrase_norm, token_count = normalized - out[phrase_norm] = PhraseCandidate( - phrase_text=phrase_text, - phrase_norm=phrase_norm, - token_count=token_count, - source_spacy_ner=True, - spacy_label=ent.label_, - ) - - for chunk in doc.noun_chunks: - normalized = normalize_candidate_phrase(chunk.text, config, strip_leading_article=True) - if normalized is None: - continue - phrase_text, phrase_norm, token_count = normalized - out[phrase_norm] = PhraseCandidate( - phrase_text=phrase_text, - phrase_norm=phrase_norm, - token_count=token_count, - source_spacy_noun_chunk=True, - ) - return out - - -def extract_capitalized_phrases(original_text: str, config: EbookSearchConfig) -> dict[str, PhraseCandidate]: - """Extract capitalized phrase runs that often carry fictional terms.""" - out: dict[str, PhraseCandidate] = {} - for match in CAPITALIZED_PHRASE_RE.finditer(original_text): - phrase_text = match.group(0).strip() - normalized = normalize_candidate_phrase( - phrase_text, - config, - max_tokens=config.phrase_max_entity_tokens, - ) - if normalized is None: - continue - display_text, phrase_norm, token_count = normalized - out[phrase_norm] = PhraseCandidate( - phrase_text=display_text, - phrase_norm=phrase_norm, - token_count=token_count, - source_capitalized=True, - ) - return out - - -def extract_metadata_candidates( - metadata: Mapping[str, object] | None, - config: EbookSearchConfig, -) -> dict[str, PhraseCandidate]: - """Extract phrases from book metadata values such as title, author, and series.""" - if metadata is None: - return {} - - out: dict[str, PhraseCandidate] = {} - for value in metadata.values(): - if value is None: - continue - phrase_text = str(value).strip() - normalized = normalize_candidate_phrase( - phrase_text, - config, - max_tokens=config.phrase_max_entity_tokens, - ) - if normalized is None: - continue - display_text, phrase_norm, token_count = normalized - out[phrase_norm] = PhraseCandidate( - phrase_text=display_text, - phrase_norm=phrase_norm, - token_count=token_count, - source_metadata=True, - ) - return out - - -def merge_candidate_sources(*sources: Mapping[str, PhraseCandidate]) -> dict[str, PhraseCandidate]: - """Merge candidate dictionaries by normalized phrase.""" - merged: dict[str, PhraseCandidate] = {} - for source in sources: - for phrase_norm, item in source.items(): - existing = merged.setdefault( - phrase_norm, - PhraseCandidate( - phrase_text=item.phrase_text, - phrase_norm=phrase_norm, - token_count=item.token_count, - ), - ) - merge_candidate(existing, item) - return merged - - -def merge_candidate(existing: PhraseCandidate, item: PhraseCandidate) -> None: - """Merge one candidate into an existing candidate object.""" - existing.source_raw_ngram = existing.source_raw_ngram or item.source_raw_ngram - existing.source_yake = existing.source_yake or item.source_yake - existing.source_spacy_ner = existing.source_spacy_ner or item.source_spacy_ner - existing.source_spacy_noun_chunk = existing.source_spacy_noun_chunk or item.source_spacy_noun_chunk - existing.source_capitalized = existing.source_capitalized or item.source_capitalized - existing.source_metadata = existing.source_metadata or item.source_metadata - existing.raw_count += item.raw_count - if item.yake_score is not None: - existing.yake_score = item.yake_score - if item.spacy_label: - existing.spacy_label = item.spacy_label - - -def enrich_with_frequency_and_chapter_counts( - candidates: Mapping[str, PhraseCandidate], - chapters: Sequence[str], -) -> dict[str, PhraseCandidate]: - """Add raw occurrence and chapter-spread counts to candidates.""" - if not candidates: - return {} - - candidate_sets_by_size: dict[int, set[str]] = defaultdict(set) - for phrase_norm, candidate in candidates.items(): - candidate_sets_by_size[candidate.token_count].add(phrase_norm) - - total_counts: defaultdict[str, int] = defaultdict(int) - chapter_counts: defaultdict[str, int] = defaultdict(int) - for chapter in chapters: - seen_in_chapter: set[str] = set() - chapter_tokens = tokenize(chapter) - for ngram_size, candidate_norms in candidate_sets_by_size.items(): - for start in range(len(chapter_tokens) - ngram_size + 1): - phrase_norm = " ".join(chapter_tokens[start : start + ngram_size]) - if phrase_norm not in candidate_norms: - continue - total_counts[phrase_norm] += 1 - seen_in_chapter.add(phrase_norm) - for phrase_norm in seen_in_chapter: - chapter_counts[phrase_norm] += 1 - - enriched = dict(candidates) - for phrase_norm, candidate in enriched.items(): - candidate.raw_count = max(candidate.raw_count, total_counts[phrase_norm]) - candidate.chapter_count = chapter_counts[phrase_norm] - return enriched - - -def filter_storable_candidates( - candidates: Mapping[str, PhraseCandidate], - config: EbookSearchConfig, -) -> tuple[dict[str, PhraseCandidate], int, int, int]: - """Remove candidates that should not be persisted.""" - min_raw_count = minimum_candidate_raw_count(config) - filtered: dict[str, PhraseCandidate] = {} - too_short = 0 - too_rare = 0 - too_common = 0 - for phrase_norm, candidate in candidates.items(): - if candidate.token_count < config.phrase_min_tokens: - too_short += 1 - continue - if candidate.raw_count < min_raw_count: - too_rare += 1 - continue - if is_most_common_word_phrase(phrase_norm): - too_common += 1 - continue - filtered[phrase_norm] = candidate - return filtered, too_short, too_rare, too_common - - -def minimum_candidate_raw_count(config: EbookSearchConfig) -> int: - """Return the minimum occurrence count required before storing a candidate.""" - return max(config.phrase_raw_ngram_min_count, 1) - - -def is_most_common_word_phrase(phrase_norm: str) -> bool: - """Return whether every token in a normalized phrase is a common word.""" - tokens_ = phrase_norm.split() - common_words = get_most_common_words() - return bool(tokens_) and all(token in common_words for token in tokens_) - - -def score_candidate(candidate: PhraseCandidate, config: EbookSearchConfig) -> float: - """Score a phrase candidate before LLM judging.""" - score = source_score(candidate) + frequency_score(candidate, config) + token_count_score(candidate, config) - if candidate.phrase_norm in get_ignored_phrases(): - score -= 100.0 - if has_bad_start(candidate.phrase_norm): - score -= BAD_START_SCORE_PENALTY - if has_bad_end(candidate.phrase_norm): - score -= BAD_END_SCORE_PENALTY - return score - - -def has_bad_start(phrase_norm: str) -> bool: - """Return whether a normalized phrase starts with a bad starting token.""" - tokens_ = phrase_norm.split() - return bool(tokens_ and tokens_[0] in get_bad_starts()) - - -def has_bad_end(phrase_norm: str) -> bool: - """Return whether a normalized phrase ends with a bad ending token.""" - tokens_ = phrase_norm.split() - return bool(tokens_ and tokens_[-1] in get_bad_ends()) - - -def source_score(candidate: PhraseCandidate) -> float: - """Return the score contribution from extraction sources.""" - return sum( - weight - for enabled, weight in ( - (candidate.source_yake, 2.0), - (candidate.source_spacy_ner, 2.5), - (candidate.source_spacy_noun_chunk, 1.5), - (candidate.source_capitalized, 2.0), - (candidate.source_raw_ngram, 0.5), - ) - if enabled - ) - - -def frequency_score(candidate: PhraseCandidate, config: EbookSearchConfig) -> float: - """Return the score contribution from frequency and chapter spread.""" - return sum( - weight - for count, threshold, weight in ( - (candidate.raw_count, config.phrase_raw_count_score_threshold, 1.0), - (candidate.raw_count, config.phrase_raw_count_high_score_threshold, 1.0), - (candidate.chapter_count, config.phrase_chapter_count_score_threshold, 1.0), - (candidate.chapter_count, config.phrase_chapter_count_high_score_threshold, 1.0), - ) - if count >= threshold - ) - - -def token_count_score(candidate: PhraseCandidate, config: EbookSearchConfig) -> float: - """Return the score contribution from phrase length.""" - if candidate.token_count == 1: - return -0.5 - if candidate.token_count in {2, 3, 4}: - return 0.5 - if candidate.token_count > config.phrase_max_tokens: - return -1.0 - return 0.0 - - -def get_sample_contexts(book_text: str, phrase_norm: str, max_contexts: int = 5) -> list[str]: - """Return normalized context snippets containing a candidate phrase.""" - norm_text = normalize_text(book_text) - contexts: list[str] = [] - start = 0 - while len(contexts) < max_contexts: - index = norm_text.find(phrase_norm, start) - if index == -1: - break - left = max(0, index - 300) - right = min(len(norm_text), index + len(phrase_norm) + 300) - contexts.append(norm_text[left:right]) - start = index + len(phrase_norm) - return contexts - - -def judge_candidate_with_llm(candidate: PhraseCandidate, config: EbookSearchConfig) -> LLMJudgment: - """Ask the configured chat model to judge one pre-extracted candidate.""" - payload = { - "phrase": candidate.phrase_norm, - "token_count": candidate.token_count, - "sources": candidate_source_names(candidate), - "raw_count": candidate.raw_count, - "chapter_count": candidate.chapter_count, - "contexts": list(candidate.sample_contexts), - } - messages = [ - { - "role": "system", - "content": ( - "Judge whether a candidate phrase from a book should be protected for RAG retrieval. " - "Do not extract new phrases. Reject common grammar fragments, ordinary nonspecific phrases, " - "unstable fragments, and phrases kept only because they are frequent. Keep people, places, " - "organizations, factions, events, technologies, fictional conditions, magic systems, formal titles, " - "named concepts, and recurring world-specific terms. Return only a JSON object with keys: keep, " - "canonical, category, aliases, confidence, importance, allow_nested, suppress_children, reason." - ), - }, - {"role": "user", "content": json.dumps(payload, ensure_ascii=True)}, - ] - return parse_llm_judgment(request_chat_completion(config, messages), config) - - -def candidate_source_names(candidate: PhraseCandidate) -> list[str]: - """Return enabled source names for an extracted candidate.""" - names: list[str] = [] - if candidate.source_raw_ngram: - names.append("raw_ngram") - if candidate.source_yake: - names.append("yake") - if candidate.source_spacy_ner: - names.append("spacy_ner") - if candidate.source_spacy_noun_chunk: - names.append("spacy_noun_chunk") - if candidate.source_capitalized: - names.append("capitalized") - if candidate.source_metadata: - names.append("metadata") - return names - - -def parse_llm_judgment(content: str, config: EbookSearchConfig) -> LLMJudgment: - """Parse and validate an LLM phrase-judge response.""" - body = json.loads(extract_json_object(content)) - if not isinstance(body, dict): - msg = "LLM phrase judge response is not a JSON object" - raise TypeError(msg) - - aliases = body.get("aliases", ()) - if not isinstance(aliases, list | tuple): - aliases = () - return LLMJudgment( - keep=bool(body.get("keep", False)), - canonical=optional_text(body.get("canonical")), - category=optional_text(body.get("category")), - aliases=tuple(str(alias) for alias in aliases if isinstance(alias, str) and alias.strip()), - confidence=clamped_float(body.get("confidence"), default=0.0), - importance=clamped_float(body.get("importance"), default=0.5), - allow_nested=bool(body.get("allow_nested", config.phrase_default_allow_nested)), - suppress_children=bool(body.get("suppress_children", config.phrase_default_suppress_children)), - reason=optional_text(body.get("reason")), - ) - - -def extract_json_object(content: str) -> str: - """Extract a JSON object from plain or fenced model output.""" - stripped = content.strip() - if stripped.startswith("{") and stripped.endswith("}"): - return stripped - match = JSON_OBJECT_RE.search(stripped) - if match is None: - msg = "LLM phrase judge response did not contain a JSON object" - raise ValueError(msg) - return match.group(0) - - -def optional_text(value: object) -> str | None: - """Return stripped text for a nullable JSON value.""" - if not isinstance(value, str): - return None - stripped = value.strip() - return stripped or None - - -def clamped_float(value: object, *, default: float) -> float: - """Coerce a JSON number into the 0.0 to 1.0 range.""" - if not isinstance(value, int | float): - return default - return min(max(float(value), 0.0), 1.0) - - -def extract_phrase_candidates_for_book( - book_text: str, - chapters: Sequence[str], - config: EbookSearchConfig, - *, - nlp: SpacyLanguage | None = None, - metadata: Mapping[str, object] | None = None, -) -> list[PhraseCandidate]: - """Extract, score, and limit phrase candidates for one book.""" - started_at = perf_counter() - logger.info( - "ebook_phrase_candidate_extract_start chapters=%s chars=%s min_tokens=%s max_tokens=%s max_candidates=%s", - len(chapters), - len(book_text), - config.phrase_min_tokens, - config.phrase_max_tokens, - config.protected_phrase_max_candidates_per_book, - ) - raw_started_at = perf_counter() - raw = extract_raw_ngrams(book_text, config) - logger.info( - "ebook_phrase_candidate_extract_raw_complete candidates=%s duration_ms=%.1f", - len(raw), - (perf_counter() - raw_started_at) * 1000, - ) - yake_started_at = perf_counter() - yake_candidates = extract_yake_candidates(book_text, config) - logger.info( - "ebook_phrase_candidate_extract_yake_complete candidates=%s duration_ms=%.1f", - len(yake_candidates), - (perf_counter() - yake_started_at) * 1000, - ) - spacy_candidates: dict[str, PhraseCandidate] = {} - if nlp is not None: - spacy_started_at = perf_counter() - for chapter in chapters: - spacy_candidates = merge_candidate_sources(spacy_candidates, extract_spacy_candidates(chapter, nlp, config)) - logger.info( - "ebook_phrase_candidate_extract_spacy_complete candidates=%s duration_ms=%.1f", - len(spacy_candidates), - (perf_counter() - spacy_started_at) * 1000, - ) - capitalized_started_at = perf_counter() - capitalized = extract_capitalized_phrases(book_text, config) - logger.info( - "ebook_phrase_candidate_extract_capitalized_complete candidates=%s duration_ms=%.1f", - len(capitalized), - (perf_counter() - capitalized_started_at) * 1000, - ) - metadata_candidates = extract_metadata_candidates(metadata, config) - - candidates = merge_candidate_sources(raw, yake_candidates, spacy_candidates, capitalized, metadata_candidates) - enriched_started_at = perf_counter() - candidates = enrich_with_frequency_and_chapter_counts(candidates, chapters) - pre_filter_count = len(candidates) - candidates, filtered_too_short, filtered_too_rare, filtered_too_common = filter_storable_candidates( - candidates, - config, - ) - for candidate in candidates.values(): - candidate.candidate_score = score_candidate(candidate, config) - - limited = sorted(candidates.values(), key=lambda item: item.candidate_score, reverse=True)[ - : config.protected_phrase_max_candidates_per_book - ] - logger.info( - "ebook_phrase_candidate_extract_complete raw=%s yake=%s spacy=%s capitalized=%s metadata=%s " - "merged=%s filtered_too_short=%s filtered_too_rare=%s filtered_too_common=%s min_uses=%s " - "storable=%s limited=%s enrich_score_ms=%.1f duration_ms=%.1f", - len(raw), - len(yake_candidates), - len(spacy_candidates), - len(capitalized), - len(metadata_candidates), - pre_filter_count, - filtered_too_short, - filtered_too_rare, - filtered_too_common, - minimum_candidate_raw_count(config), - len(candidates), - len(limited), - (perf_counter() - enriched_started_at) * 1000, - (perf_counter() - started_at) * 1000, - ) - return limited - - -def generate_candidate_phrases_for_book( - session: Session, - book_id: int, - series_id: int | None, - book_text: str, - chapters: Sequence[str], - config: EbookSearchConfig, - *, - nlp: SpacyLanguage | None = None, - metadata: Mapping[str, object] | None = None, -) -> list[EbookCandidatePhrase]: - """Extract and store candidate phrases for one book without LLM judging.""" - started_at = perf_counter() - limited_candidates = extract_phrase_candidates_for_book( - book_text, - chapters, - config, - nlp=nlp, - metadata=metadata, - ) - save_started_at = perf_counter() - pruned_count = prune_unstorable_unjudged_candidate_phrases(session, book_id, config) - logger.info( - "ebook_candidate_phrase_save_start book_id=%s candidates=%s pruned_unstorable=%s", - book_id, - len(limited_candidates), - pruned_count, - ) - rows = [ - save_candidate_to_db(session, book_id, series_id, candidate, judgment=None) for candidate in limited_candidates - ] - session.flush() - logger.info( - "ebook_candidate_phrase_generation_complete book_id=%s candidates=%s save_ms=%.1f duration_ms=%.1f", - book_id, - len(rows), - (perf_counter() - save_started_at) * 1000, - (perf_counter() - started_at) * 1000, - ) - return rows - - -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.""" - 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), - ) - deleted += prune_common_word_candidate_phrases(session, book_id) - return deleted - - -def prune_common_word_candidate_phrases(session: Session, book_id: int) -> int: - """Delete old unjudged candidates made only from common words.""" - rows = list( - session.scalars( - select(EbookCandidatePhrase).where( - EbookCandidatePhrase.book_id == book_id, - EbookCandidatePhrase.llm_judged.is_(False), - ) - ) - ) - ids = [row.id for row in rows if row.id is not None and is_most_common_word_phrase(row.phrase_norm)] - if not ids: - return 0 - deleted = rowcount(session.execute(delete(EbookCandidatePhrase).where(EbookCandidatePhrase.id.in_(ids)))) - logger.info("ebook_candidate_phrase_common_word_pruned book_id=%s deleted=%s", book_id, deleted) - return deleted - - -def build_protected_phrases_for_book( - session: Session, - book_id: int, - series_id: int | None, - book_text: str, - chapters: Sequence[str], - config: EbookSearchConfig, - *, - nlp: SpacyLanguage | None = None, - metadata: Mapping[str, object] | None = None, -) -> list[EbookProtectedPhrase]: - """Extract candidates, judge stored candidates, and store accepted protected phrases.""" - candidates = generate_candidate_phrases_for_book( - session, - book_id, - series_id, - book_text, - chapters, - config, - nlp=nlp, - metadata=metadata, - ) - judged_count, protected = judge_candidate_phrases_for_book( - session, - book_id, - series_id, - book_text, - config, - ) - - logger.info( - "ebook_protected_phrase_build_complete book_id=%s candidates=%s judged=%s protected=%s", - book_id, - len(candidates), - judged_count, - len(protected), - ) - return protected - - -def generate_candidate_phrases_for_books( - session: Session, - config: EbookSearchConfig, - *, - limit: int | None = None, - nlp: SpacyLanguage | None = None, -) -> PhraseCandidateGenerationResult: - """Create or refresh candidate phrases for indexed books without calling the LLM judge.""" - sources = list(session.scalars(select(EbookSource).order_by(EbookSource.id))) - books_seen = len(sources) - books_built = 0 - candidate_count = 0 - logger.info( - "ebook_candidate_phrase_generation_start books_seen=%s limit=%s min_tokens=%s max_tokens=%s " - "max_candidates_per_book=%s", - books_seen, - limit, - config.phrase_min_tokens, - config.phrase_max_tokens, - config.protected_phrase_max_candidates_per_book, - ) - - for book_number, source in enumerate(sources, start=1): - if limit is not None and books_built >= limit: - logger.info("ebook_candidate_phrase_generation_book_limit_reached limit=%s", limit) - break - book_started_at = perf_counter() - logger.info( - "ebook_candidate_phrase_generation_book_start source_id=%s book_number=%s books_seen=%s title=%r", - source.id, - book_number, - books_seen, - source.title, - ) - - try: - chapters = load_book_chapter_texts(session, source.id) - if not chapters: - logger.warning("ebook_candidate_phrase_generation_book_empty source_id=%s", source.id) - continue - book_text = "\n\n".join(chapters) - logger.info( - "ebook_candidate_phrase_generation_book_loaded source_id=%s chapters=%s chars=%s", - source.id, - len(chapters), - len(book_text), - ) - - candidates = generate_candidate_phrases_for_book( - session, - source.id, - series_id=None, - book_text=book_text, - chapters=chapters, - config=config, - nlp=nlp, - metadata=metadata_for_source(source), - ) - session.commit() - except Exception: - session.rollback() - logger.exception( - "ebook_candidate_phrase_generation_book_failed source_id=%s book_number=%s books_seen=%s", - source.id, - book_number, - books_seen, - ) - raise - books_built += 1 - candidate_count += len(candidates) - logger.info( - "ebook_candidate_phrase_generation_book_committed source_id=%s candidates=%s books_built=%s " - "candidate_total=%s duration_ms=%.1f", - source.id, - len(candidates), - books_built, - candidate_count, - (perf_counter() - book_started_at) * 1000, - ) - - logger.info( - "ebook_candidate_phrase_generation_complete books_seen=%s books_built=%s candidate_total=%s", - books_seen, - books_built, - candidate_count, - ) - return PhraseCandidateGenerationResult( - books_seen=books_seen, - books_built=books_built, - candidate_phrases=candidate_count, - ) - - -def recalculate_candidate_phrases_for_book( - session: Session, - source: EbookSource, - config: EbookSearchConfig, - *, - nlp: SpacyLanguage | None = None, -) -> PhraseRecalculationResult: - """Remove all book phrase data, regenerate candidates, and commit the completed book.""" - started_at = perf_counter() - logger.info( - "ebook_candidate_phrase_recalculation_start source_id=%s title=%r", - source.id, - source.title, - ) - try: - deleted = delete_phrase_data_for_book(session, source.id) - chapters = load_book_chapter_texts(session, source.id) - if not chapters: - logger.warning("ebook_candidate_phrase_recalculation_book_empty source_id=%s", source.id) - session.commit() - return PhraseRecalculationResult( - book_id=source.id, - deleted_candidates=deleted.deleted_candidates, - deleted_protected_phrases=deleted.deleted_protected_phrases, - deleted_aliases=deleted.deleted_aliases, - deleted_mentions=deleted.deleted_mentions, - candidate_phrases=0, - ) - - candidates = generate_candidate_phrases_for_book( - session, - source.id, - series_id=None, - book_text="\n\n".join(chapters), - chapters=chapters, - config=config, - nlp=nlp, - metadata=metadata_for_source(source), - ) - session.commit() - except Exception: - session.rollback() - logger.exception("ebook_candidate_phrase_recalculation_failed source_id=%s", source.id) - raise - - result = PhraseRecalculationResult( - book_id=source.id, - deleted_candidates=deleted.deleted_candidates, - deleted_protected_phrases=deleted.deleted_protected_phrases, - deleted_aliases=deleted.deleted_aliases, - deleted_mentions=deleted.deleted_mentions, - candidate_phrases=len(candidates), - ) - logger.info( - "ebook_candidate_phrase_recalculation_complete source_id=%s deleted_candidates=%s " - "deleted_protected=%s deleted_aliases=%s deleted_mentions=%s candidates=%s duration_ms=%.1f", - source.id, - result.deleted_candidates, - result.deleted_protected_phrases, - result.deleted_aliases, - result.deleted_mentions, - result.candidate_phrases, - (perf_counter() - started_at) * 1000, - ) - return result - - -def delete_phrase_data_for_book(session: Session, book_id: int) -> PhraseRecalculationResult: - """Delete all candidate, protected, alias, and mention phrase data for one book.""" - protected_ids = list( - session.scalars(select(EbookProtectedPhrase.id).where(EbookProtectedPhrase.book_id == book_id)) - ) - 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.""" - count = getattr(result, "rowcount", 0) - return int(count if count is not None and count >= 0 else 0) - - -def judge_candidate_phrases_for_books( - session: Session, - config: EbookSearchConfig, - *, - limit: int | None = None, -) -> PhraseJudgmentBackfillResult: - """Judge stored candidate phrases and promote accepted phrases for indexed books.""" - sources = list(session.scalars(select(EbookSource).order_by(EbookSource.id))) - books_seen = len(sources) - books_judged = 0 - books_failed = 0 - candidates_judged = 0 - protected_count = 0 - mention_count = 0 - logger.info( - "ebook_candidate_phrase_judgment_start books_seen=%s limit=%s llm_candidates_per_book=%s " - "target_protected_per_book=%s confidence_threshold=%.2f min_tokens=%s min_uses=%s", - books_seen, - limit, - config.protected_phrase_llm_candidates_per_book, - config.phrase_target_protected_per_book, - config.protected_phrase_confidence_threshold, - config.phrase_min_tokens, - minimum_candidate_raw_count(config), - ) - - for book_number, source in enumerate(sources, start=1): - if limit is not None and books_judged >= limit: - logger.info("ebook_candidate_phrase_judgment_book_limit_reached limit=%s", limit) - break - book_started_at = perf_counter() - pruned_count = prune_common_word_candidate_phrases(session, source.id) - if pruned_count: - session.flush() - - unjudged_count = session.scalar( - select(func.count(EbookCandidatePhrase.id)).where( - EbookCandidatePhrase.book_id == source.id, - EbookCandidatePhrase.llm_judged.is_(False), - EbookCandidatePhrase.token_count >= config.phrase_min_tokens, - EbookCandidatePhrase.raw_count >= minimum_candidate_raw_count(config), - ) - ) - if not unjudged_count: - logger.info( - "ebook_candidate_phrase_judgment_book_skip_no_unjudged source_id=%s book_number=%s books_seen=%s", - source.id, - book_number, - books_seen, - ) - continue - logger.info( - "ebook_candidate_phrase_judgment_book_start source_id=%s book_number=%s books_seen=%s title=%r unjudged=%s", - source.id, - book_number, - books_seen, - source.title, - unjudged_count, - ) - - try: - chapters = load_book_chapter_texts(session, source.id) - if not chapters: - logger.warning("ebook_candidate_phrase_judgment_book_empty source_id=%s", source.id) - continue - - judged, protected = judge_candidate_phrases_for_book( - session, - source.id, - series_id=None, - book_text="\n\n".join(chapters), - config=config, - ) - if judged == 0: - logger.info( - "ebook_candidate_phrase_judgment_book_skip_no_judgments source_id=%s unjudged=%s", - source.id, - unjudged_count, - ) - continue - - mentions = index_chunk_phrase_mentions_for_book(session, source.id, config) if protected else 0 - session.commit() - except Exception: - session.rollback() - books_failed += 1 - logger.exception( - "ebook_candidate_phrase_judgment_book_failed source_id=%s book_number=%s books_seen=%s books_failed=%s", - source.id, - book_number, - books_seen, - books_failed, - ) - continue - books_judged += 1 - candidates_judged += judged - protected_count += len(protected) - mention_count += mentions - logger.info( - "ebook_candidate_phrase_judgment_book_committed source_id=%s judged=%s protected=%s mentions=%s " - "books_judged=%s candidates_judged_total=%s protected_total=%s duration_ms=%.1f", - source.id, - judged, - len(protected), - mentions, - books_judged, - candidates_judged, - protected_count, - (perf_counter() - book_started_at) * 1000, - ) - - logger.info( - "ebook_candidate_phrase_judgment_complete books_seen=%s books_judged=%s books_failed=%s " - "candidates_judged=%s protected=%s mentions=%s", - books_seen, - books_judged, - books_failed, - candidates_judged, - protected_count, - mention_count, - ) - return PhraseJudgmentBackfillResult( - books_seen=books_seen, - books_judged=books_judged, - books_failed=books_failed, - candidates_judged=candidates_judged, - protected_phrases=protected_count, - phrase_mentions=mention_count, - ) - - -def judge_candidate_phrases_for_book( - session: Session, - book_id: int, - series_id: int | None, - book_text: str, - config: EbookSearchConfig, -) -> tuple[int, list[EbookProtectedPhrase]]: - """Judge unjudged candidate phrase rows for one book.""" - judgment_limit = config.protected_phrase_llm_candidates_per_book - if judgment_limit <= 0: - logger.info("ebook_candidate_phrase_judgment_skipped_llm_limit_zero book_id=%s", book_id) - return 0, [] - - existing_protected = session.scalar( - select(func.count(EbookProtectedPhrase.id)).where(EbookProtectedPhrase.book_id == book_id) - ) - target_remaining: int | None = None - if config.phrase_target_protected_per_book > 0: - target_remaining = max(config.phrase_target_protected_per_book - int(existing_protected or 0), 0) - if target_remaining == 0: - logger.info( - "ebook_candidate_phrase_judgment_skipped_target_met book_id=%s existing_protected=%s target=%s", - book_id, - existing_protected, - config.phrase_target_protected_per_book, - ) - return 0, [] - - rows = list( - 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) - ) - ) - rows = [row for row in rows if not is_most_common_word_phrase(row.phrase_norm)] - - logger.info( - "ebook_candidate_phrase_judgment_candidates_loaded book_id=%s candidates=%s existing_protected=%s " - "target_remaining=%s judgment_limit=%s", - book_id, - len(rows), - existing_protected, - target_remaining, - judgment_limit, - ) - judged_count = 0 - protected: list[EbookProtectedPhrase] = [] - for row_number, row in enumerate(rows, start=1): - row_started_at = perf_counter() - candidate = phrase_candidate_from_row(row) - candidate.sample_contexts = tuple(row.sample_contexts or get_sample_contexts(book_text, candidate.phrase_norm)) - logger.info( - "ebook_candidate_phrase_judgment_candidate_start book_id=%s candidate_id=%s row_number=%s rows=%s " - "phrase=%r score=%.3f raw_count=%s chapter_count=%s", - book_id, - row.id, - row_number, - len(rows), - candidate.phrase_norm, - candidate.candidate_score, - candidate.raw_count, - candidate.chapter_count, - ) - judgment = judge_candidate_with_llm(candidate, config) - candidate_row = save_candidate_to_db(session, book_id, series_id, candidate, judgment=judgment) - judged_count += 1 - logger.info( - "ebook_candidate_phrase_judgment_candidate_complete book_id=%s candidate_id=%s phrase=%r keep=%s " - "confidence=%.3f importance=%.3f category=%r duration_ms=%.1f", - book_id, - row.id, - candidate.phrase_norm, - judgment.keep, - judgment.confidence, - judgment.importance, - judgment.category, - (perf_counter() - row_started_at) * 1000, - ) - if not judgment.keep or judgment.confidence < config.protected_phrase_confidence_threshold: - continue - accepted_norm = normalize_text(judgment.canonical or candidate.phrase_text) - accepted_token_count = len(accepted_norm.split()) - if accepted_token_count < config.phrase_min_tokens: - logger.info( - "ebook_candidate_phrase_judgment_candidate_skip_short_canonical book_id=%s candidate_id=%s " - "phrase=%r canonical=%r token_count=%s min_tokens=%s", - book_id, - row.id, - candidate.phrase_norm, - accepted_norm, - accepted_token_count, - config.phrase_min_tokens, - ) - continue - if is_most_common_word_phrase(accepted_norm): - logger.info( - "ebook_candidate_phrase_judgment_candidate_skip_common_canonical book_id=%s candidate_id=%s " - "phrase=%r canonical=%r", - book_id, - row.id, - candidate.phrase_norm, - accepted_norm, - ) - continue - protected.append(upsert_protected_phrase(session, book_id, series_id, candidate, judgment, candidate_row)) - if target_remaining is not None and len(protected) >= target_remaining: - break - - session.flush() - return judged_count, protected - - -def phrase_candidate_from_row(row: EbookCandidatePhrase) -> PhraseCandidate: - """Recreate an in-memory candidate from a persisted candidate row.""" - 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=tuple(row.sample_contexts or ()), - ) - - -def metadata_for_source(source: EbookSource) -> dict[str, object | None]: - """Return phrase extraction metadata for one indexed source.""" - return { - "title": source.title, - "author": source.author, - "language": source.language, - "publisher": source.publisher, - "identifier": source.identifier, - } - - -def build_missing_protected_phrases( - session: Session, - config: EbookSearchConfig, - *, - limit: int | None = None, - nlp: SpacyLanguage | None = None, -) -> ProtectedPhraseBackfillResult: - """Create protected phrases and phrase mentions for indexed books that do not have them yet.""" - sources = list(session.scalars(select(EbookSource).order_by(EbookSource.id))) - books_seen = len(sources) - books_built = 0 - protected_count = 0 - mention_count = 0 - logger.info( - "ebook_protected_phrase_backfill_start books_seen=%s limit=%s", - books_seen, - limit, - ) - - for book_number, source in enumerate(sources, start=1): - if limit is not None and books_built >= limit: - logger.info("ebook_protected_phrase_backfill_book_limit_reached limit=%s", limit) - break - book_started_at = perf_counter() - existing_count = session.scalar( - select(func.count(EbookProtectedPhrase.id)).where(EbookProtectedPhrase.book_id == source.id) - ) - if existing_count: - logger.info( - "ebook_protected_phrase_backfill_book_skip_existing source_id=%s book_number=%s books_seen=%s " - "existing=%s", - source.id, - book_number, - books_seen, - existing_count, - ) - continue - logger.info( - "ebook_protected_phrase_backfill_book_start source_id=%s book_number=%s books_seen=%s title=%r", - source.id, - book_number, - books_seen, - source.title, - ) - - try: - chapters = load_book_chapter_texts(session, source.id) - if not chapters: - logger.warning("ebook_protected_phrase_backfill_book_empty source_id=%s", source.id) - continue - - protected = build_protected_phrases_for_book( - session, - source.id, - series_id=None, - book_text="\n\n".join(chapters), - chapters=chapters, - config=config, - nlp=nlp, - metadata=metadata_for_source(source), - ) - mentions = index_chunk_phrase_mentions_for_book(session, source.id, config) - session.commit() - except Exception: - session.rollback() - logger.exception( - "ebook_protected_phrase_backfill_book_failed source_id=%s book_number=%s books_seen=%s", - source.id, - book_number, - books_seen, - ) - raise - books_built += 1 - protected_count += len(protected) - mention_count += mentions - logger.info( - "ebook_protected_phrase_backfill_book_committed source_id=%s protected=%s mentions=%s books_built=%s " - "protected_total=%s mention_total=%s duration_ms=%.1f", - source.id, - len(protected), - mentions, - books_built, - protected_count, - mention_count, - (perf_counter() - book_started_at) * 1000, - ) - - logger.info( - "ebook_protected_phrase_backfill_complete books_seen=%s books_built=%s protected=%s mentions=%s", - books_seen, - books_built, - protected_count, - mention_count, - ) - return ProtectedPhraseBackfillResult( - books_seen=books_seen, - books_built=books_built, - protected_phrases=protected_count, - phrase_mentions=mention_count, - ) - - -def load_book_chapter_texts(session: Session, book_id: int) -> list[str]: - """Reconstruct chapter-like text blocks from indexed chunks for phrase extraction.""" - 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 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.""" - row = session.scalar( - select(EbookCandidatePhrase).where( - EbookCandidatePhrase.book_id == book_id, - EbookCandidatePhrase.phrase_norm == candidate.phrase_norm, - ) - ) - if row is None: - row = EbookCandidatePhrase(book_id=book_id, phrase_norm=candidate.phrase_norm) - row.llm_judged = False - session.add(row) - 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_spacy_ner = candidate.source_spacy_ner - row.source_spacy_noun_chunk = candidate.source_spacy_noun_chunk - row.source_capitalized = candidate.source_capitalized - row.source_metadata = candidate.source_metadata - row.spacy_label = candidate.spacy_label - 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) - if judgment is not None: - row.llm_judged = True - row.llm_keep = judgment.keep - row.llm_confidence = judgment.confidence - row.llm_category = judgment.category - row.llm_reason = judgment.reason - return row - - -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.""" - 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) - - row = session.scalar( - select(EbookProtectedPhrase).where( - EbookProtectedPhrase.book_id == book_id, - EbookProtectedPhrase.phrase_norm == phrase_norm, - ) - ) - if row is None: - row = EbookProtectedPhrase(book_id=book_id, phrase_norm=phrase_norm) - session.add(row) - row.series_id = series_id - row.phrase_text = phrase_text - row.canonical_id = make_canonical_id(judgment, phrase_norm) - row.phrase_type = judgment.category - row.token_count = len(phrase_norm.split()) - row.confidence = judgment.confidence - row.importance = judgment.importance - row.allow_nested = judgment.allow_nested - row.suppress_children = judgment.suppress_children - row.source_candidate_id = source_candidate.id - session.flush() - - 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.""" - alias_norm = normalize_text(alias_text) - if not alias_norm or alias_norm == phrase.phrase_norm: - return None - - row = session.scalar( - select(EbookPhraseAlias).where( - EbookPhraseAlias.phrase_id == phrase.id, - EbookPhraseAlias.alias_norm == alias_norm, - ) - ) - if row is None: - row = EbookPhraseAlias(phrase_id=phrase.id, alias_norm=alias_norm) - session.add(row) - row.alias_text = alias_text - row.confidence = 1.0 - return row - - -def make_canonical_id(judgment: LLMJudgment, phrase_norm: str) -> str: - """Create a deterministic canonical id from a judgment category and phrase.""" - 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.""" - slug = re.sub(r"[^a-z0-9]+", "_", normalize_text(value).replace("'", "")) - return slug.strip("_") or "unknown" - - -def load_phrase_lookup( - session: Session, - config: EbookSearchConfig, - *, - book_id: int | None = None, - series_id: int | None = None, -) -> PhraseLookup: - """Load protected phrases and aliases into RAM lookup maps.""" - norm_to_ids: defaultdict[str, list[int]] = defaultdict(list) - alias_to_ids: defaultdict[str, list[int]] = defaultdict(list) - max_tokens = config.phrase_max_tokens - - phrase_statement = select( - EbookProtectedPhrase.id, - EbookProtectedPhrase.phrase_norm, - EbookProtectedPhrase.token_count, - ) - scope_filter = protected_phrase_scope_filter(book_id=book_id, series_id=series_id) - if scope_filter is not None: - phrase_statement = phrase_statement.where(scope_filter) - - for row in session.execute(phrase_statement): - phrase_id = int(row.id) - phrase_norm = str(row.phrase_norm) - norm_to_ids[phrase_norm].append(phrase_id) - max_tokens = max(max_tokens, int(row.token_count)) - - alias_statement = select( - EbookPhraseAlias.alias_norm, - EbookPhraseAlias.phrase_id, - ).join(EbookProtectedPhrase, EbookProtectedPhrase.id == EbookPhraseAlias.phrase_id) - if scope_filter is not None: - alias_statement = alias_statement.where(scope_filter) - - for row in session.execute(alias_statement): - alias_norm = str(row.alias_norm) - alias_to_ids[alias_norm].append(int(row.phrase_id)) - max_tokens = max(max_tokens, len(alias_norm.split())) - - return PhraseLookup( - norm_to_phrase_ids={key: tuple(values) for key, values in norm_to_ids.items()}, - alias_to_phrase_ids={key: tuple(values) for key, values in alias_to_ids.items()}, - min_tokens=config.phrase_min_tokens, - max_tokens=max_tokens, - ) - - -def protected_phrase_scope_filter(*, book_id: int | None, series_id: int | None) -> object | None: - """Build a SQLAlchemy filter for optional phrase book and series scope.""" - conditions = [] - if book_id is not None: - conditions.append(or_(EbookProtectedPhrase.book_id.is_(None), EbookProtectedPhrase.book_id == book_id)) - if series_id is not None: - conditions.append(or_(EbookProtectedPhrase.series_id.is_(None), EbookProtectedPhrase.series_id == series_id)) - if not conditions: - return None - return and_(*conditions) - - -def generate_query_ngrams( - tokens_: Sequence[str], - min_n: int, - max_n: int, -) -> Iterator[tuple[str, int, int]]: - """Generate normalized query windows from longest to shortest.""" - capped_max_n = min(max_n, len(tokens_)) - for ngram_size in range(capped_max_n, min_n - 1, -1): - for start in range(len(tokens_) - ngram_size + 1): - end = start + ngram_size - phrase_norm = " ".join(tokens_[start:end]) - if phrase_norm in get_ignored_phrases(): - continue - yield phrase_norm, start, end - - -def detect_phrase_candidates(query_text: str, lookup: PhraseLookup) -> list[PhraseMatch]: - """Detect protected phrase windows in a user query using RAM hash lookups.""" - return detect_phrase_candidates_from_tokens(tokenize_with_offsets(query_text), lookup) - - -def detect_phrase_candidates_in_text(text: str, lookup: PhraseLookup) -> list[PhraseMatch]: - """Detect protected phrase windows in arbitrary text with character offsets.""" - return detect_phrase_candidates_from_tokens(tokenize_with_offsets(text), lookup) - - -def detect_phrase_candidates_from_tokens(tokens_: Sequence[NormalizedToken], lookup: PhraseLookup) -> list[PhraseMatch]: - """Detect protected phrase windows from already-normalized tokens.""" - matches: list[PhraseMatch] = [] - seen: set[tuple[int | None, str, int, int]] = set() - token_texts = [token.text for token in tokens_] - for phrase_norm, start, end in generate_query_ngrams(token_texts, min_n=lookup.min_tokens, max_n=lookup.max_tokens): - phrase_ids = lookup.norm_to_phrase_ids.get(phrase_norm, ()) - alias_ids = lookup.alias_to_phrase_ids.get(phrase_norm, ()) - for phrase_id in (*phrase_ids, *alias_ids): - key = (phrase_id, phrase_norm, start, end) - if key in seen: - continue - seen.add(key) - matches.append( - PhraseMatch( - phrase_norm=phrase_norm, - phrase_id=phrase_id, - start_token=start, - end_token=end, - token_count=end - start, - start_char=tokens_[start].start_char, - end_char=tokens_[end - 1].end_char, - ) - ) - return matches - - -def hydrate_matches(session: Session, matches: Sequence[PhraseMatch]) -> list[HydratedPhraseMatch]: - """Fetch protected phrase metadata for raw phrase matches.""" - if not matches: - return [] - - phrase_ids = sorted({match.phrase_id for match in matches if match.phrase_id is not None}) - if not phrase_ids: - return [] - - rows = { - row.id: row - for row in session.scalars(select(EbookProtectedPhrase).where(EbookProtectedPhrase.id.in_(phrase_ids))) - } - hydrated: list[HydratedPhraseMatch] = [] - for match in matches: - if match.phrase_id is None: - continue - phrase = rows.get(match.phrase_id) - if phrase is None: - continue - hydrated.append( - HydratedPhraseMatch( - phrase_id=phrase.id, - matched_norm=match.phrase_norm, - phrase_text=phrase.phrase_text, - phrase_norm=phrase.phrase_norm, - canonical_id=phrase.canonical_id, - phrase_type=phrase.phrase_type, - token_count=match.token_count, - confidence=phrase.confidence, - importance=phrase.importance, - allow_nested=phrase.allow_nested, - suppress_children=phrase.suppress_children, - start_token=match.start_token, - end_token=match.end_token, - start_char=match.start_char, - end_char=match.end_char, - book_id=phrase.book_id, - series_id=phrase.series_id, - ) - ) - return hydrated - - -def overlaps(first: HydratedPhraseMatch, second: HydratedPhraseMatch) -> bool: - """Return whether two token spans overlap.""" - return not (first.end_token <= second.start_token or first.start_token >= second.end_token) - - -def is_inside(child: HydratedPhraseMatch, parent: HydratedPhraseMatch) -> bool: - """Return whether one token span is strictly inside another.""" - return ( - child.start_token >= parent.start_token - and child.end_token <= parent.end_token - and (child.start_token, child.end_token, child.phrase_id) - != (parent.start_token, parent.end_token, parent.phrase_id) - ) - - -def rank_match(match: HydratedPhraseMatch) -> tuple[float, float, int]: - """Rank phrase matches by importance, confidence, then token count.""" - return (match.importance, match.confidence, match.token_count) - - -def should_suppress(candidate: HydratedPhraseMatch, kept: HydratedPhraseMatch) -> bool: - """Return whether an already-kept match should suppress a candidate.""" - if not overlaps(candidate, kept): - return False - if candidate.canonical_id == kept.canonical_id: - return rank_match(kept) >= rank_match(candidate) - if is_inside(candidate, kept) and kept.suppress_children and not candidate.allow_nested: - return True - return not candidate.allow_nested and rank_match(kept) > rank_match(candidate) - - -def resolve_overlaps(matches: Sequence[HydratedPhraseMatch]) -> list[HydratedPhraseMatch]: - """Resolve overlapping phrase matches without relying only on longest match.""" - sorted_matches = sorted( - matches, - key=lambda match: (match.start_token, -match.token_count, -match.importance, -match.confidence), - ) - kept: list[HydratedPhraseMatch] = [] - for candidate in sorted_matches: - if any(should_suppress(candidate, existing) for existing in kept): - continue - kept.append(candidate) - return kept - - -def detect_protected_phrases_for_query( - session: Session, - query_text: str, - config: EbookSearchConfig, - *, - lookup: PhraseLookup | None = None, - book_id: int | None = None, - series_id: int | None = None, -) -> list[HydratedPhraseMatch]: - """Run the full online protected-phrase query-detection pipeline.""" - active_lookup = ( - lookup if lookup is not None else load_phrase_lookup(session, config, book_id=book_id, series_id=series_id) - ) - return resolve_overlaps(hydrate_matches(session, detect_phrase_candidates(query_text, active_lookup))) - - -def index_chunk_phrase_mentions_for_book( - session: Session, - book_id: int, - config: EbookSearchConfig, - *, - series_id: int | None = None, - lookup: PhraseLookup | None = None, -) -> int: - """Rebuild chunk phrase mentions for all chunks in one book.""" - active_lookup = ( - lookup if lookup is not None else load_phrase_lookup(session, config, book_id=book_id, series_id=series_id) - ) - session.execute(delete(EbookChunkPhraseMention).where(EbookChunkPhraseMention.book_id == book_id)) - chunks = session.scalars(select(EbookChunk).where(EbookChunk.source_id == book_id).order_by(EbookChunk.id)) - count = 0 - for chunk in chunks: - count += index_chunk_phrase_mentions(session, chunk, lookup=active_lookup) - session.flush() - logger.info("ebook_chunk_phrase_mentions_indexed book_id=%s mentions=%s", book_id, count) - return count - - -def index_chunk_phrase_mentions(session: Session, chunk: EbookChunk, *, lookup: PhraseLookup) -> int: - """Store protected phrase mentions for one chunk.""" - raw_matches = detect_phrase_candidates_in_text(chunk.text, lookup) - hydrated = resolve_overlaps(hydrate_matches(session, raw_matches)) - for match in hydrated: - session.add( - EbookChunkPhraseMention( - chunk_id=chunk.id, - phrase_id=match.phrase_id, - book_id=match.book_id if match.book_id is not None else chunk.source_id, - series_id=match.series_id, - start_char=match.start_char if match.start_char is not None else 0, - end_char=match.end_char, - ) - ) - return len(hydrated) - - -def phrase_hit_counts_for_chunks( - session: Session, - *, - chunk_ids: Sequence[int], - phrase_ids: Sequence[int], -) -> dict[int, int]: - """Return phrase-hit counts by chunk id using indexed chunk mentions.""" - if not chunk_ids or not phrase_ids: - return {} - - statement = ( - select( - EbookChunkPhraseMention.chunk_id, - func.count(EbookChunkPhraseMention.phrase_id).label("phrase_hit_count"), - ) - .where( - EbookChunkPhraseMention.chunk_id.in_(chunk_ids), - EbookChunkPhraseMention.phrase_id.in_(phrase_ids), - ) - .group_by(EbookChunkPhraseMention.chunk_id) - ) - return {int(row.chunk_id): int(row.phrase_hit_count) for row in session.execute(statement)} diff --git a/python/ebook_search/protected_phrases/matching.py b/python/ebook_search/protected_phrases/matching.py new file mode 100644 index 0000000..7fecbe3 --- /dev/null +++ b/python/ebook_search/protected_phrases/matching.py @@ -0,0 +1,487 @@ +"""Runtime protected-phrase matching and chunk mention indexing.""" + +from __future__ import annotations + +import logging +from collections import defaultdict +from typing import TYPE_CHECKING + +from sqlalchemy import and_, delete, func, or_, select + +from python.ebook_search.protected_phrases.config import get_ignored_phrases +from python.ebook_search.protected_phrases.models import ( + ChunkPhraseHit, + HydratedPhraseMatch, + PhraseLookup, + PhraseMatch, +) +from python.ebook_search.protected_phrases.text_normalization import tokenize_with_offsets +from python.orm.richie import ( + EbookChunk, + EbookChunkPhraseMention, + EbookPhraseAlias, + EbookProtectedPhrase, +) + +if TYPE_CHECKING: + from collections.abc import Iterator, Sequence + + from sqlalchemy.orm import Session + + from python.ebook_search.config import EbookSearchConfig + from python.ebook_search.protected_phrases.text_normalization import NormalizedToken + +logger = logging.getLogger(__name__) + + +def load_phrase_lookup( + session: Session, + config: EbookSearchConfig, + *, + book_id: int | None = None, + series_id: int | None = None, +) -> PhraseLookup: + """Load protected phrases and aliases into RAM lookup maps. + + Args: + session (Session): Active database session. + config (EbookSearchConfig): Runtime phrase-tuning settings. + book_id (int | None): Optional book scope to restrict loaded phrases. + series_id (int | None): Optional series scope to restrict loaded phrases. + + Returns: + PhraseLookup: Normalized phrase and alias maps with the token-window bounds to test. + """ + norm_to_ids: defaultdict[str, list[int]] = defaultdict(list) + alias_to_ids: defaultdict[str, list[int]] = defaultdict(list) + max_tokens = config.phrase_max_tokens + + phrase_statement = select( + EbookProtectedPhrase.id, + EbookProtectedPhrase.phrase_norm, + EbookProtectedPhrase.token_count, + ) + scope_filter = protected_phrase_scope_filter(book_id=book_id, series_id=series_id) + if scope_filter is not None: + phrase_statement = phrase_statement.where(scope_filter) + + for row in session.execute(phrase_statement): + phrase_id = int(row.id) + phrase_norm = str(row.phrase_norm) + norm_to_ids[phrase_norm].append(phrase_id) + max_tokens = max(max_tokens, int(row.token_count)) + + alias_statement = select( + EbookPhraseAlias.alias_norm, + EbookPhraseAlias.phrase_id, + ).join(EbookProtectedPhrase, EbookProtectedPhrase.id == EbookPhraseAlias.phrase_id) + if scope_filter is not None: + alias_statement = alias_statement.where(scope_filter) + + for row in session.execute(alias_statement): + alias_norm = str(row.alias_norm) + alias_to_ids[alias_norm].append(int(row.phrase_id)) + max_tokens = max(max_tokens, len(alias_norm.split())) + + return PhraseLookup( + norm_to_phrase_ids={key: tuple(values) for key, values in norm_to_ids.items()}, + alias_to_phrase_ids={key: tuple(values) for key, values in alias_to_ids.items()}, + min_tokens=config.phrase_min_tokens, + max_tokens=max_tokens, + ) + + +def protected_phrase_scope_filter(*, book_id: int | None, series_id: int | None) -> object | None: + """Build a SQLAlchemy filter for optional phrase book and series scope. + + Args: + book_id (int | None): Optional book scope to include alongside global phrases. + series_id (int | None): Optional series scope to include alongside global phrases. + + Returns: + object | None: A combined SQLAlchemy filter clause, or ``None`` when no scope is given. + """ + conditions = [] + if book_id is not None: + conditions.append(or_(EbookProtectedPhrase.book_id.is_(None), EbookProtectedPhrase.book_id == book_id)) + if series_id is not None: + conditions.append(or_(EbookProtectedPhrase.series_id.is_(None), EbookProtectedPhrase.series_id == series_id)) + if not conditions: + return None + return and_(*conditions) + + +def generate_query_ngrams( + tokens_: Sequence[str], + min_n: int, + max_n: int, +) -> Iterator[tuple[str, int, int]]: + """Generate normalized query windows from longest to shortest. + + Args: + tokens_ (Sequence[str]): Normalized query tokens. + min_n (int): Smallest window size to yield. + max_n (int): Largest window size to yield, capped at the token count. + + Yields: + tuple[str, int, int]: Normalized window text with its start and end token indices. + """ + capped_max_n = min(max_n, len(tokens_)) + for ngram_size in range(capped_max_n, min_n - 1, -1): + for start in range(len(tokens_) - ngram_size + 1): + end = start + ngram_size + phrase_norm = " ".join(tokens_[start:end]) + if phrase_norm in get_ignored_phrases(): + continue + yield phrase_norm, start, end + + +def detect_phrase_candidates(query_text: str, lookup: PhraseLookup) -> list[PhraseMatch]: + """Detect protected phrase windows in a user query using RAM hash lookups. + + Args: + query_text (str): User query text to scan. + lookup (PhraseLookup): In-memory phrase and alias lookup maps. + + Returns: + list[PhraseMatch]: Unhydrated phrase matches found in the query. + """ + return detect_phrase_candidates_from_tokens(tokenize_with_offsets(query_text), lookup) + + +def detect_phrase_candidates_in_text(text: str, lookup: PhraseLookup) -> list[PhraseMatch]: + """Detect protected phrase windows in arbitrary text with character offsets. + + Args: + text (str): Arbitrary text, such as a chunk, to scan. + lookup (PhraseLookup): In-memory phrase and alias lookup maps. + + Returns: + list[PhraseMatch]: Unhydrated phrase matches found in the text. + """ + return detect_phrase_candidates_from_tokens(tokenize_with_offsets(text), lookup) + + +def detect_phrase_candidates_from_tokens(tokens_: Sequence[NormalizedToken], lookup: PhraseLookup) -> list[PhraseMatch]: + """Detect protected phrase windows from already-normalized tokens. + + Args: + tokens_ (Sequence[NormalizedToken]): Normalized tokens with character offsets. + lookup (PhraseLookup): In-memory phrase and alias lookup maps. + + Returns: + list[PhraseMatch]: Deduplicated unhydrated phrase matches with token and character spans. + """ + matches: list[PhraseMatch] = [] + seen: set[tuple[int | None, str, int, int]] = set() + token_texts = [token.text for token in tokens_] + for phrase_norm, start, end in generate_query_ngrams(token_texts, min_n=lookup.min_tokens, max_n=lookup.max_tokens): + phrase_ids = lookup.norm_to_phrase_ids.get(phrase_norm, ()) + alias_ids = lookup.alias_to_phrase_ids.get(phrase_norm, ()) + for phrase_id in (*phrase_ids, *alias_ids): + key = (phrase_id, phrase_norm, start, end) + if key in seen: + continue + seen.add(key) + matches.append( + PhraseMatch( + phrase_norm=phrase_norm, + phrase_id=phrase_id, + start_token=start, + end_token=end, + token_count=end - start, + start_char=tokens_[start].start_char, + end_char=tokens_[end - 1].end_char, + ) + ) + return matches + + +def hydrate_matches(session: Session, matches: Sequence[PhraseMatch]) -> list[HydratedPhraseMatch]: + """Fetch protected phrase metadata for raw phrase matches. + + Args: + session (Session): Active database session. + matches (Sequence[PhraseMatch]): Unhydrated matches to enrich. + + Returns: + list[HydratedPhraseMatch]: Matches with protected-phrase metadata attached. + """ + if not matches: + return [] + + phrase_ids = sorted({match.phrase_id for match in matches if match.phrase_id is not None}) + if not phrase_ids: + return [] + + rows = { + row.id: row + for row in session.scalars(select(EbookProtectedPhrase).where(EbookProtectedPhrase.id.in_(phrase_ids))) + } + hydrated: list[HydratedPhraseMatch] = [] + for match in matches: + if match.phrase_id is None: + continue + phrase = rows.get(match.phrase_id) + if phrase is None: + continue + hydrated.append( + HydratedPhraseMatch( + phrase_id=phrase.id, + matched_norm=match.phrase_norm, + phrase_text=phrase.phrase_text, + phrase_norm=phrase.phrase_norm, + canonical_id=phrase.canonical_id, + phrase_type=phrase.phrase_type, + token_count=match.token_count, + confidence=phrase.confidence, + importance=phrase.importance, + allow_nested=phrase.allow_nested, + suppress_children=phrase.suppress_children, + start_token=match.start_token, + end_token=match.end_token, + start_char=match.start_char, + end_char=match.end_char, + book_id=phrase.book_id, + series_id=phrase.series_id, + ) + ) + return hydrated + + +def overlaps(first: HydratedPhraseMatch, second: HydratedPhraseMatch) -> bool: + """Return whether two token spans overlap. + + Args: + first (HydratedPhraseMatch): First match to compare. + second (HydratedPhraseMatch): Second match to compare. + + Returns: + bool: True when the two token spans share at least one token position. + """ + return not (first.end_token <= second.start_token or first.start_token >= second.end_token) + + +def is_inside(child: HydratedPhraseMatch, parent: HydratedPhraseMatch) -> bool: + """Return whether one token span is strictly inside another. + + Args: + child (HydratedPhraseMatch): Candidate nested match. + parent (HydratedPhraseMatch): Candidate enclosing match. + + Returns: + bool: True when ``child`` lies within ``parent`` and is not the same span. + """ + return ( + child.start_token >= parent.start_token + and child.end_token <= parent.end_token + and (child.start_token, child.end_token, child.phrase_id) + != (parent.start_token, parent.end_token, parent.phrase_id) + ) + + +def rank_match(match: HydratedPhraseMatch) -> tuple[float, float, int]: + """Rank phrase matches by importance, confidence, then token count. + + Args: + match (HydratedPhraseMatch): Match to build a sort key for. + + Returns: + tuple[float, float, int]: A comparable key of importance, confidence, and token count. + """ + return (match.importance, match.confidence, match.token_count) + + +def should_suppress(candidate: HydratedPhraseMatch, kept: HydratedPhraseMatch) -> bool: + """Return whether an already-kept match should suppress a candidate. + + Args: + candidate (HydratedPhraseMatch): Match being considered for keeping. + kept (HydratedPhraseMatch): Match already kept that may suppress the candidate. + + Returns: + bool: True when the candidate should be dropped in favor of the kept match. + """ + if not overlaps(candidate, kept): + return False + if candidate.canonical_id == kept.canonical_id: + return rank_match(kept) >= rank_match(candidate) + if is_inside(candidate, kept) and kept.suppress_children and not candidate.allow_nested: + return True + return not candidate.allow_nested and rank_match(kept) > rank_match(candidate) + + +def resolve_overlaps(matches: Sequence[HydratedPhraseMatch]) -> list[HydratedPhraseMatch]: + """Resolve overlapping phrase matches without relying only on longest match. + + Args: + matches (Sequence[HydratedPhraseMatch]): Hydrated matches that may overlap. + + Returns: + list[HydratedPhraseMatch]: The kept, non-suppressed matches. + """ + sorted_matches = sorted( + matches, + key=lambda match: (match.start_token, -match.token_count, -match.importance, -match.confidence), + ) + kept: list[HydratedPhraseMatch] = [] + for candidate in sorted_matches: + if any(should_suppress(candidate, existing) for existing in kept): + continue + kept.append(candidate) + return kept + + +def detect_protected_phrases_for_query( + session: Session, + query_text: str, + config: EbookSearchConfig, + *, + lookup: PhraseLookup | None = None, + book_id: int | None = None, + series_id: int | None = None, +) -> list[HydratedPhraseMatch]: + """Run the full online protected-phrase query-detection pipeline. + + Args: + session (Session): Active database session. + query_text (str): User query text to detect phrases in. + config (EbookSearchConfig): Runtime phrase-tuning settings. + lookup (PhraseLookup | None): Optional preloaded lookup; loaded on demand when ``None``. + book_id (int | None): Optional book scope for lookup loading. + series_id (int | None): Optional series scope for lookup loading. + + Returns: + list[HydratedPhraseMatch]: Hydrated, overlap-resolved phrase matches for the query. + """ + active_lookup = ( + lookup if lookup is not None else load_phrase_lookup(session, config, book_id=book_id, series_id=series_id) + ) + return resolve_overlaps(hydrate_matches(session, detect_phrase_candidates(query_text, active_lookup))) + + +def index_chunk_phrase_mentions_for_book( + session: Session, + book_id: int, + config: EbookSearchConfig, + *, + series_id: int | None = None, + lookup: PhraseLookup | None = None, +) -> int: + """Rebuild chunk phrase mentions for all chunks in one book. + + Args: + session (Session): Active database session. + book_id (int): Book whose chunk mentions are rebuilt. + config (EbookSearchConfig): Runtime phrase-tuning settings. + series_id (int | None): Optional series scope for lookup loading. + lookup (PhraseLookup | None): Optional preloaded lookup; loaded on demand when ``None``. + + Returns: + int: Total number of chunk phrase mentions indexed for the book. + """ + active_lookup = ( + lookup if lookup is not None else load_phrase_lookup(session, config, book_id=book_id, series_id=series_id) + ) + session.execute(delete(EbookChunkPhraseMention).where(EbookChunkPhraseMention.book_id == book_id)) + chunks = session.scalars(select(EbookChunk).where(EbookChunk.source_id == book_id).order_by(EbookChunk.id)) + count = 0 + for chunk in chunks: + count += index_chunk_phrase_mentions(session, chunk, lookup=active_lookup) + session.flush() + logger.info("ebook_chunk_phrase_mentions_indexed book_id=%s mentions=%s", book_id, count) + return count + + +def index_chunk_phrase_mentions(session: Session, chunk: EbookChunk, *, lookup: PhraseLookup) -> int: + """Store protected phrase mentions for one chunk. + + Args: + session (Session): Active database session. + chunk (EbookChunk): Chunk whose text is scanned for phrase mentions. + lookup (PhraseLookup): In-memory phrase and alias lookup maps. + + Returns: + int: Number of phrase mentions stored for the chunk. + """ + raw_matches = detect_phrase_candidates_in_text(chunk.text, lookup) + hydrated = resolve_overlaps(hydrate_matches(session, raw_matches)) + for match in hydrated: + session.add( + EbookChunkPhraseMention( + chunk_id=chunk.id, + phrase_id=match.phrase_id, + book_id=match.book_id if match.book_id is not None else chunk.source_id, + series_id=match.series_id, + start_char=match.start_char if match.start_char is not None else 0, + end_char=match.end_char, + ) + ) + return len(hydrated) + + +def phrase_hits_for_chunks( + session: Session, + *, + chunk_ids: Sequence[int], + phrase_ids: Sequence[int], +) -> dict[int, tuple[ChunkPhraseHit, ...]]: + """Return matched protected phrases with mention counts by chunk id using indexed chunk mentions. + + Args: + session (Session): Active database session. + chunk_ids (Sequence[int]): Chunk ids to look up mentions for. + phrase_ids (Sequence[int]): Protected phrase ids to restrict the results to. + + Returns: + dict[int, tuple[ChunkPhraseHit, ...]]: Phrase hits per chunk id, ordered by mention count. + """ + if not chunk_ids or not phrase_ids: + return {} + + mention_count = func.count(EbookChunkPhraseMention.phrase_id).label("mention_count") + statement = ( + select( + EbookChunkPhraseMention.chunk_id, + EbookProtectedPhrase.id.label("phrase_id"), + EbookProtectedPhrase.phrase_text, + mention_count, + ) + .join(EbookProtectedPhrase, EbookProtectedPhrase.id == EbookChunkPhraseMention.phrase_id) + .where( + EbookChunkPhraseMention.chunk_id.in_(chunk_ids), + EbookChunkPhraseMention.phrase_id.in_(phrase_ids), + ) + .group_by(EbookChunkPhraseMention.chunk_id, EbookProtectedPhrase.id, EbookProtectedPhrase.phrase_text) + .order_by(EbookChunkPhraseMention.chunk_id, mention_count.desc(), EbookProtectedPhrase.phrase_text) + ) + hits: defaultdict[int, list[ChunkPhraseHit]] = defaultdict(list) + for row in session.execute(statement): + hits[int(row.chunk_id)].append( + ChunkPhraseHit( + phrase_id=int(row.phrase_id), + phrase_text=str(row.phrase_text), + mention_count=int(row.mention_count), + ) + ) + return {chunk_id: tuple(chunk_hits) for chunk_id, chunk_hits in hits.items()} + + +def phrase_hit_counts_for_chunks( + session: Session, + *, + chunk_ids: Sequence[int], + phrase_ids: Sequence[int], +) -> dict[int, int]: + """Return phrase-hit counts by chunk id using indexed chunk mentions. + + Args: + session (Session): Active database session. + chunk_ids (Sequence[int]): Chunk ids to count mentions for. + phrase_ids (Sequence[int]): Protected phrase ids to restrict the counts to. + + Returns: + dict[int, int]: Total mention count per chunk id. + """ + hits = phrase_hits_for_chunks(session, chunk_ids=chunk_ids, phrase_ids=phrase_ids) + return {chunk_id: sum(hit.mention_count for hit in chunk_hits) for chunk_id, chunk_hits in hits.items()} diff --git a/python/ebook_search/protected_phrases/models.py b/python/ebook_search/protected_phrases/models.py new file mode 100644 index 0000000..2336fe6 --- /dev/null +++ b/python/ebook_search/protected_phrases/models.py @@ -0,0 +1,262 @@ +"""Dataclasses shared by protected phrase extraction, judging, matching, and backfills.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + from collections.abc import Mapping + + +@dataclass(slots=True) +class PhraseCandidate: + """A phrase candidate with merged extraction-source metadata. + + Attributes: + phrase_text (str): Display text for the phrase. + phrase_norm (str): Normalized phrase used as the merge key. + token_count (int): Number of normalized tokens in the phrase. + source_raw_ngram (bool): Whether the raw n-gram extractor produced the phrase. + source_yake (bool): Whether YAKE keyword extraction produced the phrase. + source_spacy_ner (bool): Whether spaCy named-entity recognition produced the phrase. + source_spacy_noun_chunk (bool): Whether spaCy noun chunking produced the phrase. + source_capitalized (bool): Whether the capitalized-run extractor produced the phrase. + source_metadata (bool): Whether book metadata produced the phrase. + spacy_label (str | None): spaCy entity label when NER produced the phrase. + raw_count (int): Occurrences counted across the book text. + chapter_count (int): Number of chapters containing the phrase. + yake_score (float | None): Raw YAKE score when available; lower is better. + candidate_score (float): Combined pre-judging score. + sample_contexts (list[str]): Normalized context snippets around occurrences. + """ + + phrase_text: str + phrase_norm: str + token_count: int + source_raw_ngram: bool = False + source_yake: bool = False + source_spacy_ner: bool = False + source_spacy_noun_chunk: bool = False + source_capitalized: bool = False + source_metadata: bool = False + spacy_label: str | None = None + raw_count: int = 0 + chapter_count: int = 0 + yake_score: float | None = None + candidate_score: float = 0.0 + sample_contexts: list[str] = field(default_factory=list) + + +@dataclass(frozen=True, slots=True) +class LLMJudgment: + """A structured phrase judgment returned by the LLM judge. + + Attributes: + keep (bool): Whether the judge accepted the phrase for protection. + canonical (str | None): Canonical phrase text chosen by the judge. + category (str | None): Phrase category such as person, place, or event. + aliases (tuple[str, ...]): Alternate surface forms for the phrase. + confidence (float): Judge confidence between 0.0 and 1.0. + importance (float): Judge importance between 0.0 and 1.0. + allow_nested (bool): Whether the phrase may match inside a larger kept match. + suppress_children (bool): Whether the phrase suppresses matches nested inside it. + reason (str | None): Free-text explanation from the judge. + """ + + keep: bool + canonical: str | None + category: str | None + aliases: tuple[str, ...] + confidence: float + importance: float = 0.5 + allow_nested: bool = False + suppress_children: bool = True + reason: str | None = None + + +@dataclass(frozen=True, slots=True) +class PhraseLookup: + """In-memory lookup maps used for constant-time phrase-window checks. + + Attributes: + norm_to_phrase_ids (Mapping[str, tuple[int, ...]]): Normalized phrase to protected phrase ids. + alias_to_phrase_ids (Mapping[str, tuple[int, ...]]): Normalized alias to protected phrase ids. + min_tokens (int): Smallest token-window size to test. + max_tokens (int): Largest token-window size to test. + """ + + norm_to_phrase_ids: Mapping[str, tuple[int, ...]] + alias_to_phrase_ids: Mapping[str, tuple[int, ...]] + min_tokens: int + max_tokens: int + + +@dataclass(frozen=True, slots=True) +class PhraseMatch: + """An unhydrated query or chunk phrase match. + + Attributes: + phrase_norm (str): Normalized text of the matched window. + start_token (int): Index of the first matched token. + end_token (int): Index one past the last matched token. + token_count (int): Number of tokens in the match. + phrase_id (int | None): Matched protected phrase id when known. + start_char (int | None): Start character offset in the source text. + end_char (int | None): End character offset in the source text. + """ + + phrase_norm: str + start_token: int + end_token: int + token_count: int + phrase_id: int | None = None + start_char: int | None = None + end_char: int | None = None + + +@dataclass(frozen=True, slots=True) +class HydratedPhraseMatch: + """A phrase match with protected-phrase metadata attached. + + Attributes: + phrase_id (int): Protected phrase id. + matched_norm (str): Normalized window text that matched. + phrase_text (str): Display text of the protected phrase. + phrase_norm (str): Normalized text of the protected phrase. + canonical_id (str): Deterministic ``category:slug`` identifier. + phrase_type (str | None): Phrase category. + token_count (int): Number of tokens in the match. + confidence (float): Stored judge confidence. + importance (float): Stored judge importance. + allow_nested (bool): Whether the phrase may match inside a larger kept match. + suppress_children (bool): Whether the phrase suppresses matches nested inside it. + start_token (int): Index of the first matched token. + end_token (int): Index one past the last matched token. + start_char (int | None): Start character offset in the source text. + end_char (int | None): End character offset in the source text. + book_id (int | None): Book scope of the phrase. + series_id (int | None): Series scope of the phrase. + """ + + phrase_id: int + matched_norm: str + phrase_text: str + phrase_norm: str + canonical_id: str + phrase_type: str | None + token_count: int + confidence: float + importance: float + allow_nested: bool + suppress_children: bool + start_token: int + end_token: int + start_char: int | None = None + end_char: int | None = None + book_id: int | None = None + series_id: int | None = None + + +@dataclass(frozen=True, slots=True) +class ChunkPhraseHit: + """One protected phrase with its mention count inside one retrieved chunk. + + Attributes: + phrase_id (int): Protected phrase id. + phrase_text (str): Display text of the protected phrase. + mention_count (int): Indexed mentions of the phrase in the chunk. + """ + + phrase_id: int + phrase_text: str + mention_count: int + + +@dataclass(frozen=True, slots=True) +class PhraseCandidateGenerationResult: + """Summary of candidate phrase extraction for indexed books. + + Attributes: + books_seen (int): Indexed books examined. + books_built (int): Books that had candidates generated and committed. + candidate_phrases (int): Candidate phrases stored across all books. + """ + + books_seen: int + books_built: int + candidate_phrases: int + + +@dataclass(frozen=True, slots=True) +class PhraseJudgmentBackfillResult: + """Summary of LLM judging for stored candidate phrases. + + Attributes: + books_seen (int): Indexed books examined. + books_judged (int): Books with judgments committed. + books_failed (int): Books rolled back after an error. + candidates_judged (int): Candidate phrases sent to the LLM judge. + protected_phrases (int): Protected phrases promoted from candidates. + phrase_mentions (int): Chunk phrase mentions indexed across all books. + """ + + books_seen: int + books_judged: int + books_failed: int + candidates_judged: int + protected_phrases: int + phrase_mentions: int + + +@dataclass(frozen=True, slots=True) +class BookJudgmentResult: + """Outcome of judging one book's candidate phrases. + + Attributes: + judged (int): Candidate phrases sent to the LLM judge. + protected (int): Protected phrases promoted from candidates. + mentions (int): Chunk phrase mentions indexed for the book. + committed (bool): Whether the book's judgments were committed. + failed (bool): Whether the book was rolled back after an error. + """ + + judged: int = 0 + protected: int = 0 + mentions: int = 0 + committed: bool = False + failed: bool = False + + +@dataclass(frozen=True, slots=True) +class BookCandidateResult: + """Outcome of generating one book's candidate phrases. + + Attributes: + candidates (int): Candidate phrases stored for the book. + built (bool): Whether candidate generation was committed. + """ + + candidates: int = 0 + built: bool = False + + +@dataclass(frozen=True, slots=True) +class PhraseRecalculationResult: + """Summary of phrase cleanup and candidate regeneration for one book. + + Attributes: + book_id (int): Book the recalculation ran against. + deleted_candidates (int): Candidate phrase rows deleted. + deleted_protected_phrases (int): Protected phrase rows deleted. + deleted_aliases (int): Phrase alias rows deleted. + deleted_mentions (int): Chunk phrase mention rows deleted. + candidate_phrases (int): Candidate phrases regenerated after cleanup. + """ + + book_id: int + deleted_candidates: int + deleted_protected_phrases: int + deleted_aliases: int + deleted_mentions: int + candidate_phrases: int diff --git a/python/ebook_search/protected_phrases/store.py b/python/ebook_search/protected_phrases/store.py new file mode 100644 index 0000000..2a24f77 --- /dev/null +++ b/python/ebook_search/protected_phrases/store.py @@ -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) diff --git a/python/ebook_search/protected_phrases/text_normalization.py b/python/ebook_search/protected_phrases/text_normalization.py index b651e9d..89535a0 100644 --- a/python/ebook_search/protected_phrases/text_normalization.py +++ b/python/ebook_search/protected_phrases/text_normalization.py @@ -2,10 +2,10 @@ from __future__ import annotations -import logging +import re from dataclasses import dataclass -logger = logging.getLogger(__name__) +JSON_OBJECT_RE = re.compile(r"\{.*\}", re.DOTALL) @dataclass(frozen=True, slots=True) @@ -18,17 +18,38 @@ class NormalizedToken: def normalize_text(text: str) -> str: - """Normalize text for phrase storage and lookup.""" + """Normalize text for phrase storage and lookup. + + Args: + text (str): Raw text to normalize. + + Returns: + str: Normalized tokens joined by single spaces. + """ return " ".join(token.text for token in tokenize_with_offsets(text)) def tokenize(text: str) -> list[str]: - """Normalize and split text into phrase-detection tokens.""" + """Normalize and split text into phrase-detection tokens. + + Args: + text (str): Raw text to tokenize. + + Returns: + list[str]: Normalized token strings. + """ return [token.text for token in tokenize_with_offsets(text)] def tokenize_with_offsets(text: str) -> list[NormalizedToken]: - """Normalize text into tokens while preserving original character offsets.""" + """Normalize text into tokens while preserving original character offsets. + + Args: + text (str): Raw text to tokenize. + + Returns: + list[NormalizedToken]: Normalized tokens with their source character spans. + """ tokens: list[NormalizedToken] = [] current: list[str] = [] start_char: int | None = None @@ -51,7 +72,14 @@ def tokenize_with_offsets(text: str) -> list[NormalizedToken]: def normalize_char(char: str) -> str: - """Normalize one character into a token character or a separator.""" + """Normalize one character into a token character or a separator. + + Args: + char (str): Single source character to normalize. + + Returns: + str: The normalized token character, or a space acting as a separator. + """ if char in {"\u2019", "\u2018"}: return "'" if char in {"-", "\u2013", "\u2014"}: diff --git a/python/ebook_search/search.py b/python/ebook_search/search.py index cd9ca1a..c59f87a 100644 --- a/python/ebook_search/search.py +++ b/python/ebook_search/search.py @@ -20,8 +20,7 @@ from python.ebook_search.bm25_corpus import ( score_bm25_corpus, ) from python.ebook_search.embeddings import MODEL_DIMENSIONS, embed_query, get_embedding_table -from python.ebook_search.protected_phrases.lib import ( - HydratedPhraseMatch, +from python.ebook_search.protected_phrases.matching import ( detect_protected_phrases_for_query, phrase_hits_for_chunks, ) @@ -40,6 +39,7 @@ if TYPE_CHECKING: from sqlalchemy.engine import Engine from python.ebook_search.config import EbookSearchConfig + from python.ebook_search.protected_phrases.models import HydratedPhraseMatch logger = logging.getLogger(__name__)