142 lines
6.6 KiB
Python
142 lines
6.6 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, *, phrase_matching):
|
|
"""Record that phrase detection started and return no matches."""
|
|
assert phrase_matching is True
|
|
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.detect_protected_phrases_for_query")
|
|
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(
|
|
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_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:
|
|
"""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.detect_protected_phrases_for_query")
|
|
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(
|
|
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_awaited_once()
|
|
assert boost_mock.await_args.kwargs == {"phrase_matching": False}
|