diff --git a/python/ebook_search/api/routes/search.py b/python/ebook_search/api/routes/search.py index a086021..874a7a5 100644 --- a/python/ebook_search/api/routes/search.py +++ b/python/ebook_search/api/routes/search.py @@ -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") diff --git a/python/ebook_search/protected_phrases/extraction.py b/python/ebook_search/protected_phrases/extraction.py index a0307f5..d54b00d 100644 --- a/python/ebook_search/protected_phrases/extraction.py +++ b/python/ebook_search/protected_phrases/extraction.py @@ -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 diff --git a/python/ebook_search/protected_phrases/generate_ngrams.py b/python/ebook_search/protected_phrases/generate_ngrams.py index 2cfdad7..fb4a906 100644 --- a/python/ebook_search/protected_phrases/generate_ngrams.py +++ b/python/ebook_search/protected_phrases/generate_ngrams.py @@ -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 diff --git a/python/ebook_search/search.py b/python/ebook_search/search.py index ed3f49e..ab578d7 100644 --- a/python/ebook_search/search.py +++ b/python/ebook_search/search.py @@ -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), ), ) diff --git a/tests/ebook_search/test_rag_pipeline.py b/tests/ebook_search/test_rag_pipeline.py index 598d79c..fa332f4 100644 --- a/tests/ebook_search/test_rag_pipeline.py +++ b/tests/ebook_search/test_rag_pipeline.py @@ -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]