"""Tests for EPUB search reranking.""" from __future__ import annotations from os import environ from typing import TYPE_CHECKING import httpx 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 if TYPE_CHECKING: from pytest_mock import MockerFixture def candidates() -> list[SearchResult]: return [ SearchResult(chunk_id=1, text="alpha", source_title="A", score=0.9), SearchResult(chunk_id=2, text="beta", source_title="B", score=0.8), SearchResult(chunk_id=3, text="gamma", source_title="C", score=0.7), ] def make_async_client(mocker: MockerFixture, fake_post) -> httpx.AsyncClient: """Build a mock async client whose post call is served by fake_post.""" client = mocker.MagicMock(spec=httpx.AsyncClient) client.post = mocker.AsyncMock(side_effect=fake_post) return client def rerank_response(payload: dict[str, object] | None = None, *, content: bytes | None = None) -> httpx.Response: return httpx.Response( 200, content=content, json=payload, request=httpx.Request("POST", "http://rerank.test/rerank"), ) def test_config_defaults_enable_reranking(mocker: MockerFixture) -> None: mocker.patch.dict(environ, {}, clear=False) environ.pop("EBOOK_SEARCH_RERANK_ENABLED", None) environ.pop("EBOOK_SEARCH_RERANK_BASE_URL", None) environ.pop("EBOOK_SEARCH_RERANK_MODEL", None) environ.pop("EBOOK_SEARCH_RERANK_CANDIDATES", None) environ.pop("EBOOK_SEARCH_RERANK_TIMEOUT_SECONDS", None) config = load_rerank_config() assert config.enabled is True assert config.base_url == "http://192.168.90.25:8001" assert config.model == "qwen3-reranker-06b" assert config.candidates == 24 assert config.timeout_seconds == 30 async def test_reranking_disabled_returns_original_fused_order(mocker: MockerFixture) -> None: config = EbookSearchConfig(rerank=RerankConfig(enabled=False), top_k=2) 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] async def test_reranking_enabled_reorders_candidates(mocker: MockerFixture) -> None: def fake_post(_url: str, *, json: dict[str, object], timeout: float) -> httpx.Response: assert timeout == 30 assert json == { "model": "qwen3-reranker-06b", "query": "query", "documents": ["alpha", "beta", "gamma"], } return rerank_response( { "results": [ {"index": 0, "relevance_score": 0.1}, {"index": 1, "relevance_score": 0.9}, {"index": 2, "relevance_score": 0.4}, ] } ) client = make_async_client(mocker, fake_post) results = await rerank_chunks(client, "query", candidates(), RerankConfig()) assert [result.chunk_id for result in results] == [2, 1, 3] assert [round(result.score, 3) for result in results] == [0.78, 0.37, 0.28] assert [result.rerank_score for result in results] == [0.9, 0.1, 0.4] async def test_reranking_cannot_ignore_hybrid_score(mocker: MockerFixture) -> None: candidates = [ SearchResult(chunk_id=1, text="strong hybrid", source_title="A", score=1.0), SearchResult(chunk_id=2, text="weak hybrid", source_title="B", score=0.1), ] def fake_post(_url: str, **_kwargs: object) -> httpx.Response: return rerank_response( { "results": [ {"index": 0, "relevance_score": 0.7}, {"index": 1, "relevance_score": 1.0}, ] } ) client = make_async_client(mocker, fake_post) results = await rerank_chunks(client, "query", candidates, RerankConfig()) assert [result.chunk_id for result in results] == [1, 2] assert results[0].score == pytest.approx(0.79) assert results[1].score == 0.7 assert results[1].rerank_score == 1.0 async def test_vllm_rerank_timeout_raises(mocker: MockerFixture) -> None: def fake_rerank_chunks( _client: httpx.AsyncClient, _query: str, _candidates: list[SearchResult], _config: RerankConfig, ) -> list[SearchResult]: message = "timeout" raise httpx.TimeoutException(message) mocker.patch("python.ebook_search.search.rerank_chunks", side_effect=fake_rerank_chunks) config = EbookSearchConfig(rerank=RerankConfig(enabled=True), top_k=2) with pytest.raises(httpx.TimeoutException, match="timeout"): await apply_rerank(mocker.Mock(), "query", candidates(), config, rerank=True) async def test_malformed_vllm_rerank_json_does_not_crash_search(mocker: MockerFixture) -> None: def fake_post(_url: str, **_kwargs: object) -> httpx.Response: return rerank_response(content=b"not-json") client = make_async_client(mocker, fake_post) results = await rerank_chunks(client, "query", candidates()[:1], RerankConfig()) assert results[0].score == 0.3 async def test_vllm_rerank_scores_are_clamped(mocker: MockerFixture) -> None: def fake_post(_url: str, **_kwargs: object) -> httpx.Response: return rerank_response( { "results": [ {"index": 0, "relevance_score": -1}, {"index": 1, "relevance_score": 2}, ] } ) client = make_async_client(mocker, fake_post) results = await rerank_chunks(client, "query", candidates()[:2], RerankConfig()) assert {result.chunk_id: result.rerank_score for result in results} == {1: 0.0, 2: 1.0}