Files
dotfiles/python/ebook_search/protected_phrases/lib.py
T

1934 lines
68 KiB
Python

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