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