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

perf(ebook-search): run phrase detection in parallel with retrieval

Browse Source 
Move protected phrase detection into the retrieval gather so it runs
concurrently with vector and BM25 candidates instead of sequentially
before them. Make the search API accept real bool form fields for
rerank/phrase_matching, gate phrase matching on both the request and
config kill switch, and reflow log f-strings for readability.
This commit is contained in:
Richie
2026-07-24 11:38:51 -04:00
parent 8f1a69529c
commit 0028237579
5 changed files with 113 additions and 37 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/ebook_search/api/routes/search.py
+5 -4
View File
@@ -79,8 +79,9 @@ async def search(
engine: AppEngine,
client: AppHttpClient,
query: Annotated[str, Form()],
rerank: Annotated[str | None, Form()] = None,
phrase_matching: Annotated[str | None, Form()] = None,
*,
rerank: Annotated[bool, Form()] = False,
phrase_matching: Annotated[bool, Form()] = False,
) -> HTMLResponse:
"""Run a search and render HTMX results."""
try:
@@ -89,8 +90,8 @@ async def search(
client,
query,
config,
rerank=rerank == "true",
phrase_matching=phrase_matching == "true",
rerank=rerank,
phrase_matching=phrase_matching,
)
except Exception as error:
logger.exception("ebook_search_request_failed")
python/ebook_search/protected_phrases/extraction.py
+7 -6
View File
@@ -688,14 +688,14 @@ def extract_phrase_candidates_for_book(
raw_started_at = perf_counter()
raw = extract_raw_ngrams_by_chapter(chapters, config)
logger.info(
f"ebook_phrase_candidate_extract_raw_complete candidates={len(raw)} duration_ms={(perf_counter() - "
f"raw_started_at) * 1000:.1f}"
f"ebook_phrase_candidate_extract_raw_complete candidates={len(raw)} "
f"duration_ms={(perf_counter() - raw_started_at) * 1000:.1f}"
)
yake_started_at = perf_counter()
yake_candidates = extract_yake_candidates(book_text, config)
logger.info(
f"ebook_phrase_candidate_extract_yake_complete candidates={len(yake_candidates)} duration_ms={(perf_counter() - "
f"yake_started_at) * 1000:.1f}"
f"ebook_phrase_candidate_extract_yake_complete candidates={len(yake_candidates)} "
f"duration_ms={(perf_counter() - yake_started_at) * 1000:.1f}"
)
capitalized_started_at = perf_counter()
capitalized = extract_capitalized_phrases(book_text, config)
@@ -728,7 +728,8 @@ def extract_phrase_candidates_for_book(
f"ebook_phrase_candidate_extract_complete raw={len(raw)} yake={len(yake_candidates)} "
f"capitalized={len(capitalized)} metadata={len(metadata_candidates)} {pre_filter_count=} {filtered_too_short=} "
f"{filtered_too_rare=} {filtered_too_common=} {filtered_junk=} min_uses={minimum_candidate_raw_count(config)} "
f"storable={len(candidates)} limited={len(limited)} enrich_score_ms={(perf_counter() - enriched_started_at) * "
f"1000:.1f} duration_ms={(perf_counter() - started_at) * 1000:.1f}"
f"storable={len(candidates)} limited={len(limited)} "
f"enrich_score_ms={(perf_counter() - enriched_started_at) * 1000:.1f} "
f"duration_ms={(perf_counter() - started_at) * 1000:.1f}"
)
return limited
python/ebook_search/protected_phrases/generate_ngrams.py
+6 -6
View File
@@ -68,8 +68,8 @@ async def generate_candidate_phrases_for_books(
source_ids = (await session.scalars(source_query)).all()
books_seen = len(source_ids)
logger.info(
f"ebook_candidate_phrase_generation_start {books_seen=} {config.phrase_min_tokens=} {config.phrase_max_tokens=} "
f"{config.protected_phrase_max_candidates_per_book=}"
f"ebook_candidate_phrase_generation_start {books_seen=} {config.phrase_min_tokens=} "
f"{config.phrase_max_tokens=} {config.protected_phrase_max_candidates_per_book=}"
)
pool = get_extraction_pool(config.protected_phrase_extraction_workers)
@@ -244,8 +244,8 @@ async def generate_candidate_phrases_for_book(
await session.rollback()
raise
logger.info(
f"ebook_candidate_phrase_generation_book_duration {book_id=} {saved_count=} duration_ms={(perf_counter() - "
f"started_at) * 1000:.1f}"
f"ebook_candidate_phrase_generation_book_duration {book_id=} {saved_count=} "
f"duration_ms={(perf_counter() - started_at) * 1000:.1f}"
)
return saved_count
@@ -289,7 +289,7 @@ async def store_candidate_phrases_for_book(
)
saved_count = await bulk_upsert_unjudged_candidates(session, book_id, series_id, limited_candidates)
logger.info(
f"ebook_candidate_phrase_save_complete {book_id=} {saved_count=} save_ms={(perf_counter() - save_started_at) * "
f"1000:.1f}"
f"ebook_candidate_phrase_save_complete {book_id=} {saved_count=} "
f"save_ms={(perf_counter() - save_started_at) * 1000:.1f}"
)
return saved_count
python/ebook_search/search.py
+34 -19
View File
@@ -84,10 +84,11 @@ class SearchResponse:
@dataclass(frozen=True)
class RetrievalResponse:
"""Parallel retrieval output for vector and BM25 candidates."""
"""Parallel retrieval output for vector, BM25, and protected phrase candidates."""
vector_results: list[SearchResult]
lexical_results: list[SearchResult]
phrase_matches: list[HydratedPhraseMatch]
timings: tuple[RuntimeStep, ...]
@@ -97,28 +98,26 @@ async def search_ebooks(
query: str,
config: EbookSearchConfig,
*,
rerank: bool = False,
phrase_matching: bool | None = None,
rerank: bool,
phrase_matching: bool,
) -> SearchResponse:
"""Run hybrid vector/BM25 search and optional reranking."""
"""Run hybrid vector/BM25 search and optional reranking.
Phrase matching only runs when both the request asks for it and
``config.phrase_matching_enabled`` allows it.
"""
if not query.strip():
logger.info("ebook_search_empty_query")
return SearchResponse(query=query, results=[], rank_label="Hybrid")
phrase_matching_enabled = config.phrase_matching_enabled if phrase_matching is None else phrase_matching
logger.info(f"ebook_search_start query_length={len(query)} {rerank=} {phrase_matching_enabled=}")
phrase_matching = phrase_matching and config.phrase_matching_enabled
logger.info(f"ebook_search_start query_length={len(query)} {rerank=} {phrase_matching=}")
timings: list[RuntimeStep] = []
if phrase_matching_enabled:
phrase_matches, timing = await async_timed_result(
"Protected phrase detection", query_phrase_matches(engine, query, config)
)
else:
phrase_matches, timing = timed_result("Protected phrase detection skipped", skip_phrase_matches)
timings.append(timing)
retrieval, timing = await async_timed_result(
"Hybrid retrieval",
parallel_retrieval(engine, client, query, config),
parallel_retrieval(engine, client, query, config, phrase_matching=phrase_matching),
)
phrase_matches = retrieval.phrase_matches
timings.extend(retrieval.timings)
timings.append(timing)
fused, timing = timed_result(
@@ -129,7 +128,7 @@ async def search_ebooks(
rank_constant=config.rrf_rank_constant,
)
timings.append(timing)
if phrase_matching_enabled:
if phrase_matching:
fused, timing = await async_timed_result(
"Phrase mention boost",
apply_phrase_mention_boosts(engine, fused, phrase_matches, config.phrase_hit_boost),
@@ -145,7 +144,7 @@ async def search_ebooks(
response = replace(response, timings=tuple(timings), phrase_matches=tuple(phrase_matches))
logger.info(
f"ebook_search_complete vector_candidates={len(retrieval.vector_results)} "
f"lexical_candidates={len(retrieval.lexical_results)} fused_candidates={len(fused)} {phrase_matching_enabled=} "
f"lexical_candidates={len(retrieval.lexical_results)} fused_candidates={len(fused)} {phrase_matching=} "
f"phrase_matches={len(phrase_matches)} returned={len(response.results)} {response.rank_label=} "
f"{response.total_runtime_ms=:.1f}"
)
@@ -236,27 +235,43 @@ async def parallel_retrieval(
client: httpx.AsyncClient,
query: str,
config: EbookSearchConfig,
*,
phrase_matching: bool,
) -> RetrievalResponse:
"""Run vector and BM25 candidate retrieval concurrently with separate database sessions.
"""Run vector, BM25, and protected phrase retrieval concurrently with separate database sessions.
BM25 scoring is pure CPU work over the cached corpus, so it runs in a worker thread
instead of on the event loop.
instead of on the event loop. Protected phrase detection only depends on the query, so
it joins the gather as a third task when phrase matching is enabled.
"""
phrase_task = (
asyncio.create_task(
async_timed_result("Protected phrase detection", query_phrase_matches(engine, query, config))
)
if phrase_matching
else None
)
(vector_results, vector_timing), (lexical_results, lexical_timing) = await asyncio.gather(
async_timed_result("Embedding + vector search", vector_candidates(engine, client, query, config)),
async_timed_result("BM25 search", asyncio.to_thread(bm25_candidates, query, config)),
)
if phrase_task is not None:
phrase_matches, phrase_timing = await phrase_task
else:
phrase_matches, phrase_timing = timed_result("Protected phrase detection skipped", skip_phrase_matches)
logger.info(
f"ebook_parallel_retrieval_complete vector_candidates={len(vector_results)} "
f"lexical_candidates={len(lexical_results)}"
f"lexical_candidates={len(lexical_results)} phrase_matches={len(phrase_matches)}"
)
return RetrievalResponse(
vector_results=vector_results,
lexical_results=lexical_results,
phrase_matches=phrase_matches,
timings=(
replace(vector_timing, counts_toward_total=False),
replace(lexical_timing, counts_toward_total=False),
replace(phrase_timing, counts_toward_total=False),
),
)