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

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:
2026-07-16 13:09:34 -04:00
parent 49b93c79f6
commit ed1ea4546a
5 changed files with 113 additions and 37 deletions
+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")
@@ -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
@@ -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
+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),
),
)
+61 -2
View File
@@ -40,7 +40,9 @@ async def test_search_ebooks_runs_vector_and_bm25_in_parallel(mocker: MockerFixt
mocker.patch("python.ebook_search.search.bm25_candidates", side_effect=fake_bm25_candidates)
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
response = await search_ebooks(engine, mocker.Mock(), "what is parallel", config)
response = await search_ebooks(
engine, mocker.Mock(), "what is parallel", config, rerank=False, phrase_matching=False
)
timings = {step.name: step for step in response.timings}
assert [result.chunk_id for result in response.results] == [1, 2]
@@ -50,6 +52,36 @@ async def test_search_ebooks_runs_vector_and_bm25_in_parallel(mocker: MockerFixt
assert received_engines == [engine]
async def test_search_ebooks_runs_phrase_detection_in_parallel_with_retrieval(mocker: MockerFixture) -> None:
"""Phrase detection joins the retrieval gather instead of running before it."""
engine = create_async_engine("sqlite+aiosqlite:///:memory:")
phrase_started = Event()
async def fake_vector_candidates(_engine, _client, _query, _config):
"""Return vector candidates only once phrase detection has started."""
assert await asyncio.to_thread(phrase_started.wait, 2)
return [SearchResult(chunk_id=1, text="vector", source_title="Vector", vector_score=0.9)]
async def fake_query_phrase_matches(_engine, _query, _config):
"""Record that phrase detection started and return no matches."""
phrase_started.set()
return []
mocker.patch("python.ebook_search.search.vector_candidates", side_effect=fake_vector_candidates)
mocker.patch("python.ebook_search.search.bm25_candidates", return_value=[])
mocker.patch("python.ebook_search.search.query_phrase_matches", side_effect=fake_query_phrase_matches)
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
response = await search_ebooks(
engine, mocker.Mock(), "what is parallel", config, rerank=False, phrase_matching=True
)
timings = {step.name: step for step in response.timings}
assert [result.chunk_id for result in response.results] == [1]
assert timings["Protected phrase detection"].counts_toward_total is False
assert timings["Hybrid retrieval"].counts_toward_total is True
async def test_search_ebooks_skips_phrase_matching_when_disabled(mocker: MockerFixture) -> None:
"""Phrase matching can be disabled for one search request."""
engine = create_async_engine("sqlite+aiosqlite:///:memory:")
@@ -62,7 +94,34 @@ async def test_search_ebooks_skips_phrase_matching_when_disabled(mocker: MockerF
boost_mock = mocker.patch("python.ebook_search.search.apply_phrase_mention_boosts")
config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
response = await search_ebooks(engine, mocker.Mock(), "what is parallel", config, phrase_matching=False)
response = await search_ebooks(
engine, mocker.Mock(), "what is parallel", config, rerank=False, phrase_matching=False
)
timing_names = {step.name for step in response.timings}
assert [result.chunk_id for result in response.results] == [1]
assert response.phrase_matches == ()
assert "Protected phrase detection skipped" in timing_names
assert "Phrase mention boost skipped" in timing_names
detect_mock.assert_not_called()
boost_mock.assert_not_called()
async def test_search_ebooks_ignores_phrase_matching_when_config_disabled(mocker: MockerFixture) -> None:
"""The config kill switch overrides a request that asks for phrase matching."""
engine = create_async_engine("sqlite+aiosqlite:///:memory:")
mocker.patch(
"python.ebook_search.search.vector_candidates",
return_value=[SearchResult(chunk_id=1, text="vector", source_title="Vector", vector_score=0.9)],
)
mocker.patch("python.ebook_search.search.bm25_candidates", return_value=[])
detect_mock = mocker.patch("python.ebook_search.search.query_phrase_matches")
boost_mock = mocker.patch("python.ebook_search.search.apply_phrase_mention_boosts")
config = EbookSearchConfig(rerank=RerankConfig(enabled=False), phrase_matching_enabled=False)
response = await search_ebooks(
engine, mocker.Mock(), "what is parallel", config, rerank=False, phrase_matching=True
)
timing_names = {step.name for step in response.timings}
assert [result.chunk_id for result in response.results] == [1]