"""Protected phrase extraction, storage, and runtime matching.""" from __future__ import annotations import importlib import json import logging import re import tomllib from collections import defaultdict from dataclasses import dataclass from functools import cache from pathlib import Path 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.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__) PHRASE_CONFIG_DIR = Path(__file__).resolve().parent 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(frozen=True, slots=True) class NormalizedToken: """A normalized token plus its source character span.""" text: str start_char: int end_char: int @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 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 @cache def get_ignored_phrases() -> frozenset[str]: """Return ignored phrase strings loaded from TOML.""" return load_toml_string_set(PHRASE_CONFIG_DIR / "ignored_phrases.toml", "phrases") @cache def get_bad_ends() -> frozenset[str]: """Return bad phrase-ending tokens loaded from TOML.""" return load_toml_string_set(PHRASE_CONFIG_DIR / "bad_ends.toml", "tokens") @cache def get_bad_starts() -> frozenset[str]: """Return bad phrase-starting tokens loaded from TOML.""" return load_toml_string_set(PHRASE_CONFIG_DIR / "bad_starts.toml", "tokens") @cache def get_most_common_words() -> frozenset[str]: """Return the most common English words loaded from TOML.""" return load_toml_string_set(PHRASE_CONFIG_DIR / "most_common_words.toml", "words") def load_toml_string_set(path: Path, key: str) -> frozenset[str]: """Load and validate a TOML string list as a normalized immutable set.""" with path.open("rb") as file: body = tomllib.load(file) values = body.get(key) if not isinstance(values, list) or not all(isinstance(item, str) for item in values): msg = f"{path} must contain a {key!r} string list" raise ValueError(msg) return frozenset(normalize_text(value) for value in values if normalize_text(value)) def normalize_text(text: str) -> str: """Normalize text for phrase storage and lookup.""" 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.""" 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.""" tokens: list[NormalizedToken] = [] current: list[str] = [] start_char: int | None = None for index, char in enumerate(text): normalized = normalize_char(char) if normalized == " ": if current and start_char is not None: tokens.append(NormalizedToken(text="".join(current), start_char=start_char, end_char=index)) current = [] start_char = None continue if start_char is None: start_char = index current.append(normalized) if current and start_char is not None: tokens.append(NormalizedToken(text="".join(current), start_char=start_char, end_char=len(text))) return tokens def normalize_char(char: str) -> str: """Normalize one character into a token character or a separator.""" if char in {"\u2019", "\u2018"}: return "'" if char in {"-", "\u2013", "\u2014"}: return " " lowered = char.lower() if lowered in "abcdefghijklmnopqrstuvwxyz0123456789'": return lowered return " " 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 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() logger.exception( "ebook_candidate_phrase_judgment_book_failed source_id=%s book_number=%s books_seen=%s", source.id, book_number, books_seen, ) raise 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 candidates_judged=%s " "protected=%s mentions=%s", books_seen, books_judged, candidates_judged, protected_count, mention_count, ) return PhraseJudgmentBackfillResult( books_seen=books_seen, books_judged=books_judged, 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)}