feat(search): enhance phrase matching and reranking logic with improved async handling

python/ebook_search/search.py

@@ -129,18 +129,24 @@ async def search_ebooks(
         rank_constant=config.rrf_rank_constant,
     )
     timings.append(timing)
-    if phrase_matching:
+    phrase_boost_timing_name = "Phrase mention boost" if phrase_matching else "Phrase mention boost skipped"
-        fused, timing = await async_timed_result(
+    fused, timing = await async_timed_result(
-            "Phrase mention boost",
+        phrase_boost_timing_name,
-            apply_phrase_mention_boosts(engine, fused, phrase_matches, config.phrase_hit_boost),
+        apply_phrase_mention_boosts(
-        )
+            engine,
-    else:
+            fused,
-        fused, timing = timed_result("Phrase mention boost skipped", skip_phrase_mention_boosts, fused)
+            phrase_matches,
+            config.phrase_hit_boost,
+            phrase_matching=phrase_matching,
+        ),
+    )
     timings.append(timing)
-    if config.rerank.enabled and rerank:
+    rerank_enabled = config.rerank.enabled and rerank
-        response, timing = await async_timed_result("Rerank", apply_rerank(client, query, fused, config))
+    rerank_timing_name = "Rerank" if rerank_enabled else "Rerank skipped"
-    else:
+    response, timing = await async_timed_result(
-        response, timing = timed_result("Rerank skipped", skip_rerank, query, fused, config)
+        rerank_timing_name,
+        apply_rerank(client, query, fused, config, rerank=rerank_enabled),
+    )
     timings.append(timing)
     response = replace(response, timings=tuple(timings), phrase_matches=tuple(phrase_matches))
     logger.info(
@@ -171,19 +177,19 @@ async def query_phrase_matches(
         return []
 
 
-def skip_phrase_mention_boosts(candidates: list[SearchResult]) -> list[SearchResult]:
-    """Return candidates unchanged when phrase matching is disabled."""
-    logger.info(f"ebook_phrase_boost_skipped candidates={len(candidates)}")
-    return candidates
-
-
 async def apply_phrase_mention_boosts(
     engine: AsyncEngine,
     candidates: list[SearchResult],
     phrase_matches: Sequence[PhraseMatch],
     phrase_hit_boost: float,
+    *,
+    phrase_matching: bool,
 ) -> list[SearchResult]:
-    """Boost retrieved chunks that have indexed mentions for detected protected phrases."""
+    """Boost retrieved chunks that have indexed mentions for detected protected phrases when enabled."""
+    if not phrase_matching:
+        logger.info(f"ebook_phrase_boost_skipped candidates={len(candidates)}")
+        return candidates
+
     phrase_ids = sorted({match.phrase_id for match in phrase_matches})
     if not candidates or not phrase_ids or phrase_hit_boost <= 0:
         return candidates
@@ -274,23 +280,19 @@ async def parallel_retrieval(
     )
 
 
-def skip_rerank(
-    query: str,
-    candidates: list[SearchResult],
-    config: EbookSearchConfig,
-) -> SearchResponse:
-    """Return fused hybrid results without reranking."""
-    logger.info(f"ebook_rerank_skipped candidates={len(candidates)}")
-    return SearchResponse(query=query, results=candidates[: config.top_k], rank_label="Hybrid")
-
-
 async def apply_rerank(
     client: httpx.AsyncClient,
     query: str,
     candidates: list[SearchResult],
     config: EbookSearchConfig,
+    *,
+    rerank: bool,
 ) -> SearchResponse:
-    """Rerank already-fused hybrid candidates."""
+    """Rerank already-fused hybrid candidates when enabled for this request."""
+    if not rerank:
+        logger.info(f"ebook_rerank_skipped candidates={len(candidates)}")
+        return SearchResponse(query=query, results=candidates[: config.top_k], rank_label="Hybrid")
+
     reranked = await rerank_chunks(client, query, candidates[: config.rerank.candidates], config.rerank)
     logger.info(
         f"ebook_rerank_complete input_candidates={min(len(candidates), config.rerank.candidates)} "

tests/ebook_search/test_rag_pipeline.py

@@ -92,7 +92,10 @@ async def test_search_ebooks_skips_phrase_matching_when_disabled(mocker: MockerF
     )
     mocker.patch("python.ebook_search.search.bm25_candidates", return_value=[])
     detect_mock = mocker.patch("python.ebook_search.search.detect_protected_phrases_for_query")
-    boost_mock = mocker.patch("python.ebook_search.search.apply_phrase_mention_boosts")
+    boost_mock = mocker.patch(
+        "python.ebook_search.search.apply_phrase_mention_boosts",
+        side_effect=lambda _engine, candidates, *_args, **_kwargs: candidates,
+    )
     config = EbookSearchConfig(rerank=RerankConfig(enabled=False))
 
     response = await search_ebooks(
@@ -105,7 +108,8 @@ async def test_search_ebooks_skips_phrase_matching_when_disabled(mocker: MockerF
     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()
+    boost_mock.assert_awaited_once()
+    assert boost_mock.await_args.kwargs == {"phrase_matching": False}
 
 
 async def test_search_ebooks_ignores_phrase_matching_when_config_disabled(mocker: MockerFixture) -> None:
@@ -117,7 +121,10 @@ async def test_search_ebooks_ignores_phrase_matching_when_config_disabled(mocker
     )
     mocker.patch("python.ebook_search.search.bm25_candidates", return_value=[])
     detect_mock = mocker.patch("python.ebook_search.search.detect_protected_phrases_for_query")
-    boost_mock = mocker.patch("python.ebook_search.search.apply_phrase_mention_boosts")
+    boost_mock = mocker.patch(
+        "python.ebook_search.search.apply_phrase_mention_boosts",
+        side_effect=lambda _engine, candidates, *_args, **_kwargs: candidates,
+    )
     config = EbookSearchConfig(rerank=RerankConfig(enabled=False), phrase_matching_enabled=False)
 
     response = await search_ebooks(
@@ -130,4 +137,5 @@ async def test_search_ebooks_ignores_phrase_matching_when_config_disabled(mocker
     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()
+    boost_mock.assert_awaited_once()
+    assert boost_mock.await_args.kwargs == {"phrase_matching": False}

tests/ebook_search/test_rerank.py

@@ -10,7 +10,7 @@ import pytest
 
 from python.ebook_search.config import EbookSearchConfig, RerankConfig, load_rerank_config
 from python.ebook_search.rerank import rerank_chunks
-from python.ebook_search.search import SearchResult, apply_rerank, skip_rerank
+from python.ebook_search.search import SearchResult, apply_rerank
 
 if TYPE_CHECKING:
     from pytest_mock import MockerFixture
@@ -57,10 +57,10 @@ def test_config_defaults_enable_reranking(mocker: MockerFixture) -> None:
     assert config.timeout_seconds == 30
 
 
-def test_reranking_disabled_returns_original_fused_order() -> None:
+async def test_reranking_disabled_returns_original_fused_order(mocker: MockerFixture) -> None:
     config = EbookSearchConfig(rerank=RerankConfig(enabled=False), top_k=2)
 
-    response = skip_rerank("query", candidates(), config)
+    response = await apply_rerank(mocker.Mock(), "query", candidates(), config, rerank=False)
 
     assert response.rank_label == "Hybrid"
     assert [result.chunk_id for result in response.results] == [1, 2]
@@ -133,7 +133,7 @@ async def test_vllm_rerank_timeout_raises(mocker: MockerFixture) -> None:
     config = EbookSearchConfig(rerank=RerankConfig(enabled=True), top_k=2)
 
     with pytest.raises(httpx.TimeoutException, match="timeout"):
-        await apply_rerank(mocker.Mock(), "query", candidates(), config)
+        await apply_rerank(mocker.Mock(), "query", candidates(), config, rerank=True)
 
 
 async def test_malformed_vllm_rerank_json_does_not_crash_search(mocker: MockerFixture) -> None: