Richie/dotfiles
1
0
Fork 0
Code Issues 1 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

feat(ebook): migrate to async DB/HTTP and parallelize phrase pipeline

Browse Source 
Convert the ebook-search web app to async end to end and add concurrency
to the protected-phrase extraction and judging pipeline so large books no
longer block the event loop or the UI.

ORM / infra:
- Add get_async_postgres_engine and factor shared URL/connect_args building
  into build_postgres_url (reused by the sync and async engine builders)
- Add async FastAPI session helpers (get_async_db, AsyncDbSession) with
  expire_on_commit=False to avoid implicit IO under asyncio

App:
- Use AsyncEngine/AsyncSession throughout routes, search, ingest, embeddings,
  answer, rerank and LLM calls; convert handlers to async
- Share a single httpx.AsyncClient in app state for LLM requests; size the
  connection pool for concurrent phrase-judging workers
- Add judge_tasks: run per-book judging as tracked background tasks so a
  book already being judged isn't double-queued

Protected phrases:
- Add a process pool (pool.py) and worker-count config
  (extraction/judge book/phrase workers) to parallelize candidate generation
  and judging
- Split admin actions into all/missing variants for generation and judging

Config:
- Add protected_phrase_extraction_workers, phrase_judge_book_workers,
  phrase_judge_phrase_workers
