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:
@@ -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
|
||||
|
||||
@@ -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),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@@ -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]
|
||||
|
||||
Reference in New Issue
Block a user