diff --git a/evaluations/evaluations/capability_runner.py b/evaluations/evaluations/capability_runner.py index cccf2c1d..9c0488a5 100644 --- a/evaluations/evaluations/capability_runner.py +++ b/evaluations/evaluations/capability_runner.py @@ -43,6 +43,10 @@ class CapabilityRunResult: answer: str cited_uris: list[str] = field(default_factory=list) cited_chunk_ids: list[str] = field(default_factory=list) + # The database each cited chunk came from, in the order they were cited, so a + # run over several databases records which one grounded the answer. Empty + # strings where the database is unnamed, since one database names nothing. + cited_sources: list[str] = field(default_factory=list) searched_uris: list[str] = field(default_factory=list) n_searches: int = 0 n_executions: int = 0 @@ -244,10 +248,12 @@ def _result_from_run( cited_chunk_ids: list[str] = list(typed.citations) seen_cited: set[str] = set() cited_uris: list[str] = [] + cited_sources: list[str] = [] for chunk_id in cited_chunk_ids: citation = typed.citation_index.get(chunk_id) if citation is None: continue + cited_sources.append(citation.source or "") if citation.document_uri not in seen_cited: seen_cited.add(citation.document_uri) cited_uris.append(citation.document_uri) @@ -275,6 +281,7 @@ def _result_from_run( answer=answer, cited_uris=cited_uris, cited_chunk_ids=cited_chunk_ids, + cited_sources=cited_sources, searched_uris=searched_uris, # Distinct search keys, not searches. Analysis files every in-code # `search()` under one "_sandbox" key, so twenty sandbox searches read diff --git a/evaluations/evaluations/qa.py b/evaluations/evaluations/qa.py index bb8082e9..e420d8b8 100644 --- a/evaluations/evaluations/qa.py +++ b/evaluations/evaluations/qa.py @@ -351,6 +351,7 @@ async def run_qa_benchmark( ) set_eval_attribute("cited_uris", result.cited_uris) set_eval_attribute("cited_chunk_ids", result.cited_chunk_ids) + set_eval_attribute("cited_sources", result.cited_sources) set_eval_attribute("searched_uris", result.searched_uris) set_eval_attribute("n_searches", result.n_searches) set_eval_attribute("n_search_calls", result.n_search_calls) diff --git a/evaluations/tests/test_capability_runner.py b/evaluations/tests/test_capability_runner.py index d7adeba5..d14e7824 100644 --- a/evaluations/tests/test_capability_runner.py +++ b/evaluations/tests/test_capability_runner.py @@ -452,3 +452,62 @@ async def test_gold_prefix_run_answers_with_history(tmp_path): ) assert result.answer == "success (no tool calls)" + + +def test_records_the_database_each_citation_came_from(): + """A run over several databases has to record which one grounded the answer: + the distribution cannot be recovered from the report afterwards.""" + from haiku.rag.capabilities._base import EvidenceState + from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord + from haiku.rag.store.models.citation import Citation + + from evaluations.capability_runner import ToolTraffic, _result_from_run + + def cited(chunk_id: str, source: str | None) -> Citation: + return Citation( + chunk_id=chunk_id, + document_id=f"doc-{chunk_id}", + document_uri=f"test://{chunk_id}", + content="body", + source=source, + ) + + state = EvidenceState( + citations=["a1", "b1", "a2"], + citation_index={ + "a1": cited("a1", "alpha"), + "b1": cited("b1", "beta"), + "a2": cited("a2", "alpha"), + }, + evidence=CapabilityEvidenceRecord(question=1), + ) + + result = _result_from_run("answer", state, ToolTraffic(0, 0, 0, 0)) + + assert result.cited_sources == ["alpha", "beta", "alpha"] + + +def test_an_unnamed_database_records_no_source(): + """One database names nothing, so the field stays empty rather than absent.""" + from haiku.rag.capabilities._base import EvidenceState + from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord + from haiku.rag.store.models.citation import Citation + + from evaluations.capability_runner import ToolTraffic, _result_from_run + + state = EvidenceState( + citations=["c1"], + citation_index={ + "c1": Citation( + chunk_id="c1", + document_id="d1", + document_uri="test://one", + content="body", + ) + }, + evidence=CapabilityEvidenceRecord(question=1), + ) + + result = _result_from_run("answer", state, ToolTraffic(0, 0, 0, 0)) + + assert result.cited_sources == [""]