This commit is contained in:
Richie
2026-07-16 13:09:34 -04:00
parent 9bc0d67962
commit 1f5e527cbc
29 changed files with 1824 additions and 769 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/ebook_search/protected_phrases/extraction.py
+185 -79
View File
@@ -4,7 +4,7 @@ from __future__ import annotations
import logging
import re
from collections import defaultdict
from collections import Counter, defaultdict
from functools import lru_cache
from time import perf_counter
from typing import TYPE_CHECKING, Protocol
@@ -15,6 +15,7 @@ from python.ebook_search.protected_phrases.config import (
get_bad_ends,
get_bad_starts,
get_ignored_phrases,
get_junk_tokens,
get_most_common_words,
)
from python.ebook_search.protected_phrases.models import PhraseCandidate
@@ -28,16 +29,10 @@ if TYPE_CHECKING:
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",
)
BAD_START_SCORE_PENALTY = 10.0
BAD_END_SCORE_PENALTY = 10.0
MULTI_SOURCE_SCORE_BONUS = 2.0
MULTI_SOURCE_MIN_SOURCES = 2
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}")
@@ -91,21 +86,6 @@ class YakeExtractorFactory(Protocol):
"""
def strip_leading_articles(phrase_norm: str) -> str:
"""Remove one leading English article from a normalized phrase.
Args:
phrase_norm (str): Normalized phrase text to strip.
Returns:
str: The phrase with a single leading ``the``, ``a``, or ``an`` removed.
"""
tokens_ = phrase_norm.split()
if tokens_ and tokens_[0] in {"the", "a", "an"}:
tokens_ = tokens_[1:]
return " ".join(tokens_)
def normalize_candidate_phrase(
phrase_text: str,
config: EbookSearchConfig,
@@ -146,35 +126,71 @@ def normalize_candidate_phrase(
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.
def count_raw_ngrams(tokens: Sequence[str], config: EbookSearchConfig) -> Counter[str]:
"""Count every n-gram window in one normalized token block.
``tokens`` are already normalized (see :func:`tokenize`), so each window's normalized form
is the joined tokens directly. Counting into a plain :class:`Counter` rather than
:class:`PhraseCandidate` objects keeps this hot loop cheap; callers filter ignored phrases
and materialize candidates per unique phrase afterwards, which is far fewer operations than
doing either per window.
Args:
text (str): Book text to slide n-gram windows over.
tokens (Sequence[str]): Normalized tokens for one text block.
config (EbookSearchConfig): Runtime phrase-tuning settings.
Returns:
dict[str, PhraseCandidate]: Candidates keyed by normalized phrase, with raw counts.
Counter[str]: Raw occurrence counts keyed by normalized phrase.
"""
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
return Counter(
" ".join(tokens[start : start + ngram_size])
for ngram_size in range(config.phrase_min_tokens, config.phrase_max_tokens + 1)
for start in range(len(tokens) - ngram_size + 1)
)
def extract_raw_ngrams_by_chapter(
chapters: Sequence[str],
config: EbookSearchConfig,
) -> dict[str, PhraseCandidate]:
"""Extract raw n-grams across chapters, tracking both raw counts and chapter spread.
Counting each chapter separately makes chapter spread fall out of dict membership: a phrase's
``chapter_count`` is simply how many per-chapter count maps contain it, so no per-window seen
tracking is needed. This also lets the enrichment step skip re-sliding the same n-gram sizes.
Phrases below the minimum raw count are dropped here rather than materialized: most unique
n-grams occur once, and :func:`filter_storable_candidates` would discard them as too rare
anyway, so building ``PhraseCandidate`` objects for them is wasted work.
Args:
chapters (Sequence[str]): Chapter-like text blocks to slide n-gram windows over.
config (EbookSearchConfig): Runtime phrase-tuning settings.
Returns:
dict[str, PhraseCandidate]: Candidates meeting the minimum raw count, keyed by normalized
phrase, with raw and chapter counts.
"""
chapter_count_maps = [count_raw_ngrams(tokenize(chapter), config) for chapter in chapters]
total_counts: Counter[str] = Counter()
chapter_spread: Counter[str] = Counter()
for chapter_counts in chapter_count_maps:
total_counts.update(chapter_counts)
chapter_spread.update(chapter_counts.keys())
min_raw_count = minimum_candidate_raw_count(config)
ignored = get_ignored_phrases()
return {
phrase_norm: PhraseCandidate(
phrase_text=phrase_norm,
phrase_norm=phrase_norm,
token_count=phrase_norm.count(" ") + 1,
source_raw_ngram=True,
raw_count=raw_count,
chapter_count=chapter_spread[phrase_norm],
)
for phrase_norm, raw_count in total_counts.items()
if raw_count >= min_raw_count and phrase_norm not in ignored
}
@lru_cache(maxsize=2)
@@ -381,6 +397,7 @@ def merge_candidate(existing: PhraseCandidate, item: PhraseCandidate) -> None:
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
existing.chapter_count = max(existing.chapter_count, item.chapter_count)
if item.yake_score is not None:
existing.yake_score = item.yake_score
if item.spacy_label:
@@ -390,12 +407,19 @@ def merge_candidate(existing: PhraseCandidate, item: PhraseCandidate) -> None:
def enrich_with_frequency_and_chapter_counts(
candidates: Mapping[str, PhraseCandidate],
chapters: Sequence[str],
*,
counted_sizes: Iterable[int] = (),
) -> dict[str, PhraseCandidate]:
"""Add raw occurrence and chapter-spread counts to candidates.
Candidates whose ``token_count`` is in ``counted_sizes`` are left untouched: those counts
were already computed while sliding the chapters in :func:`extract_raw_ngrams_by_chapter`,
so re-sliding those n-gram sizes here would just duplicate that work.
Args:
candidates (Mapping[str, PhraseCandidate]): Candidates to enrich, keyed by normalized phrase.
chapters (Sequence[str]): Chapter-like text blocks used to count occurrences and spread.
counted_sizes (Iterable[int]): Token counts whose counts are already populated and should be skipped.
Returns:
dict[str, PhraseCandidate]: Candidates with updated ``raw_count`` and ``chapter_count`` values.
@@ -403,10 +427,40 @@ def enrich_with_frequency_and_chapter_counts(
if not candidates:
return {}
already_counted = set(counted_sizes)
candidate_sets_by_size: dict[int, set[str]] = defaultdict(set)
for phrase_norm, candidate in candidates.items():
if candidate.token_count in already_counted:
continue
candidate_sets_by_size[candidate.token_count].add(phrase_norm)
enriched = dict(candidates)
if not candidate_sets_by_size:
return enriched
total_counts, chapter_counts = count_candidate_occurrences(candidate_sets_by_size, chapters)
for phrase_norm, candidate in enriched.items():
if candidate.token_count in already_counted:
continue
candidate.raw_count = max(candidate.raw_count, total_counts[phrase_norm])
candidate.chapter_count = chapter_counts[phrase_norm]
return enriched
def count_candidate_occurrences(
candidate_sets_by_size: Mapping[int, set[str]],
chapters: Sequence[str],
) -> tuple[dict[str, int], dict[str, int]]:
"""Count total occurrences and chapter spread for candidate phrases across chapters.
Args:
candidate_sets_by_size (Mapping[int, set[str]]): Candidate normalized phrases grouped by token count.
chapters (Sequence[str]): Chapter-like text blocks to slide n-gram windows over.
Returns:
tuple[dict[str, int], dict[str, int]]: Total occurrence counts and chapter-spread counts,
each keyed by normalized phrase.
"""
total_counts: defaultdict[str, int] = defaultdict(int)
chapter_counts: defaultdict[str, int] = defaultdict(int)
for chapter in chapters:
@@ -421,18 +475,13 @@ def enrich_with_frequency_and_chapter_counts(
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
return total_counts, chapter_counts
def filter_storable_candidates(
candidates: Mapping[str, PhraseCandidate],
config: EbookSearchConfig,
) -> tuple[dict[str, PhraseCandidate], int, int, int]:
) -> tuple[dict[str, PhraseCandidate], int, int, int, int]:
"""Remove candidates that should not be persisted.
Args:
@@ -440,14 +489,15 @@ def filter_storable_candidates(
config (EbookSearchConfig): Runtime phrase-tuning settings.
Returns:
tuple[dict[str, PhraseCandidate], int, int, int]: The storable candidates followed by the counts
dropped for being too short, too rare, and too common.
tuple[dict[str, PhraseCandidate], int, int, int, int]: The storable candidates followed by the
counts dropped for being too short, too rare, too common, and junk.
"""
min_raw_count = minimum_candidate_raw_count(config)
filtered: dict[str, PhraseCandidate] = {}
too_short = 0
too_rare = 0
too_common = 0
junk = 0
for phrase_norm, candidate in candidates.items():
if candidate.token_count < config.phrase_min_tokens:
too_short += 1
@@ -455,11 +505,15 @@ def filter_storable_candidates(
if candidate.raw_count < min_raw_count:
too_rare += 1
continue
if is_most_common_word_phrase(phrase_norm):
phrase_tokens = phrase_norm.split()
if is_most_common_word_phrase(phrase_tokens):
too_common += 1
continue
if is_junk_phrase(phrase_tokens):
junk += 1
continue
filtered[phrase_norm] = candidate
return filtered, too_short, too_rare, too_common
return filtered, too_short, too_rare, too_common, junk
def minimum_candidate_raw_count(config: EbookSearchConfig) -> int:
@@ -474,18 +528,41 @@ def minimum_candidate_raw_count(config: EbookSearchConfig) -> int:
return max(config.phrase_raw_ngram_min_count, 1)
def is_most_common_word_phrase(phrase_norm: str) -> bool:
def is_most_common_word_phrase(phrase_tokens: list[str]) -> bool:
"""Return whether every token in a normalized phrase is a common word.
Args:
phrase_norm (str): Normalized phrase text to inspect.
phrase_tokens (list[str]): Normalized phrase tokens to inspect.
Returns:
bool: True when the phrase is non-empty and every token is a common word.
"""
tokens_ = phrase_norm.split()
common_words = get_most_common_words()
return bool(tokens_) and all(token in common_words for token in tokens_)
return bool(phrase_tokens) and all(token in common_words for token in phrase_tokens)
def is_junk_phrase(phrase_tokens: list[str]) -> bool:
"""Return whether a normalized phrase is lexical junk not worth LLM judging.
Judged data shows phrases containing a dialogue/action verb or a pronoun contraction are
never kept, and phrases whose tokens are mostly common words almost never are. Possessives
of proper nouns (``chapman's death``) pass because matching is by exact token, and
exactly-half-common bigrams (``data feed``) pass because the common-word rule is strict.
Args:
phrase_tokens (list[str]): Normalized phrase tokens to inspect.
Returns:
bool: True when the phrase contains a junk token or is majority common words.
"""
if not phrase_tokens:
return False
junk_tokens = get_junk_tokens()
if any(token in junk_tokens for token in phrase_tokens):
return True
common_words = get_most_common_words()
half_phrase_len = len(phrase_tokens) // 2
return sum(token in common_words for token in phrase_tokens) > half_phrase_len
def score_candidate(candidate: PhraseCandidate, config: EbookSearchConfig) -> float:
@@ -499,6 +576,8 @@ def score_candidate(candidate: PhraseCandidate, config: EbookSearchConfig) -> fl
float: Combined score from sources, frequency, and length, less any penalties.
"""
score = source_score(candidate) + frequency_score(candidate, config) + token_count_score(candidate, config)
if non_raw_source_count(candidate) >= MULTI_SOURCE_MIN_SOURCES:
score += MULTI_SOURCE_SCORE_BONUS
if candidate.phrase_norm in get_ignored_phrases():
score -= 100.0
if has_bad_start(candidate.phrase_norm):
@@ -508,6 +587,26 @@ def score_candidate(candidate: PhraseCandidate, config: EbookSearchConfig) -> fl
return score
def non_raw_source_count(candidate: PhraseCandidate) -> int:
"""Count the non-raw-ngram extraction sources that produced a candidate.
Args:
candidate (PhraseCandidate): Candidate whose enabled sources are counted.
Returns:
int: Number of enabled sources other than the raw n-gram slide.
"""
return sum(
(
candidate.source_yake,
candidate.source_spacy_ner,
candidate.source_spacy_noun_chunk,
candidate.source_capitalized,
candidate.source_metadata,
)
)
def has_bad_start(phrase_norm: str) -> bool:
"""Return whether a normalized phrase starts with a bad starting token.
@@ -517,8 +616,8 @@ def has_bad_start(phrase_norm: str) -> bool:
Returns:
bool: True when the first token is a known bad starting token.
"""
tokens_ = phrase_norm.split()
return bool(tokens_ and tokens_[0] in get_bad_starts())
phrase_tokens = phrase_norm.split()
return bool(phrase_tokens and phrase_tokens[0] in get_bad_starts())
def has_bad_end(phrase_norm: str) -> bool:
@@ -530,8 +629,8 @@ def has_bad_end(phrase_norm: str) -> bool:
Returns:
bool: True when the last token is a known bad ending token.
"""
tokens_ = phrase_norm.split()
return bool(tokens_ and tokens_[-1] in get_bad_ends())
phrase_tokens = phrase_norm.split()
return bool(phrase_tokens and phrase_tokens[-1] in get_bad_ends())
def source_score(candidate: PhraseCandidate) -> float:
@@ -550,6 +649,7 @@ def source_score(candidate: PhraseCandidate) -> float:
(candidate.source_spacy_ner, 2.5),
(candidate.source_spacy_noun_chunk, 1.5),
(candidate.source_capitalized, 2.0),
(candidate.source_metadata, 2.0),
(candidate.source_raw_ngram, 0.5),
)
if enabled
@@ -569,10 +669,10 @@ def frequency_score(candidate: PhraseCandidate, config: EbookSearchConfig) -> fl
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),
(candidate.raw_count, config.phrase_raw_count_score_threshold, 0.5),
(candidate.raw_count, config.phrase_raw_count_high_score_threshold, 0.5),
(candidate.chapter_count, config.phrase_chapter_count_score_threshold, 0.5),
(candidate.chapter_count, config.phrase_chapter_count_high_score_threshold, 0.5),
)
if count >= threshold
)
@@ -679,7 +779,7 @@ def extract_phrase_candidates_for_book(
config.protected_phrase_max_candidates_per_book,
)
raw_started_at = perf_counter()
raw = extract_raw_ngrams(book_text, config)
raw = extract_raw_ngrams_by_chapter(chapters, config)
logger.info(
"ebook_phrase_candidate_extract_raw_complete candidates=%s duration_ms=%.1f",
len(raw),
@@ -713,11 +813,16 @@ def extract_phrase_candidates_for_book(
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(
# Raw n-gram sizes were already counted per chapter above, so only enrich the remaining
# (entity-length) sizes here instead of re-sliding every size over the whole book.
candidates = enrich_with_frequency_and_chapter_counts(
candidates,
config,
chapters,
counted_sizes=range(config.phrase_min_tokens, config.phrase_max_tokens + 1),
)
pre_filter_count = len(candidates)
candidates, filtered_too_short, filtered_too_rare, filtered_too_common, filtered_junk = filter_storable_candidates(
candidates, config
)
for candidate in candidates.values():
candidate.candidate_score = score_candidate(candidate, config)
@@ -727,8 +832,8 @@ def extract_phrase_candidates_for_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",
"merged=%s filtered_too_short=%s filtered_too_rare=%s filtered_too_common=%s filtered_junk=%s "
"min_uses=%s storable=%s limited=%s enrich_score_ms=%.1f duration_ms=%.1f",
len(raw),
len(yake_candidates),
len(spacy_candidates),
@@ -738,6 +843,7 @@ def extract_phrase_candidates_for_book(
filtered_too_short,
filtered_too_rare,
filtered_too_common,
filtered_junk,
minimum_candidate_raw_count(config),
len(candidates),
len(limited),
python/ebook_search/protected_phrases/generate_ngrams.py
+191 -95
View File
@@ -2,7 +2,9 @@
from __future__ import annotations
import asyncio
import logging
from collections import deque
from time import perf_counter
from typing import TYPE_CHECKING
@@ -14,44 +16,55 @@ from python.ebook_search.protected_phrases.models import (
PhraseCandidateGenerationResult,
PhraseRecalculationResult,
)
from python.ebook_search.protected_phrases.pool import extract_phrase_candidates_in_pool, get_extraction_pool
from python.ebook_search.protected_phrases.store import (
bulk_upsert_unjudged_candidates,
delete_phrase_data_for_book,
load_book_chapter_texts,
metadata_for_source,
new_candidate_row,
prune_unstorable_unjudged_candidate_phrases,
save_candidate_to_db,
)
from python.orm.richie import EbookSource
from python.orm.richie import EbookCandidatePhrase, EbookSource
if TYPE_CHECKING:
from collections.abc import Mapping, Sequence
from concurrent.futures import Future
from sqlalchemy.orm import Session
from sqlalchemy.ext.asyncio import AsyncSession
from python.ebook_search.config import EbookSearchConfig
from python.ebook_search.protected_phrases.extraction import SpacyLanguage
from python.orm.richie import EbookCandidatePhrase
from python.ebook_search.protected_phrases.models import PhraseCandidate
logger = logging.getLogger(__name__)
def generate_candidate_phrases_for_books(
session: Session,
async def generate_candidate_phrases_for_books(
session: AsyncSession,
config: EbookSearchConfig,
*,
nlp: SpacyLanguage | None = None,
only_missing: bool = False,
) -> PhraseCandidateGenerationResult:
"""Create or refresh candidate phrases for indexed books without calling the LLM judge.
Extraction always runs concurrently in the shared process pool so a full backfill uses
multiple cores.
Args:
session (Session): Active database session.
config (EbookSearchConfig): Runtime phrase-tuning settings.
nlp (SpacyLanguage | None): Optional spaCy pipeline for entity and noun-chunk sources.
only_missing (bool): When True, only generate for books that have no candidate phrases
yet instead of refreshing every book.
Returns:
PhraseCandidateGenerationResult: Per-corpus counts of books seen, built, and candidates stored.
"""
sources = session.scalars(select(EbookSource).order_by(EbookSource.id)).all()
source_query = select(EbookSource).order_by(EbookSource.id)
if only_missing:
has_candidates = select(EbookCandidatePhrase.id).where(EbookCandidatePhrase.book_id == EbookSource.id)
source_query = source_query.where(~has_candidates.exists())
sources = (await session.scalars(source_query)).all()
books_seen = len(sources)
logger.info(
"ebook_candidate_phrase_generation_start books_seen=%s min_tokens=%s max_tokens=%s max_candidates_per_book=%s",
@@ -61,7 +74,7 @@ def generate_candidate_phrases_for_books(
config.protected_phrase_max_candidates_per_book,
)
outcomes = [generate_candidates_for_source(session, source, config, nlp=nlp) for source in sources]
outcomes = await generate_candidates_for_sources_pooled(session, sources, config)
result = PhraseCandidateGenerationResult(
books_seen=books_seen,
@@ -77,76 +90,99 @@ def generate_candidate_phrases_for_books(
return result
def generate_candidates_for_source(
session: Session,
source: EbookSource,
async def generate_candidates_for_sources_pooled(
session: AsyncSession,
sources: Sequence[EbookSource],
config: EbookSearchConfig,
*,
nlp: SpacyLanguage | None = None,
) -> BookCandidateResult:
"""Generate and store candidate phrases for one book, managing its own transaction.
) -> list[BookCandidateResult]:
"""Generate candidate phrases for many books, extracting them concurrently in worker processes.
Commits on success; rolls back and re-raises on error so callers stop the backfill.
Chapter loading and row persistence stay on the caller's session (serial), while the CPU-bound
extraction runs in the shared process pool. A bounded window of in-flight books overlaps
extraction across cores without loading every book's candidates into memory at once.
Args:
session (Session): Active database session.
source (EbookSource): Indexed book to generate candidates for.
sources (Sequence[EbookSource]): Indexed books to generate candidates for.
config (EbookSearchConfig): Runtime phrase-tuning settings.
nlp (SpacyLanguage | None): Optional spaCy pipeline for entity and noun-chunk sources.
Returns:
BookCandidateResult: Candidate count and whether the book was committed.
list[BookCandidateResult]: One result per book.
"""
pool = get_extraction_pool(config.protected_phrase_extraction_workers)
max_in_flight = max(1, config.protected_phrase_extraction_workers) * 2
pending: deque[tuple[EbookSource, Future[list[PhraseCandidate]]]] = deque()
outcomes: list[BookCandidateResult] = []
async def drain_one() -> None:
source, future = pending.popleft()
extracted = await asyncio.wrap_future(future)
outcomes.append(await store_source_candidates(session, source, extracted, config))
try:
for source in sources:
chapters = await load_book_chapter_texts(session, source.id)
if not chapters:
logger.warning("ebook_candidate_phrase_generation_book_empty source_id=%s", source.id)
outcomes.append(BookCandidateResult())
continue
future = pool.submit(
extract_phrase_candidates_for_book,
"\n\n".join(chapters),
chapters,
config,
metadata=metadata_for_source(source),
)
pending.append((source, future))
if len(pending) >= max_in_flight:
await drain_one()
while pending:
await drain_one()
except Exception:
for _, future in pending:
future.cancel()
await session.rollback()
logger.exception("ebook_candidate_phrase_generation_pooled_failed")
raise
return outcomes
async def store_source_candidates(
session: AsyncSession,
source: EbookSource,
limited_candidates: list[PhraseCandidate],
config: EbookSearchConfig,
) -> BookCandidateResult:
"""Persist and commit one book's already-extracted candidates.
Args:
session (AsyncSession): Active database session.
source (EbookSource): Book the candidates belong to.
limited_candidates (list[PhraseCandidate]): Scored candidates to persist.
config (EbookSearchConfig): Runtime phrase-tuning settings.
Returns:
BookCandidateResult: Candidate count and that the book was committed.
"""
book_started_at = perf_counter()
logger.info(
"ebook_candidate_phrase_generation_book_start source_id=%s title=%r",
source.id,
source.title,
)
try:
chapters = load_book_chapter_texts(session, source.id)
if not chapters:
logger.warning("ebook_candidate_phrase_generation_book_empty source_id=%s", source.id)
return BookCandidateResult()
book_text = "\n\n".join(chapters)
logger.info(
"ebook_candidate_phrase_generation_book_loaded source_id=%s chapters=%s chars=%s",
source.id,
len(chapters),
len(book_text),
)
candidates = generate_candidate_phrases_for_book(
session,
source.id,
series_id=None,
book_text=book_text,
chapters=chapters,
config=config,
nlp=nlp,
metadata=metadata_for_source(source),
)
session.commit()
except Exception:
session.rollback()
logger.exception("ebook_candidate_phrase_generation_book_failed source_id=%s", source.id)
raise
saved_count = await store_candidate_phrases_for_book(session, source.id, None, limited_candidates, config)
await session.commit()
logger.info(
"ebook_candidate_phrase_generation_book_committed source_id=%s candidates=%s duration_ms=%.1f",
source.id,
len(candidates),
saved_count,
(perf_counter() - book_started_at) * 1000,
)
return BookCandidateResult(candidates=len(candidates), built=True)
return BookCandidateResult(candidates=saved_count, built=True)
def recalculate_candidate_phrases_for_book(
session: Session,
async def recalculate_candidate_phrases_for_book(
session: AsyncSession,
source: EbookSource,
config: EbookSearchConfig,
*,
nlp: SpacyLanguage | None = None,
use_process_pool: bool = False,
) -> PhraseRecalculationResult:
"""Remove all book phrase data, regenerate candidates, and commit the completed book.
@@ -155,6 +191,9 @@ def recalculate_candidate_phrases_for_book(
source (EbookSource): Indexed book to recalculate.
config (EbookSearchConfig): Runtime phrase-tuning settings.
nlp (SpacyLanguage | None): Optional spaCy pipeline for entity and noun-chunk sources.
use_process_pool (bool): Run the CPU-bound extraction in a worker process so concurrent
recalculations do not serialize behind the GIL. Defaults to in-process for callers
(tests, backfills) that do not need it.
Returns:
PhraseRecalculationResult: Deleted-row counts and the number of candidates regenerated.
@@ -166,11 +205,11 @@ def recalculate_candidate_phrases_for_book(
source.title,
)
try:
deleted = delete_phrase_data_for_book(session, source.id)
chapters = load_book_chapter_texts(session, source.id)
deleted = await delete_phrase_data_for_book(session, source.id)
chapters = await 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()
await session.commit()
return PhraseRecalculationResult(
book_id=source.id,
deleted_candidates=deleted.deleted_candidates,
@@ -180,19 +219,20 @@ def recalculate_candidate_phrases_for_book(
candidate_phrases=0,
)
candidates = generate_candidate_phrases_for_book(
candidate_count = await 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),
replace_all=True,
use_process_pool=use_process_pool,
)
session.commit()
await session.commit()
except Exception:
session.rollback()
await session.rollback()
logger.exception("ebook_candidate_phrase_recalculation_failed source_id=%s", source.id)
raise
@@ -202,7 +242,7 @@ def recalculate_candidate_phrases_for_book(
deleted_protected_phrases=deleted.deleted_protected_phrases,
deleted_aliases=deleted.deleted_aliases,
deleted_mentions=deleted.deleted_mentions,
candidate_phrases=len(candidates),
candidate_phrases=candidate_count,
)
logger.info(
"ebook_candidate_phrase_recalculation_complete source_id=%s deleted_candidates=%s "
@@ -218,57 +258,113 @@ def recalculate_candidate_phrases_for_book(
return result
def generate_candidate_phrases_for_book(
session: Session,
async def generate_candidate_phrases_for_book(
session: AsyncSession,
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]:
replace_all: bool = False,
use_process_pool: bool = False,
) -> int:
"""Extract and store candidate phrases for one book without LLM judging.
Args:
session (Session): Active database session.
book_id (int): Book the candidates belong to.
series_id (int | None): Series scope for the stored candidates.
book_text (str): Full book text used for extraction.
chapters (Sequence[str]): Chapter-like text blocks used for frequency counts.
chapters (Sequence[str]): Chapter-like text blocks used for extraction and frequency counts.
config (EbookSearchConfig): Runtime phrase-tuning settings.
nlp (SpacyLanguage | None): Optional spaCy pipeline for entity and noun-chunk sources.
metadata (Mapping[str, object] | None): Optional book metadata used as a candidate source.
replace_all (bool): When the caller has already cleared this book's candidates (e.g. a
recalculation), skip the per-candidate existence lookup and bulk-insert new rows.
use_process_pool (bool): Run the CPU-bound extraction in a worker process to avoid
serializing concurrent requests behind the GIL. Ignored when ``nlp`` is set, since
the spaCy pipeline cannot be sent to a worker process.
Returns:
list[EbookCandidatePhrase]: The stored candidate phrase rows.
int: Number of candidate phrase rows stored.
"""
started_at = perf_counter()
limited_candidates = extract_phrase_candidates_for_book(
book_text,
chapters,
book_text = "\n\n".join(chapters)
if use_process_pool and nlp is None:
limited_candidates = await extract_phrase_candidates_in_pool(book_text, chapters, config, metadata=metadata)
else:
limited_candidates = extract_phrase_candidates_for_book(
book_text,
chapters,
config,
nlp=nlp,
metadata=metadata,
)
saved_count = await store_candidate_phrases_for_book(
session,
book_id,
series_id,
limited_candidates,
config,
nlp=nlp,
metadata=metadata,
replace_all=replace_all,
)
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",
"ebook_candidate_phrase_generation_book_duration book_id=%s candidates=%s duration_ms=%.1f",
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,
saved_count,
(perf_counter() - started_at) * 1000,
)
return rows
return saved_count
async def store_candidate_phrases_for_book(
session: AsyncSession,
book_id: int,
series_id: int | None,
limited_candidates: list[PhraseCandidate],
config: EbookSearchConfig,
*,
replace_all: bool = False,
) -> int:
"""Persist already-extracted candidate phrase rows for one book without committing.
Args:
session (Session): Active database session.
book_id (int): Book the candidates belong to.
series_id (int | None): Series scope for the stored candidates.
limited_candidates (list[PhraseCandidate]): Scored candidates to persist.
config (EbookSearchConfig): Runtime phrase-tuning settings.
replace_all (bool): When the caller has already cleared this book's candidates, skip the
per-candidate existence lookup and bulk-insert new rows.
Returns:
int: Number of candidate phrase rows stored.
"""
save_started_at = perf_counter()
if replace_all:
rows = [new_candidate_row(book_id, series_id, candidate) for candidate in limited_candidates]
session.add_all(rows)
await session.flush()
saved_count = len(rows)
logger.info(
"ebook_candidate_phrase_save_start book_id=%s candidates=%s mode=bulk_insert",
book_id,
len(limited_candidates),
)
else:
pruned_count = await 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,
)
saved_count = await bulk_upsert_unjudged_candidates(session, book_id, series_id, limited_candidates)
logger.info(
"ebook_candidate_phrase_save_complete book_id=%s candidates=%s save_ms=%.1f",
book_id,
saved_count,
(perf_counter() - save_started_at) * 1000,
)
return saved_count
python/ebook_search/protected_phrases/judge_ngrams.py
+235 -205
View File
@@ -2,20 +2,24 @@
from __future__ import annotations
import asyncio
import json
import logging
import re
from time import perf_counter
from typing import TYPE_CHECKING
import httpx
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession
from python.ebook_search.llm_interface import request_chat_completion
from python.ebook_search.protected_phrases.extraction import (
candidate_source_names,
get_sample_contexts,
is_junk_phrase,
is_most_common_word_phrase,
minimum_candidate_raw_count,
score_candidate,
)
from python.ebook_search.protected_phrases.matching import index_chunk_phrase_mentions_for_book
from python.ebook_search.protected_phrases.models import BookJudgmentResult, LLMJudgment, PhraseJudgmentBackfillResult
@@ -32,47 +36,65 @@ from python.ebook_search.protected_phrases.text_normalization import normalize_t
from python.orm.richie import EbookSource
if TYPE_CHECKING:
from sqlalchemy.orm import Session
from collections.abc import Sequence
from sqlalchemy.ext.asyncio import AsyncEngine
from python.ebook_search.config import EbookSearchConfig
from python.ebook_search.protected_phrases.models import PhraseCandidate
from python.orm.richie import EbookCandidatePhrase, EbookProtectedPhrase
from python.orm.richie import EbookProtectedPhrase
JSON_OBJECT_RE = re.compile(r"\{.*\}", re.DOTALL)
logger = logging.getLogger(__name__)
def judge_candidate_phrases_for_books(
session: Session,
async def judge_candidate_phrases_for_books(
engine: AsyncEngine,
config: EbookSearchConfig,
*,
source_ids: Sequence[int] | None = None,
) -> PhraseJudgmentBackfillResult:
"""Judge stored candidate phrases and promote accepted phrases for indexed books.
"""Judge candidate phrases for books, fanning LLM calls out across books and phrases.
Up to ``phrase_judge_book_workers`` books are judged at once, and within each book candidates
are judged in concurrent chunks of ``phrase_judge_phrase_workers``. Each book uses its own
short-lived sessions for reads and writes; no database connection is held while LLM calls are
in flight. For a pseudo-single-threaded run (solo testing, debugging), set both worker
settings to 1.
Args:
session (Session): Active database session.
config (EbookSearchConfig): Runtime phrase-tuning settings.
engine (AsyncEngine): Engine used to open one session per book.
config (EbookSearchConfig): Runtime phrase-tuning settings and chat configuration.
source_ids (Sequence[int] | None): Books to judge; ``None`` judges every indexed book.
Returns:
PhraseJudgmentBackfillResult: Per-corpus counts of books judged, failures, candidates,
protected phrases, and mentions.
"""
source_ids = session.scalars(select(EbookSource.id).order_by(EbookSource.id)).all()
if source_ids is None:
async with AsyncSession(engine) as session:
source_ids = list((await session.scalars(select(EbookSource.id).order_by(EbookSource.id))).all())
books_seen = len(source_ids)
book_workers = max(1, config.phrase_judge_book_workers)
phrase_workers = max(1, config.phrase_judge_phrase_workers)
logger.info(
"ebook_candidate_phrase_judgment_start books_seen=%s llm_candidates_per_book=%s "
"target_protected_per_book=%s confidence_threshold=%.2f min_tokens=%s min_uses=%s",
"ebook_candidate_phrase_judgment_start books_seen=%s book_workers=%s phrase_workers=%s "
"confidence_threshold=%.2f",
books_seen,
config.protected_phrase_llm_candidates_per_book,
config.phrase_target_protected_per_book,
book_workers,
phrase_workers,
config.protected_phrase_confidence_threshold,
config.phrase_min_tokens,
minimum_candidate_raw_count(config),
)
outcomes = [judge_book_for_backfill(session, source_id, config) for source_id in source_ids]
book_semaphore = asyncio.Semaphore(book_workers)
max_connections = book_workers * phrase_workers
limits = httpx.Limits(max_connections=max_connections, max_keepalive_connections=max_connections)
async with httpx.AsyncClient(limits=limits) as client:
outcomes = await asyncio.gather(
*(judge_one_book_async(engine, source_id, config, client, book_semaphore) for source_id in source_ids)
)
result = PhraseJudgmentBackfillResult(
books_seen=books_seen,
@@ -95,235 +117,247 @@ def judge_candidate_phrases_for_books(
return result
def judge_book_for_backfill(
session: Session,
async def judge_one_book_async(
engine: AsyncEngine,
source_id: int,
config: EbookSearchConfig,
client: httpx.AsyncClient,
book_semaphore: asyncio.Semaphore,
) -> BookJudgmentResult:
"""Judge one book's candidates and return its outcome.
"""Judge one book concurrently and persist the outcome, honoring the book-level limit.
Args:
session (Session): Active database session.
engine (AsyncEngine): Engine used to open the book's read and write sessions.
source_id (int): Book to judge candidates for.
config (EbookSearchConfig): Runtime phrase-tuning settings.
client (httpx.AsyncClient): Shared async client for LLM calls.
book_semaphore (asyncio.Semaphore): Caps how many books judge at once.
Returns:
BookJudgmentResult: The book's judgment outcome, or an empty result when nothing was unjudged.
BookJudgmentResult: The book's judgment outcome.
"""
unjudged_count = count_unjudged_candidates(session, source_id, config)
if not unjudged_count:
logger.info(
"ebook_candidate_phrase_judgment_book_skip_no_unjudged source_id=%s",
source_id,
)
return BookJudgmentResult()
logger.info(
"ebook_candidate_phrase_judgment_book_start source_id=%s unjudged=%s",
source_id,
unjudged_count,
)
return run_book_judgment(session, source_id, unjudged_count, config)
async with book_semaphore:
try:
prepared = await prepare_book_judgment(engine, source_id, config)
if prepared is None:
return BookJudgmentResult()
work_items, target_remaining = prepared
judged = await judge_book_candidates_async(client, config, source_id, work_items, target_remaining)
if not judged:
return BookJudgmentResult()
return await persist_book_judgments(engine, source_id, config, judged)
except Exception:
logger.exception("ebook_candidate_phrase_judgment_book_failed source_id=%s", source_id)
return BookJudgmentResult(failed=True)
def run_book_judgment(
session: Session,
async def prepare_book_judgment(
engine: AsyncEngine,
source_id: int,
unjudged_count: int,
config: EbookSearchConfig,
) -> BookJudgmentResult:
"""Judge and index one book's candidates, managing its own transaction.
Commits on success and rolls back on error, returning a :class:`BookJudgmentResult`
that describes the outcome rather than raising it to the caller.
) -> tuple[list[tuple[int, PhraseCandidate]], int | None] | None:
"""Load one book's candidates to judge, with sample contexts, on a short-lived read session.
Args:
session (Session): Active database session.
source_id (int): Book to judge candidates for.
unjudged_count (int): Storable unjudged candidates counted before judging, for logging.
engine (AsyncEngine): Engine used to open the read session.
source_id (int): Book to load candidates for.
config (EbookSearchConfig): Runtime phrase-tuning settings.
Returns:
BookJudgmentResult: Judged, protected, and mention counts with commit and failure flags.
"""
book_started_at = perf_counter()
try:
book_text = load_book_text(session, source_id)
if not book_text:
logger.warning("ebook_candidate_phrase_judgment_book_empty source_id=%s", source_id)
return BookJudgmentResult()
judged, protected = judge_candidate_phrases_for_book(
session,
source_id,
series_id=None,
normalized_book_text=normalize_text(book_text),
config=config,
)
if judged == 0:
logger.info(
"ebook_candidate_phrase_judgment_book_skip_no_judgments source_id=%s unjudged=%s",
source_id,
unjudged_count,
)
return BookJudgmentResult()
mentions = index_chunk_phrase_mentions_for_book(session, source_id, config) if protected else 0
session.commit()
except Exception:
session.rollback()
logger.exception("ebook_candidate_phrase_judgment_book_failed source_id=%s", source_id)
return BookJudgmentResult(failed=True)
logger.info(
"ebook_candidate_phrase_judgment_book_committed source_id=%s judged=%s protected=%s mentions=%s "
"duration_ms=%.1f",
source_id,
judged,
len(protected),
mentions,
(perf_counter() - book_started_at) * 1000,
)
return BookJudgmentResult(judged=judged, protected=len(protected), mentions=mentions, committed=True)
def judge_candidate_phrases_for_book(
session: Session,
book_id: int,
series_id: int | None,
normalized_book_text: str,
config: EbookSearchConfig,
) -> tuple[int, list[EbookProtectedPhrase]]:
"""Judge unjudged candidate phrase rows for one book.
Args:
session (Session): Active database session.
book_id (int): Book whose candidates are judged.
series_id (int | None): Series scope for promoted protected phrases.
normalized_book_text (str): Whole book text, already normalized, for context lookups.
config (EbookSearchConfig): Runtime phrase-tuning settings.
Returns:
tuple[int, list[EbookProtectedPhrase]]: Number of candidates judged and the promoted phrases.
tuple[list[tuple[int, PhraseCandidate]], int | None] | None: Candidate rows paired with
in-memory candidates and the remaining protected-phrase target, or ``None`` when the book
has nothing to judge.
"""
judgment_limit = config.protected_phrase_llm_candidates_per_book
if judgment_limit <= 0:
logger.info("ebook_candidate_phrase_judgment_skipped_llm_limit_zero book_id=%s", book_id)
return 0, []
existing_protected = count_protected_phrases(session, book_id)
target_remaining: int | None = None
if config.phrase_target_protected_per_book > 0:
target_remaining = max(config.phrase_target_protected_per_book - existing_protected, 0)
if target_remaining == 0:
logger.info(
"ebook_candidate_phrase_judgment_skipped_target_met book_id=%s existing_protected=%s target=%s",
book_id,
existing_protected,
config.phrase_target_protected_per_book,
return None
async with AsyncSession(engine) as session:
if not await count_unjudged_candidates(session, source_id, config):
logger.info("ebook_candidate_phrase_judgment_book_skip_no_unjudged source_id=%s", source_id)
return None
existing_protected = await count_protected_phrases(session, source_id)
target_remaining: int | None = None
if config.phrase_target_protected_per_book > 0:
target_remaining = max(config.phrase_target_protected_per_book - existing_protected, 0)
if target_remaining == 0:
logger.info(
"ebook_candidate_phrase_judgment_skipped_target_met source_id=%s existing_protected=%s target=%s",
source_id,
existing_protected,
config.phrase_target_protected_per_book,
)
return None
book_text = await load_book_text(session, source_id)
if not book_text:
logger.warning("ebook_candidate_phrase_judgment_book_empty source_id=%s", source_id)
return None
normalized_book_text = normalize_text(book_text)
# Stored rows may predate the current junk filters and score weights, so re-filter and
# rescore every unjudged row here instead of trusting the persisted candidate_score.
rows = await load_candidates_for_judgment(session, source_id, config)
scored_items: list[tuple[int, PhraseCandidate]] = []
skipped_junk = 0
for row in rows:
candidate = phrase_candidate_from_row(row)
if is_junk_phrase(candidate.phrase_norm.split()):
skipped_junk += 1
continue
candidate.candidate_score = score_candidate(candidate, config)
scored_items.append((row.id, candidate))
scored_items.sort(key=lambda item: item[1].candidate_score, reverse=True)
work_items = scored_items[:judgment_limit]
for _, candidate in work_items:
candidate.sample_contexts = candidate.sample_contexts or get_sample_contexts(
normalized_book_text, candidate.phrase_norm
)
return 0, []
rows = load_candidates_for_judgment(session, book_id, judgment_limit, config)
logger.info(
"ebook_candidate_phrase_judgment_candidates_loaded book_id=%s candidates=%s existing_protected=%s "
"target_remaining=%s judgment_limit=%s",
book_id,
"ebook_candidate_phrase_judgment_candidates_loaded source_id=%s candidates=%s skipped_junk=%s "
"unjudged_rows=%s existing_protected=%s target_remaining=%s judgment_limit=%s",
source_id,
len(work_items),
skipped_junk,
len(rows),
existing_protected,
target_remaining,
judgment_limit,
)
judged_count = 0
protected: list[EbookProtectedPhrase] = []
for row_number, row in enumerate(rows, start=1):
candidate, judgment, candidate_row = judge_candidate_row(
session, book_id, series_id, normalized_book_text, row, row_number, len(rows), config
)
judged_count += 1
if not should_protect_judged_candidate(row, candidate, judgment, book_id, config):
continue
protected.append(upsert_protected_phrase(session, book_id, series_id, candidate, judgment, candidate_row))
if target_remaining is not None and len(protected) >= target_remaining:
break
session.flush()
return judged_count, protected
return work_items, target_remaining
def judge_candidate_row(
session: Session,
book_id: int,
series_id: int | None,
normalized_book_text: str,
row: EbookCandidatePhrase,
row_number: int,
total_rows: int,
async def judge_book_candidates_async(
client: httpx.AsyncClient,
config: EbookSearchConfig,
) -> tuple[PhraseCandidate, LLMJudgment, EbookCandidatePhrase]:
"""Run and persist the LLM judgment for a single candidate row.
source_id: int,
work_items: list[tuple[int, PhraseCandidate]],
target_remaining: int | None,
) -> list[tuple[int, PhraseCandidate, LLMJudgment, bool]]:
"""Judge a book's candidates in concurrent chunks, stopping once the target is reached.
Promotion decisions are made in memory so judging can stop early without any database writes.
Args:
session (Session): Active database session.
book_id (int): Book the candidate belongs to.
series_id (int | None): Series scope for the saved candidate.
normalized_book_text (str): Whole book text, already normalized, for context lookups.
row (EbookCandidatePhrase): Stored candidate row to judge.
row_number (int): 1-based position of the row in the batch, for logging.
total_rows (int): Total rows in the batch, for logging.
client (httpx.AsyncClient): Shared async client for LLM calls.
config (EbookSearchConfig): Runtime phrase-tuning settings.
source_id (int): Book being judged, for logging.
work_items (list[tuple[int, PhraseCandidate]]): Candidate row ids paired with candidates,
in best-first score order.
target_remaining (int | None): Remaining protected-phrase target, or ``None`` for no cap.
Returns:
tuple[PhraseCandidate, LLMJudgment, EbookCandidatePhrase]: The candidate, its judgment,
and the persisted candidate row.
list[tuple[int, PhraseCandidate, LLMJudgment, bool]]: Judged rows with their judgment and
whether each should be promoted.
"""
row_started_at = perf_counter()
candidate = phrase_candidate_from_row(row)
candidate.sample_contexts = row.sample_contexts or get_sample_contexts(normalized_book_text, candidate.phrase_norm)
chunk_size = max(1, config.phrase_judge_phrase_workers)
judged: list[tuple[int, PhraseCandidate, LLMJudgment, bool]] = []
promoted = 0
for start in range(0, len(work_items), chunk_size):
chunk = work_items[start : start + chunk_size]
judgments = await asyncio.gather(*(judge_candidate_async(client, config, candidate) for _, candidate in chunk))
for (candidate_id, candidate), judgment in zip(chunk, judgments, strict=True):
promote = (target_remaining is None or promoted < target_remaining) and should_protect_judged_candidate(
candidate, judgment, source_id, config, candidate_id=candidate_id
)
if promote:
promoted += 1
judged.append((candidate_id, candidate, judgment, promote))
if target_remaining is not None and promoted >= target_remaining:
break
return judged
async def judge_candidate_async(
client: httpx.AsyncClient,
config: EbookSearchConfig,
candidate: PhraseCandidate,
) -> LLMJudgment:
"""Judge one candidate with the LLM over the shared async client.
Args:
client (httpx.AsyncClient): Shared async client for LLM calls.
config (EbookSearchConfig): Runtime phrase-tuning settings.
candidate (PhraseCandidate): Candidate to judge.
Returns:
LLMJudgment: The parsed judgment.
"""
content = await request_chat_completion(client, config, build_judge_messages(candidate))
return parse_llm_judgment(content, config)
async def persist_book_judgments(
engine: AsyncEngine,
source_id: int,
config: EbookSearchConfig,
judged: list[tuple[int, PhraseCandidate, LLMJudgment, bool]],
) -> BookJudgmentResult:
"""Persist one book's judgments and promotions in a single committed transaction.
Args:
engine (AsyncEngine): Engine used to open the write session.
source_id (int): Book being persisted.
config (EbookSearchConfig): Runtime phrase-tuning settings.
judged (list[tuple[int, PhraseCandidate, LLMJudgment, bool]]): Judged candidates with their
judgment and promotion flag.
Returns:
BookJudgmentResult: The book's committed counts, or a failed result on error.
"""
book_started_at = perf_counter()
async with AsyncSession(engine, expire_on_commit=False) as session:
try:
protected: list[EbookProtectedPhrase] = []
for candidate_id, candidate, judgment, promote in judged:
candidate_row = await save_candidate_to_db(session, source_id, None, candidate, judgment=judgment)
if promote:
protected.append(
await upsert_protected_phrase(session, source_id, None, candidate, judgment, candidate_row)
)
logger.info(
"ebook_candidate_phrase_judgment_candidate_complete source_id=%s candidate_id=%s phrase=%r "
"keep=%s confidence=%.3f category=%r promoted=%s",
source_id,
candidate_id,
candidate.phrase_norm,
judgment.keep,
judgment.confidence,
judgment.category,
promote,
)
await session.flush()
mentions = await index_chunk_phrase_mentions_for_book(session, source_id, config) if protected else 0
await session.commit()
except Exception:
await session.rollback()
logger.exception("ebook_candidate_phrase_judgment_book_persist_failed source_id=%s", source_id)
return BookJudgmentResult(failed=True)
logger.info(
"ebook_candidate_phrase_judgment_candidate_start book_id=%s candidate_id=%s row_number=%s rows=%s "
"phrase=%r score=%.3f raw_count=%s chapter_count=%s",
book_id,
row.id,
row_number,
total_rows,
candidate.phrase_norm,
candidate.candidate_score,
candidate.raw_count,
candidate.chapter_count,
"ebook_candidate_phrase_judgment_book_committed source_id=%s judged=%s protected=%s mentions=%s "
"duration_ms=%.1f",
source_id,
len(judged),
len(protected),
mentions,
(perf_counter() - book_started_at) * 1000,
)
judgment = judge_candidate_with_llm(candidate, config)
candidate_row = save_candidate_to_db(session, book_id, series_id, candidate, judgment=judgment)
logger.info(
"ebook_candidate_phrase_judgment_candidate_complete book_id=%s candidate_id=%s phrase=%r keep=%s "
"confidence=%.3f importance=%.3f category=%r duration_ms=%.1f",
book_id,
row.id,
candidate.phrase_norm,
judgment.keep,
judgment.confidence,
judgment.importance,
judgment.category,
(perf_counter() - row_started_at) * 1000,
)
return candidate, judgment, candidate_row
return BookJudgmentResult(judged=len(judged), protected=len(protected), mentions=mentions, committed=True)
def should_protect_judged_candidate(
row: EbookCandidatePhrase,
candidate: PhraseCandidate,
judgment: LLMJudgment,
book_id: int,
config: EbookSearchConfig,
*,
candidate_id: int,
) -> bool:
"""Report whether a judged candidate qualifies to become a protected phrase.
Args:
row (EbookCandidatePhrase): Stored candidate row the judgment came from.
candidate (PhraseCandidate): In-memory candidate that was judged.
judgment (LLMJudgment): Judge decision for the candidate.
book_id (int): Book the candidate belongs to, for logging.
config (EbookSearchConfig): Runtime phrase-tuning settings.
candidate_id (int): Stored candidate row id the judgment came from, for logging.
Returns:
bool: True when the judged candidate should be promoted to a protected phrase.
@@ -331,25 +365,26 @@ def should_protect_judged_candidate(
if not judgment.keep or judgment.confidence < config.protected_phrase_confidence_threshold:
return False
accepted_norm = normalize_text(judgment.canonical or candidate.phrase_text)
accepted_token_count = len(accepted_norm.split())
accepted_tokens = accepted_norm.split()
accepted_token_count = len(accepted_tokens)
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_id,
candidate.phrase_norm,
accepted_norm,
accepted_token_count,
config.phrase_min_tokens,
)
return False
if is_most_common_word_phrase(accepted_norm):
if is_most_common_word_phrase(accepted_tokens):
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_id,
candidate.phrase_norm,
accepted_norm,
)
@@ -357,18 +392,14 @@ def should_protect_judged_candidate(
return True
"""LLM judging of extracted candidate phrases."""
def judge_candidate_with_llm(candidate: PhraseCandidate, config: EbookSearchConfig) -> LLMJudgment:
"""Ask the configured chat model to judge one pre-extracted candidate.
def build_judge_messages(candidate: PhraseCandidate) -> list[dict[str, str]]:
"""Build the chat messages used to judge one candidate phrase.
Args:
candidate (PhraseCandidate): Candidate to send to the LLM judge.
config (EbookSearchConfig): Runtime phrase-tuning settings and chat configuration.
candidate (PhraseCandidate): Candidate to describe for the judge.
Returns:
LLMJudgment: The parsed structured judgment for the candidate.
list[dict[str, str]]: OpenAI-style system and user messages.
"""
payload = {
"phrase": candidate.phrase_norm,
@@ -378,7 +409,7 @@ def judge_candidate_with_llm(candidate: PhraseCandidate, config: EbookSearchConf
"chapter_count": candidate.chapter_count,
"contexts": candidate.sample_contexts,
}
messages = [
return [
{
"role": "system",
"content": (
@@ -392,7 +423,6 @@ def judge_candidate_with_llm(candidate: PhraseCandidate, config: EbookSearchConf
},
{"role": "user", "content": json.dumps(payload, ensure_ascii=True)},
]
return parse_llm_judgment(request_chat_completion(config, messages), config)
def parse_llm_judgment(content: str, config: EbookSearchConfig) -> LLMJudgment:
python/ebook_search/protected_phrases/matching.py
+37 -33
View File
@@ -26,7 +26,7 @@ from python.orm.richie import (
if TYPE_CHECKING:
from collections.abc import Iterator, Sequence
from sqlalchemy.orm import Session
from sqlalchemy.ext.asyncio import AsyncSession
from python.ebook_search.config import EbookSearchConfig
from python.ebook_search.protected_phrases.text_normalization import NormalizedToken
@@ -34,8 +34,8 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
def load_phrase_lookup(
session: Session,
async def load_phrase_lookup(
session: AsyncSession,
config: EbookSearchConfig,
*,
book_id: int | None = None,
@@ -44,7 +44,7 @@ def load_phrase_lookup(
"""Load protected phrases and aliases into RAM lookup maps.
Args:
session (Session): Active database session.
session (AsyncSession): Active database session.
config (EbookSearchConfig): Runtime phrase-tuning settings.
book_id (int | None): Optional book scope to restrict loaded phrases.
series_id (int | None): Optional series scope to restrict loaded phrases.
@@ -65,7 +65,7 @@ def load_phrase_lookup(
if scope_filter is not None:
phrase_statement = phrase_statement.where(scope_filter)
for row in session.execute(phrase_statement):
for row in await session.execute(phrase_statement):
phrase_id = int(row.id)
phrase_norm = str(row.phrase_norm)
norm_to_ids[phrase_norm].append(phrase_id)
@@ -78,7 +78,7 @@ def load_phrase_lookup(
if scope_filter is not None:
alias_statement = alias_statement.where(scope_filter)
for row in session.execute(alias_statement):
for row in await 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()))
@@ -197,11 +197,11 @@ def detect_phrase_candidates_from_tokens(tokens_: Sequence[NormalizedToken], loo
return matches
def hydrate_matches(session: Session, matches: Sequence[PhraseMatch]) -> list[HydratedPhraseMatch]:
async def hydrate_matches(session: AsyncSession, matches: Sequence[PhraseMatch]) -> list[HydratedPhraseMatch]:
"""Fetch protected phrase metadata for raw phrase matches.
Args:
session (Session): Active database session.
session (AsyncSession): Active database session.
matches (Sequence[PhraseMatch]): Unhydrated matches to enrich.
Returns:
@@ -216,7 +216,7 @@ def hydrate_matches(session: Session, matches: Sequence[PhraseMatch]) -> list[Hy
rows = {
row.id: row
for row in session.scalars(select(EbookProtectedPhrase).where(EbookProtectedPhrase.id.in_(phrase_ids)))
for row in await session.scalars(select(EbookProtectedPhrase).where(EbookProtectedPhrase.id.in_(phrase_ids)))
}
hydrated: list[HydratedPhraseMatch] = []
for match in matches:
@@ -332,8 +332,8 @@ def resolve_overlaps(matches: Sequence[HydratedPhraseMatch]) -> list[HydratedPhr
return kept
def detect_protected_phrases_for_query(
session: Session,
async def detect_protected_phrases_for_query(
session: AsyncSession,
query_text: str,
config: EbookSearchConfig,
*,
@@ -344,7 +344,7 @@ def detect_protected_phrases_for_query(
"""Run the full online protected-phrase query-detection pipeline.
Args:
session (Session): Active database session.
session (AsyncSession): Active database session.
query_text (str): User query text to detect phrases in.
config (EbookSearchConfig): Runtime phrase-tuning settings.
lookup (PhraseLookup | None): Optional preloaded lookup; loaded on demand when ``None``.
@@ -355,13 +355,15 @@ def detect_protected_phrases_for_query(
list[HydratedPhraseMatch]: Hydrated, overlap-resolved phrase matches for the query.
"""
active_lookup = (
lookup if lookup is not None else load_phrase_lookup(session, config, book_id=book_id, series_id=series_id)
lookup
if lookup is not None
else await 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)))
return resolve_overlaps(await hydrate_matches(session, detect_phrase_candidates(query_text, active_lookup)))
def index_chunk_phrase_mentions_for_book(
session: Session,
async def index_chunk_phrase_mentions_for_book(
session: AsyncSession,
book_id: int,
config: EbookSearchConfig,
*,
@@ -371,7 +373,7 @@ def index_chunk_phrase_mentions_for_book(
"""Rebuild chunk phrase mentions for all chunks in one book.
Args:
session (Session): Active database session.
session (AsyncSession): Active database session.
book_id (int): Book whose chunk mentions are rebuilt.
config (EbookSearchConfig): Runtime phrase-tuning settings.
series_id (int | None): Optional series scope for lookup loading.
@@ -381,23 +383,25 @@ def index_chunk_phrase_mentions_for_book(
int: Total number of chunk phrase mentions indexed for the book.
"""
active_lookup = (
lookup if lookup is not None else load_phrase_lookup(session, config, book_id=book_id, series_id=series_id)
lookup
if lookup is not None
else await 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))
await session.execute(delete(EbookChunkPhraseMention).where(EbookChunkPhraseMention.book_id == book_id))
chunks = await 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()
count += await index_chunk_phrase_mentions(session, chunk, lookup=active_lookup)
await 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:
async def index_chunk_phrase_mentions(session: AsyncSession, chunk: EbookChunk, *, lookup: PhraseLookup) -> int:
"""Store protected phrase mentions for one chunk.
Args:
session (Session): Active database session.
session (AsyncSession): Active database session.
chunk (EbookChunk): Chunk whose text is scanned for phrase mentions.
lookup (PhraseLookup): In-memory phrase and alias lookup maps.
@@ -405,7 +409,7 @@ def index_chunk_phrase_mentions(session: Session, chunk: EbookChunk, *, lookup:
int: Number of phrase mentions stored for the chunk.
"""
raw_matches = detect_phrase_candidates_in_text(chunk.text, lookup)
hydrated = resolve_overlaps(hydrate_matches(session, raw_matches))
hydrated = resolve_overlaps(await hydrate_matches(session, raw_matches))
for match in hydrated:
session.add(
EbookChunkPhraseMention(
@@ -420,8 +424,8 @@ def index_chunk_phrase_mentions(session: Session, chunk: EbookChunk, *, lookup:
return len(hydrated)
def phrase_hits_for_chunks(
session: Session,
async def phrase_hits_for_chunks(
session: AsyncSession,
*,
chunk_ids: Sequence[int],
phrase_ids: Sequence[int],
@@ -429,7 +433,7 @@ def phrase_hits_for_chunks(
"""Return matched protected phrases with mention counts by chunk id using indexed chunk mentions.
Args:
session (Session): Active database session.
session (AsyncSession): Active database session.
chunk_ids (Sequence[int]): Chunk ids to look up mentions for.
phrase_ids (Sequence[int]): Protected phrase ids to restrict the results to.
@@ -456,7 +460,7 @@ def phrase_hits_for_chunks(
.order_by(EbookChunkPhraseMention.chunk_id, mention_count.desc(), EbookProtectedPhrase.phrase_text)
)
hits: defaultdict[int, list[ChunkPhraseHit]] = defaultdict(list)
for row in session.execute(statement):
for row in await session.execute(statement):
hits[int(row.chunk_id)].append(
ChunkPhraseHit(
phrase_id=int(row.phrase_id),
@@ -467,8 +471,8 @@ def phrase_hits_for_chunks(
return {chunk_id: tuple(chunk_hits) for chunk_id, chunk_hits in hits.items()}
def phrase_hit_counts_for_chunks(
session: Session,
async def phrase_hit_counts_for_chunks(
session: AsyncSession,
*,
chunk_ids: Sequence[int],
phrase_ids: Sequence[int],
@@ -476,12 +480,12 @@ def phrase_hit_counts_for_chunks(
"""Return phrase-hit counts by chunk id using indexed chunk mentions.
Args:
session (Session): Active database session.
session (AsyncSession): Active database session.
chunk_ids (Sequence[int]): Chunk ids to count mentions for.
phrase_ids (Sequence[int]): Protected phrase ids to restrict the counts to.
Returns:
dict[int, int]: Total mention count per chunk id.
"""
hits = phrase_hits_for_chunks(session, chunk_ids=chunk_ids, phrase_ids=phrase_ids)
hits = await phrase_hits_for_chunks(session, chunk_ids=chunk_ids, phrase_ids=phrase_ids)
return {chunk_id: sum(hit.mention_count for hit in chunk_hits) for chunk_id, chunk_hits in hits.items()}
python/ebook_search/protected_phrases/models.py
+23
View File
@@ -188,6 +188,29 @@ class PhraseCandidateGenerationResult:
candidate_phrases: int
@dataclass(frozen=True, slots=True)
class CorpusPhraseStats:
"""Corpus-wide candidate and protected phrase counts for the admin page.
Attributes:
total_books (int): Indexed books in the corpus.
books_with_candidates (int): Books that have candidate phrases generated.
books_fully_judged (int): Books with candidates where every candidate has been judged.
candidate_phrases (int): Candidate phrases stored across all books.
judged_candidates (int): Candidate phrases that have been LLM judged.
unjudged_candidates (int): Candidate phrases still waiting for judgment.
protected_phrases (int): Protected phrases promoted across all books.
"""
total_books: int
books_with_candidates: int
books_fully_judged: int
candidate_phrases: int
judged_candidates: int
unjudged_candidates: int
protected_phrases: int
@dataclass(frozen=True, slots=True)
class PhraseJudgmentBackfillResult:
"""Summary of LLM judging for stored candidate phrases.
python/ebook_search/protected_phrases/pool.py
+101
View File
@@ -0,0 +1,101 @@
"""Process pool for offloading CPU-bound phrase extraction off the request thread.
Phrase extraction is pure-Python CPU work (n-gram sliding, YAKE), so running it inline in a
sync request handler serializes concurrent recalculations behind the GIL. Submitting it to a
``ProcessPoolExecutor`` lets concurrent extractions run in parallel across cores instead. A
``spawn`` context is used so workers do not inherit the parent's database engine, connections,
or server threads.
"""
from __future__ import annotations
import asyncio
import logging
import multiprocessing
import os
from concurrent.futures import ProcessPoolExecutor
from threading import Lock
from typing import TYPE_CHECKING
from python.ebook_search.protected_phrases.extraction import extract_phrase_candidates_for_book
if TYPE_CHECKING:
from collections.abc import Mapping, Sequence
from python.ebook_search.config import EbookSearchConfig
from python.ebook_search.protected_phrases.models import PhraseCandidate
logger = logging.getLogger(__name__)
class _ExtractionPool:
"""Lazily created process-wide extraction pool and the lock guarding it."""
def __init__(self) -> None:
self.lock = Lock()
self.pool: ProcessPoolExecutor | None = None
_extraction_pool = _ExtractionPool()
def get_extraction_pool(max_workers: int) -> ProcessPoolExecutor:
"""Return the shared extraction process pool, creating it on first use.
Args:
max_workers (int): Desired worker count; values below 1 fall back to the CPU count.
Returns:
ProcessPoolExecutor: The shared pool for phrase extraction.
"""
with _extraction_pool.lock:
if _extraction_pool.pool is None:
workers = max_workers if max_workers > 0 else (os.cpu_count() or 1)
_extraction_pool.pool = ProcessPoolExecutor(
max_workers=workers,
mp_context=multiprocessing.get_context("spawn"),
)
logger.info("ebook_phrase_extraction_pool_started workers=%s", workers)
return _extraction_pool.pool
def shutdown_extraction_pool() -> None:
"""Shut down the shared extraction pool if it was started."""
with _extraction_pool.lock:
if _extraction_pool.pool is not None:
_extraction_pool.pool.shutdown(wait=False, cancel_futures=True)
_extraction_pool.pool = None
logger.info("ebook_phrase_extraction_pool_shutdown")
async def extract_phrase_candidates_in_pool(
book_text: str,
chapters: Sequence[str],
config: EbookSearchConfig,
*,
metadata: Mapping[str, object] | None,
) -> list[PhraseCandidate]:
"""Run book phrase extraction in a worker process and await the result.
Only the CPU-bound extraction runs in the worker; the caller keeps all database work in the
request process. The spaCy pipeline is not supported here because it is not picklable, so
this always runs the non-spaCy extraction path.
Args:
book_text (str): Full book text used for extraction.
chapters (Sequence[str]): Chapter-like text blocks used for frequency counts.
config (EbookSearchConfig): Runtime phrase-tuning settings.
metadata (Mapping[str, object] | None): Optional book metadata used as a candidate source.
Returns:
list[PhraseCandidate]: Scored candidates sorted best-first and capped per book.
"""
pool = get_extraction_pool(config.protected_phrase_extraction_workers)
future = pool.submit(
extract_phrase_candidates_for_book,
book_text,
list(chapters),
config,
metadata=dict(metadata) if metadata is not None else None,
)
return await asyncio.wrap_future(future)
python/ebook_search/protected_phrases/store.py
+244 -72
View File
@@ -11,7 +11,11 @@ from sqlalchemy.dialects.postgresql import insert as pg_insert
from sqlalchemy.dialects.sqlite import insert as sqlite_insert
from python.ebook_search.protected_phrases.extraction import minimum_candidate_raw_count
from python.ebook_search.protected_phrases.models import PhraseCandidate, PhraseRecalculationResult
from python.ebook_search.protected_phrases.models import (
CorpusPhraseStats,
PhraseCandidate,
PhraseRecalculationResult,
)
from python.ebook_search.protected_phrases.text_normalization import normalize_text
from python.orm.richie import (
EbookCandidatePhrase,
@@ -27,7 +31,7 @@ if TYPE_CHECKING:
from sqlalchemy.dialects.postgresql.dml import Insert as PostgresInsert
from sqlalchemy.dialects.sqlite.dml import Insert as SqliteInsert
from sqlalchemy.orm import Session
from sqlalchemy.ext.asyncio import AsyncSession
from python.ebook_search.config import EbookSearchConfig
from python.ebook_search.protected_phrases.models import LLMJudgment
@@ -36,14 +40,14 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
def dialect_insert(session: Session, table: type[TableBase]) -> PostgresInsert | SqliteInsert:
def dialect_insert(session: AsyncSession, table: type[TableBase]) -> PostgresInsert | SqliteInsert:
"""Return a dialect-specific INSERT construct that supports ``ON CONFLICT DO UPDATE``.
Production runs on PostgreSQL while tests run on SQLite; both support upserts with
compatible SQLAlchemy constructs, so the correct one is chosen from the bound dialect.
Args:
session (Session): Active database session whose bind selects the dialect.
session (AsyncSession): Active database session whose bind selects the dialect.
table (type[TableBase]): Mapped table to insert into.
Returns:
@@ -54,35 +58,33 @@ def dialect_insert(session: Session, table: type[TableBase]) -> PostgresInsert |
return pg_insert(table)
def load_book_text(session: Session, book_id: int) -> str:
async def load_book_text(session: AsyncSession, book_id: int) -> str:
"""Load a book's indexed chunk text as one string for phrase extraction.
Args:
session (Session): Active database session.
session (AsyncSession): Active database session.
book_id (int): Book whose chunk text is loaded.
Returns:
str: The book's chunk text joined into a single string.
"""
texts = session.scalars(
texts = await session.scalars(
select(EbookChunk.text).where(EbookChunk.source_id == book_id).order_by(EbookChunk.chunk_index)
)
return "\n\n".join(stripped for text in texts if (stripped := text.strip()))
def load_book_chapter_texts(session: Session, book_id: int) -> list[str]:
async def load_book_chapter_texts(session: AsyncSession, book_id: int) -> list[str]:
"""Reconstruct chapter-like text blocks from indexed chunks for phrase extraction.
Args:
session (Session): Active database session.
session (AsyncSession): Active database session.
book_id (int): Book whose chunks are grouped into chapters.
Returns:
list[str]: Non-empty chapter-like text blocks in chunk order.
"""
rows = session.execute(
rows = await session.execute(
select(EbookChunk.chapter_id, EbookChunk.text)
.where(EbookChunk.source_id == book_id)
.order_by(EbookChunk.chunk_index)
@@ -127,11 +129,11 @@ def metadata_for_source(source: EbookSource) -> dict[str, object | None]:
}
def metadata_for_source_id(session: Session, source_id: int) -> dict[str, object | None]:
async def metadata_for_source_id(session: AsyncSession, source_id: int) -> dict[str, object | None]:
"""Return phrase extraction metadata for one indexed source by id.
Args:
session (Session): Active database session.
session (AsyncSession): Active database session.
source_id (int): Id of the indexed source to read metadata from.
Returns:
@@ -140,70 +142,128 @@ def metadata_for_source_id(session: Session, source_id: int) -> dict[str, object
Raises:
ValueError: If no source exists with the given id.
"""
source = session.get(EbookSource, source_id)
source = await session.get(EbookSource, source_id)
if source is None:
msg = f"No indexed source with id {source_id}"
raise ValueError(msg)
return metadata_for_source(source)
def count_protected_phrases(session: Session, book_id: int) -> int:
async def count_protected_phrases(session: AsyncSession, book_id: int) -> int:
"""Count stored protected phrases for one book.
Args:
session (Session): Active database session.
session (AsyncSession): Active database session.
book_id (int): Book whose protected phrases are counted.
Returns:
int: Number of protected phrases stored for the book.
"""
return session.scalars(
select(func.count(EbookProtectedPhrase.id)).where(EbookProtectedPhrase.book_id == book_id)
return (
await session.scalars(
select(func.count(EbookProtectedPhrase.id)).where(EbookProtectedPhrase.book_id == book_id)
)
).one()
def count_unjudged_candidates(session: Session, book_id: int, config: EbookSearchConfig) -> int:
async def count_unjudged_candidates(session: AsyncSession, book_id: int, config: EbookSearchConfig) -> int:
"""Count storable candidate rows for a book that have not yet been judged.
Args:
session (Session): Active database session.
session (AsyncSession): Active database session.
book_id (int): Book whose unjudged candidates are counted.
config (EbookSearchConfig): Runtime phrase-tuning settings supplying storage thresholds.
Returns:
int: Number of storable, unjudged candidate rows for the book.
"""
return session.scalars(
select(func.count(EbookCandidatePhrase.id)).where(
EbookCandidatePhrase.book_id == book_id,
EbookCandidatePhrase.llm_judged.is_(False),
EbookCandidatePhrase.token_count >= config.phrase_min_tokens,
EbookCandidatePhrase.raw_count >= minimum_candidate_raw_count(config),
return (
await session.scalars(
select(func.count(EbookCandidatePhrase.id)).where(
EbookCandidatePhrase.book_id == book_id,
EbookCandidatePhrase.llm_judged.is_(False),
EbookCandidatePhrase.token_count >= config.phrase_min_tokens,
EbookCandidatePhrase.raw_count >= minimum_candidate_raw_count(config),
)
)
).one()
def load_candidates_for_judgment(
session: Session,
book_id: int,
judgment_limit: int,
config: EbookSearchConfig,
) -> Sequence[EbookCandidatePhrase]:
"""Load the top unjudged candidate rows for a book.
Common-word phrases are filtered before storage (``filter_storable_candidates``), so
no re-check is needed here.
async def corpus_phrase_stats(session: AsyncSession) -> CorpusPhraseStats:
"""Summarize candidate and protected phrase coverage across the whole corpus.
Args:
session (Session): Active database session.
session (AsyncSession): Active database session.
Returns:
CorpusPhraseStats: Corpus-wide phrase counts and per-book coverage counts.
"""
total_books = (await session.scalars(select(func.count(EbookSource.id)))).one()
candidate_phrases, judged_candidates, books_with_candidates, books_with_unjudged = (
await session.execute(
select(
func.count(EbookCandidatePhrase.id),
func.count(EbookCandidatePhrase.id).filter(EbookCandidatePhrase.llm_judged.is_(True)),
func.count(func.distinct(EbookCandidatePhrase.book_id)),
func.count(func.distinct(EbookCandidatePhrase.book_id)).filter(
EbookCandidatePhrase.llm_judged.is_(False)
),
)
)
).one()
protected_phrases = (await session.scalars(select(func.count(EbookProtectedPhrase.id)))).one()
return CorpusPhraseStats(
total_books=total_books,
books_with_candidates=books_with_candidates,
books_fully_judged=books_with_candidates - books_with_unjudged,
candidate_phrases=candidate_phrases,
judged_candidates=judged_candidates,
unjudged_candidates=candidate_phrases - judged_candidates,
protected_phrases=protected_phrases,
)
async def book_ids_pending_first_judgment(session: AsyncSession) -> list[int]:
"""Return books that have candidate phrases but no judged candidates yet.
Args:
session (AsyncSession): Active database session.
Returns:
list[int]: Book ids with candidates where judging has never run, ordered by id.
"""
judged_books = select(EbookCandidatePhrase.book_id).where(EbookCandidatePhrase.llm_judged.is_(True)).distinct()
return list(
(
await session.scalars(
select(EbookCandidatePhrase.book_id)
.where(EbookCandidatePhrase.book_id.not_in(judged_books))
.distinct()
.order_by(EbookCandidatePhrase.book_id)
)
).all()
)
async def load_candidates_for_judgment(
session: AsyncSession,
book_id: int,
config: EbookSearchConfig,
) -> Sequence[EbookCandidatePhrase]:
"""Load every storable unjudged candidate row for a book.
Rows may have been stored before the current junk filters and score weights existed, so
callers re-check :func:`is_junk_phrase` and rescore before selecting what to judge.
Args:
session (AsyncSession): Active database session.
book_id (int): Book whose candidates are loaded.
judgment_limit (int): Maximum number of candidate rows to return.
config (EbookSearchConfig): Runtime phrase-tuning settings supplying storage thresholds.
Returns:
Sequence[EbookCandidatePhrase]: Top storable, unjudged candidate rows ordered by score.
Sequence[EbookCandidatePhrase]: Storable, unjudged candidate rows ordered by stored score.
"""
return session.scalars(
query = (
select(EbookCandidatePhrase)
.where(
EbookCandidatePhrase.book_id == book_id,
@@ -216,8 +276,8 @@ def load_candidates_for_judgment(
EbookCandidatePhrase.raw_count.desc(),
EbookCandidatePhrase.id,
)
.limit(judgment_limit)
).all()
)
return (await session.scalars(query)).all()
def phrase_candidate_from_row(row: EbookCandidatePhrase) -> PhraseCandidate:
@@ -248,25 +308,23 @@ def phrase_candidate_from_row(row: EbookCandidatePhrase) -> PhraseCandidate:
)
def save_candidate_to_db(
session: Session,
def candidate_row_values(
book_id: int,
series_id: int | None,
candidate: PhraseCandidate,
*,
judgment: LLMJudgment | None,
) -> EbookCandidatePhrase:
"""Insert or update one candidate phrase row.
) -> dict[str, object]:
"""Build the column values for one candidate phrase upsert.
Args:
session (Session): Active database session.
book_id (int): Book the candidate belongs to.
series_id (int | None): Series scope stored on the row.
candidate (PhraseCandidate): Candidate whose fields are written to the row.
judgment (LLMJudgment | None): Judgment to record, or ``None`` to leave the row unjudged.
Returns:
EbookCandidatePhrase: The inserted or updated candidate row.
dict[str, object]: Column values keyed by column name.
"""
values: dict[str, object] = {
"book_id": book_id,
@@ -296,6 +354,30 @@ def save_candidate_to_db(
llm_category=judgment.category,
llm_reason=judgment.reason,
)
return values
async def save_candidate_to_db(
session: AsyncSession,
book_id: int,
series_id: int | None,
candidate: PhraseCandidate,
*,
judgment: LLMJudgment | None,
) -> EbookCandidatePhrase:
"""Insert or update one candidate phrase row.
Args:
session (AsyncSession): Active database session.
book_id (int): Book the candidate belongs to.
series_id (int | None): Series scope stored on the row.
candidate (PhraseCandidate): Candidate whose fields are written to the row.
judgment (LLMJudgment | None): Judgment to record, or ``None`` to leave the row unjudged.
Returns:
EbookCandidatePhrase: The inserted or updated candidate row.
"""
values = candidate_row_values(book_id, series_id, candidate, judgment=judgment)
# Preserve an existing judgment when this call is only refreshing candidate fields.
skip_update = {"book_id", "phrase_norm"}
@@ -306,11 +388,94 @@ def save_candidate_to_db(
index_elements=["book_id", "phrase_norm"],
set_={column: insert_statement.excluded[column] for column in values if column not in skip_update},
).returning(EbookCandidatePhrase)
return session.scalars(statement, execution_options={"populate_existing": True}).one()
return (await session.scalars(statement, execution_options={"populate_existing": True})).one()
def upsert_protected_phrase(
session: Session,
BULK_CANDIDATE_UPSERT_CHUNK = 1000
async def bulk_upsert_unjudged_candidates(
session: AsyncSession,
book_id: int,
series_id: int | None,
candidates: Sequence[PhraseCandidate],
) -> int:
"""Insert or update many freshly extracted candidate rows in chunked multi-row upserts.
Saving one row per statement costs one database round trip per candidate, which dominated
generation time for full books, so candidates are written ``BULK_CANDIDATE_UPSERT_CHUNK``
rows per statement instead. Existing judgments and sample contexts are never overwritten:
fresh extractions carry no contexts, and ``llm_judged`` plus the ``llm_*`` columns are left
out of the conflict update. Candidates must have unique ``phrase_norm`` values, as produced
by extraction, since one multi-row upsert cannot touch the same row twice.
Args:
session (Session): Active database session.
book_id (int): Book the candidates belong to.
series_id (int | None): Series scope stored on the rows.
candidates (Sequence[PhraseCandidate]): Freshly extracted candidates to persist.
Returns:
int: Number of candidate rows written.
"""
values = [
candidate_row_values(book_id, series_id, candidate, judgment=None)
for candidate in candidates
if not candidate.sample_contexts
]
if len(values) != len(candidates):
msg = "bulk_upsert_unjudged_candidates only accepts freshly extracted candidates without sample contexts"
raise ValueError(msg)
skip_update = {"book_id", "phrase_norm", "llm_judged"}
for chunk_start in range(0, len(values), BULK_CANDIDATE_UPSERT_CHUNK):
chunk = values[chunk_start : chunk_start + BULK_CANDIDATE_UPSERT_CHUNK]
insert_statement = dialect_insert(session, EbookCandidatePhrase).values(chunk)
statement = insert_statement.on_conflict_do_update(
index_elements=["book_id", "phrase_norm"],
set_={column: insert_statement.excluded[column] for column in chunk[0] if column not in skip_update},
)
await session.execute(statement)
return len(values)
def new_candidate_row(book_id: int, series_id: int | None, candidate: PhraseCandidate) -> EbookCandidatePhrase:
"""Build a fresh unjudged candidate row without checking for an existing one.
Unlike :func:`save_candidate_to_db`, this does no lookup, so it is only safe when the caller
guarantees there is no existing row for ``(book_id, candidate.phrase_norm)`` — for example
right after :func:`delete_phrase_data_for_book` has cleared the book.
Args:
book_id (int): Book the candidate belongs to.
series_id (int | None): Series scope stored on the row.
candidate (PhraseCandidate): Candidate whose fields are written to the row.
Returns:
EbookCandidatePhrase: A new, unattached candidate row.
"""
row = EbookCandidatePhrase(book_id=book_id, phrase_norm=candidate.phrase_norm)
row.llm_judged = False
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)
return row
async def upsert_protected_phrase(
session: AsyncSession,
book_id: int,
series_id: int | None,
candidate: PhraseCandidate,
@@ -320,7 +485,7 @@ def upsert_protected_phrase(
"""Insert or update one accepted protected phrase and its aliases.
Args:
session (Session): Active database session.
session (AsyncSession): Active database session.
book_id (int): Book the protected phrase belongs to.
series_id (int | None): Series scope stored on the phrase.
candidate (PhraseCandidate): Candidate the phrase was promoted from.
@@ -360,18 +525,22 @@ def upsert_protected_phrase(
column: insert_statement.excluded[column] for column in values if column not in {"book_id", "phrase_norm"}
},
).returning(EbookProtectedPhrase)
row = session.scalars(statement, execution_options={"populate_existing": True}).one()
row = (await session.scalars(statement, execution_options={"populate_existing": True})).one()
for alias_text in judgment.aliases:
upsert_phrase_alias(session, row, alias_text)
await upsert_phrase_alias(session, row, alias_text)
return row
def upsert_phrase_alias(session: Session, phrase: EbookProtectedPhrase, alias_text: str) -> EbookPhraseAlias | None:
async def upsert_phrase_alias(
session: AsyncSession,
phrase: EbookProtectedPhrase,
alias_text: str,
) -> EbookPhraseAlias | None:
"""Insert or update one protected phrase alias.
Args:
session (Session): Active database session.
session (AsyncSession): Active database session.
phrase (EbookProtectedPhrase): Protected phrase the alias points to.
alias_text (str): Alias surface form to store.
@@ -395,7 +564,7 @@ def upsert_phrase_alias(session: Session, phrase: EbookProtectedPhrase, alias_te
"confidence": insert_statement.excluded.confidence,
},
).returning(EbookPhraseAlias)
return session.scalars(statement, execution_options={"populate_existing": True}).one()
return (await session.scalars(statement, execution_options={"populate_existing": True})).one()
def make_canonical_id(judgment: LLMJudgment, phrase_norm: str) -> str:
@@ -426,15 +595,15 @@ def slugify_identifier(value: str) -> str:
return slug.strip("_") or "unknown"
def prune_unstorable_unjudged_candidate_phrases(
session: Session,
async def prune_unstorable_unjudged_candidate_phrases(
session: AsyncSession,
book_id: int,
config: EbookSearchConfig,
) -> int:
"""Delete old unjudged candidate rows that no longer satisfy storage filters.
Args:
session (Session): Active database session.
session (AsyncSession): Active database session.
book_id (int): Book whose stale candidates are pruned.
config (EbookSearchConfig): Runtime phrase-tuning settings supplying storage thresholds.
@@ -442,7 +611,7 @@ def prune_unstorable_unjudged_candidate_phrases(
int: Number of candidate rows deleted.
"""
deleted = rowcount(
session.execute(
await session.execute(
delete(EbookCandidatePhrase).where(
EbookCandidatePhrase.book_id == book_id,
EbookCandidatePhrase.llm_judged.is_(False),
@@ -464,40 +633,42 @@ def prune_unstorable_unjudged_candidate_phrases(
return deleted
def delete_phrase_data_for_book(session: Session, book_id: int) -> PhraseRecalculationResult:
async def delete_phrase_data_for_book(session: AsyncSession, book_id: int) -> PhraseRecalculationResult:
"""Delete all candidate, protected, alias, and mention phrase data for one book.
Args:
session (Session): Active database session.
session (AsyncSession): Active database session.
book_id (int): Book whose phrase data is deleted.
Returns:
PhraseRecalculationResult: Deleted-row counts with ``candidate_phrases`` set to 0.
"""
protected_ids = session.scalars(
select(EbookProtectedPhrase.id).where(EbookProtectedPhrase.book_id == book_id)
protected_ids = (
await session.scalars(select(EbookProtectedPhrase.id).where(EbookProtectedPhrase.book_id == book_id))
).all()
deleted_aliases = 0
if protected_ids:
deleted_aliases = rowcount(
session.execute(delete(EbookPhraseAlias).where(EbookPhraseAlias.phrase_id.in_(protected_ids)))
await session.execute(delete(EbookPhraseAlias).where(EbookPhraseAlias.phrase_id.in_(protected_ids)))
)
deleted_mentions = rowcount(
session.execute(delete(EbookChunkPhraseMention).where(EbookChunkPhraseMention.book_id == book_id))
await 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)))
await session.execute(
delete(EbookChunkPhraseMention).where(EbookChunkPhraseMention.phrase_id.in_(protected_ids))
)
)
deleted_protected = rowcount(
session.execute(delete(EbookProtectedPhrase).where(EbookProtectedPhrase.book_id == book_id))
await session.execute(delete(EbookProtectedPhrase).where(EbookProtectedPhrase.book_id == book_id))
)
deleted_candidates = rowcount(
session.execute(delete(EbookCandidatePhrase).where(EbookCandidatePhrase.book_id == book_id))
await session.execute(delete(EbookCandidatePhrase).where(EbookCandidatePhrase.book_id == book_id))
)
session.flush()
await session.flush()
logger.info(
"ebook_candidate_phrase_data_deleted book_id=%s candidates=%s protected=%s aliases=%s mentions=%s",
book_id,
@@ -515,6 +686,7 @@ def delete_phrase_data_for_book(session: Session, book_id: int) -> PhraseRecalcu
candidate_phrases=0,
)
def rowcount(result: object) -> int:
"""Return a safe integer rowcount from a SQLAlchemy execution result.