dotfiles/tests/ebook_search/test_rerank.py

"""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, skip_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 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


def test_reranking_disabled_returns_original_fused_order() -> None:
    config = EbookSearchConfig(rerank=RerankConfig(enabled=False), top_k=2)

    response = skip_rerank("query", candidates(), config)

    assert response.rank_label == "Hybrid"
    assert [result.chunk_id for result in response.results] == [1, 2]


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},
                ]
            }
        )

    mocker.patch.object(httpx, "post", side_effect=fake_post)

    results = rerank_chunks("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]


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},
                ]
            }
        )

    mocker.patch.object(httpx, "post", side_effect=fake_post)

    results = rerank_chunks("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


def test_vllm_rerank_timeout_raises(mocker: MockerFixture) -> None:
    def fake_rerank_chunks(
        _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"):
        apply_rerank("query", candidates(), config)


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")

    mocker.patch.object(httpx, "post", side_effect=fake_post)

    results = rerank_chunks("query", candidates()[:1], RerankConfig())

    assert results[0].score == 0.3


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},
                ]
            }
        )

    mocker.patch.object(httpx, "post", side_effect=fake_post)

    results = rerank_chunks("query", candidates()[:2], RerankConfig())

    assert {result.chunk_id: result.rerank_score for result in results} == {1: 0.0, 2: 1.0}