diff --git a/python/ebook_search/search.py b/python/ebook_search/search.py index 360e61c..ba0a408 100644 --- a/python/ebook_search/search.py +++ b/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: - fused, timing = await async_timed_result( - "Phrase mention boost", - apply_phrase_mention_boosts(engine, fused, phrase_matches, config.phrase_hit_boost), - ) - else: - fused, timing = timed_result("Phrase mention boost skipped", skip_phrase_mention_boosts, fused) + phrase_boost_timing_name = "Phrase mention boost" if phrase_matching else "Phrase mention boost skipped" + fused, timing = await async_timed_result( + phrase_boost_timing_name, + apply_phrase_mention_boosts( + engine, + fused, + phrase_matches, + config.phrase_hit_boost, + phrase_matching=phrase_matching, + ), + ) timings.append(timing) - if config.rerank.enabled and rerank: - response, timing = await async_timed_result("Rerank", apply_rerank(client, query, fused, config)) - else: - response, timing = timed_result("Rerank skipped", skip_rerank, query, fused, config) + rerank_enabled = config.rerank.enabled and rerank + rerank_timing_name = "Rerank" if rerank_enabled else "Rerank skipped" + response, timing = await async_timed_result( + 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)} " diff --git a/tests/ebook_search/test_rag_pipeline.py b/tests/ebook_search/test_rag_pipeline.py index 64f6d25..2df3917 100644 --- a/tests/ebook_search/test_rag_pipeline.py +++ b/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} diff --git a/tests/ebook_search/test_rerank.py b/tests/ebook_search/test_rerank.py index 0fcf907..f32c45c 100644 --- a/tests/ebook_search/test_rerank.py +++ b/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: