| from __future__ import annotations
|
|
|
| from pathlib import Path
|
| import sys
|
|
|
| import pytest
|
|
|
|
|
| ROOT = Path(__file__).resolve().parents[1]
|
| PKG_ROOT = ROOT / "src" / "ai_agent"
|
| for p in (ROOT, PKG_ROOT):
|
| sp = str(p)
|
| if sp not in sys.path:
|
| sys.path.insert(0, sp)
|
|
|
| from ai_agent.retriever.reranker import CrossEncoderReranker
|
|
|
|
|
| class _DummyResponse:
|
| def __init__(self, payload: dict):
|
| self._payload = payload
|
|
|
| def raise_for_status(self) -> None:
|
| return None
|
|
|
| def json(self) -> dict:
|
| return self._payload
|
|
|
|
|
| def test_extract_rankings_results_shape() -> None:
|
| payload = {
|
| "results": [
|
| {"index": 2, "relevance_score": 0.2},
|
| {"index": 0, "relevance_score": 0.9},
|
| {"index": 1, "relevance_score": 0.5},
|
| ]
|
| }
|
| out = CrossEncoderReranker._extract_rankings(payload, doc_count=3)
|
| assert out == [(0, 0.9), (1, 0.5), (2, 0.2)]
|
|
|
|
|
| def test_rerank_uses_remote_endpoint(monkeypatch: pytest.MonkeyPatch) -> None:
|
| monkeypatch.setenv("EPFL_API_KEY_EMBEDDER", "test-key")
|
|
|
| captured: dict = {}
|
|
|
| def _fake_post(url, headers, json, timeout):
|
| captured["url"] = url
|
| captured["headers"] = headers
|
| captured["json"] = json
|
| captured["timeout"] = timeout
|
| return _DummyResponse(
|
| {
|
| "results": [
|
| {"index": 1, "relevance_score": 0.8},
|
| {"index": 0, "relevance_score": 0.1},
|
| ]
|
| }
|
| )
|
|
|
| monkeypatch.setattr("ai_agent.retriever.reranker.requests.post", _fake_post)
|
|
|
| rr = CrossEncoderReranker(
|
| backend="remote",
|
| model_name="BAAI/bge-reranker-v2-m3",
|
| base_url="https://inference-rcp.epfl.ch/v1",
|
| api_key_env="EPFL_API_KEY_EMBEDDER",
|
| )
|
| ranked = rr.rerank("segment lungs", ["doc a", "doc b"], top_k=2)
|
|
|
| assert ranked == [(1, 0.8), (0, 0.1)]
|
| assert captured["url"] == "https://inference-rcp.epfl.ch/v1/rerank"
|
| assert captured["headers"]["Authorization"] == "Bearer test-key"
|
| assert captured["json"]["model"] == "BAAI/bge-reranker-v2-m3"
|
| assert captured["json"]["top_n"] == 2
|
|
|
|
|
| def test_rerank_raises_without_api_key(monkeypatch: pytest.MonkeyPatch) -> None:
|
| monkeypatch.delenv("EPFL_API_KEY_EMBEDDER", raising=False)
|
| rr = CrossEncoderReranker(api_key_env="EPFL_API_KEY_EMBEDDER")
|
| with pytest.raises(RuntimeError, match="Missing reranker API key"):
|
| rr.rerank("q", ["a"], top_k=1)
|
|
|
