diff --git a/evaluations/evaluations/benchmark.py b/evaluations/evaluations/benchmark.py index 1249660c..a550bc9f 100644 --- a/evaluations/evaluations/benchmark.py +++ b/evaluations/evaluations/benchmark.py @@ -235,11 +235,10 @@ async def run_retrieval_benchmark( seen = set() identifiers = [] for result in chunks: - # Use arxiv_id from metadata if present, otherwise use URI - doc_id = result.document_meta.get("arxiv_id") or result.document_uri - if doc_id and doc_id not in seen: - identifiers.append(doc_id) - seen.add(doc_id) + uri = result.document_uri + if uri and uri not in seen: + identifiers.append(uri) + seen.add(uri) return identifiers diff --git a/evaluations/tests/test_benchmark.py b/evaluations/tests/test_benchmark.py index 782a9c5e..2e49a836 100644 --- a/evaluations/tests/test_benchmark.py +++ b/evaluations/tests/test_benchmark.py @@ -515,21 +515,21 @@ class TestRetrievalTarget: assert searches[0]["include_images"] is False @pytest.mark.asyncio - async def test_prefers_metadata_identifier_over_uri(self, tmp_path: Path) -> None: + async def test_ranks_each_document_once(self, tmp_path: Path) -> None: from haiku.rag.store.models.chunk import SearchResult from evaluations.benchmark import run_retrieval_benchmark + def _result(uri: str, score: float) -> SearchResult: + return SearchResult(content="x", score=score, document_uri=uri) + class FakeRag: async def search(self, **kwargs) -> list[SearchResult]: return [ - SearchResult( - content="x", - score=1.0, - document_id="doc-1", - document_uri="uri-other", - document_meta={"arxiv_id": "uri-x"}, - ) + _result("uri-other", 1.0), + _result("uri-x", 0.9), + _result("uri-other", 0.8), + _result("uri-x", 0.7), ] with patch("evaluations.benchmark.HaikuRAG") as mock_haiku: @@ -538,8 +538,9 @@ class TestRetrievalTarget: self._spec(), AppConfig(), db_path=tmp_path / "test.lancedb" ) + # uri-x is the only relevant document and ranks second of two assert result is not None - assert result["map"] == 1.0 + assert result["map"] == 0.5 class TestEvaluateDatasetCaseIds: