Richie
0028237579
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.
133 lines
6.2 KiB
Python
133 lines
6.2 KiB
Python
"""Tests for the ebook search RAG pipeline orchestration."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
from threading import Event
|
|
from typing import TYPE_CHECKING
|
|
|
|
from sqlalchemy.ext.asyncio import create_async_engine
|
|
|
|
from python.ebook_search.config import EbookSearchConfig, RerankConfig
|
|
from python.ebook_search.search import SearchResult, search_ebooks
|
|
|
|
if TYPE_CHECKING:
|
|
from pytest_mock import MockerFixture
|
|
|
|
|
|
async def test_search_ebooks_runs_vector_and_bm25_in_parallel(mocker: MockerFixture) -> None:
|
|
engine = create_async_engine("sqlite+aiosqlite:///:memory:")
|
|
vector_started = Event()
|
|
bm25_started = Event()
|
|
received_engines: list[object] = []
|
|
|
|
async def fake_vector_candidates(received_engine, _client, query, _config):
|
|
"""Return vector candidates after confirming BM25 has started."""
|
|
received_engines.append(received_engine)
|
|
assert query == "what is parallel"
|
|
vector_started.set()
|
|
assert await asyncio.to_thread(bm25_started.wait, 2)
|
|
return [SearchResult(chunk_id=1, text="vector", source_title="Vector", vector_score=0.9)]
|
|
|
|
def fake_bm25_candidates(query, _config):
|
|
"""Return BM25 candidates after confirming vector search has started."""
|
|
assert query == "what is parallel"
|
|
bm25_started.set()
|
|
assert vector_started.wait(timeout=2)
|
|
return [SearchResult(chunk_id=2, text="bm25", source_title="BM25", bm25_score=2.0)]
|
|
|
|
mocker.patch("python.ebook_search.search.vector_candidates", side_effect=fake_vector_candidates)
|
|
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, 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]
|
|
assert timings["Embedding + vector search"].counts_toward_total is False
|
|
assert timings["BM25 search"].counts_toward_total is False
|
|
assert timings["Hybrid retrieval"].counts_toward_total is True
|
|
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:")
|
|
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))
|
|
|
|
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]
|
|
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()
|