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-24 11:38:50 -04:00
parent 38c01ec121
commit 2706c4417d
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),