diff --git a/src/kb/eval/run_retrieval_only.py b/src/kb/eval/run_retrieval_only.py index 0fa2f00..53b07ee 100644 --- a/src/kb/eval/run_retrieval_only.py +++ b/src/kb/eval/run_retrieval_only.py @@ -153,7 +153,7 @@ async def _run(args: argparse.Namespace) -> int: question, top_k_dense=top_k_dense, top_k_sparse=top_k_sparse, - rerank_top_k=top_k_dense, + rerank_top_k=rerank_top_k, filters={"project": args.project}, ) if hits: diff --git a/src/kb/query/engine.py b/src/kb/query/engine.py index f2ee395..7c71f48 100644 --- a/src/kb/query/engine.py +++ b/src/kb/query/engine.py @@ -160,6 +160,33 @@ def _format_sources(hits: list[dict[str, Any]]) -> str: return "\n\n".join(out) +def _combine_sources_for_verification( + graph_sources: list[dict[str, Any]], retrieval_sources: list[dict[str, Any]] +) -> list[dict[str, Any]]: + """Flatten graph and retrieval sources in the same order as prompt numbering.""" + return [*graph_sources, *retrieval_sources] + + +def _verification_sources( + graph_result: dict[str, Any] | None, serializable_hits: list[dict[str, Any]] +) -> list[dict[str, Any]]: + """Build verifier inputs in the same order as the prompt source list.""" + graph_sources: list[dict[str, Any]] = [] + if graph_result and graph_result.get("mentions"): + for m in graph_result["mentions"]: + graph_sources.append( + { + "text": m.get("excerpt", ""), + "metadata": { + "file_id": m.get("file_id", ""), + "page_start": m.get("page_start", 0), + "page_end": m.get("page_end", m.get("page_start", 0)), + }, + } + ) + return _combine_sources_for_verification(graph_sources, serializable_hits) + + def _bookend_reorder(hits: list[Any]) -> list[Any]: """Reorder so highest-relevance chunks sit at the START and END of the prompt. @@ -788,7 +815,7 @@ def _section_score(h: Any) -> float: started = time.time() checks = await verify_citations( answer=answer_text, - sources=serializable_hits, + sources=_verification_sources(graph_result, serializable_hits), model=pipeline.get(cfg, "llm.synthesize.model"), ) verify_summary = verification_summary(checks) diff --git a/tests/test_eval_retrieval_only.py b/tests/test_eval_retrieval_only.py new file mode 100644 index 0000000..2f72657 --- /dev/null +++ b/tests/test_eval_retrieval_only.py @@ -0,0 +1,77 @@ +"""Retrieval-only eval wiring.""" + +from __future__ import annotations + +import asyncio +from argparse import Namespace + +import yaml + +from kb.eval import run_retrieval_only as eval_mod +from kb.vector.base import SearchHit + + +class FakeStore: + def __init__(self) -> None: + self.calls: list[dict] = [] + + async def ensure_collection(self, domain: str) -> None: + return None + + async def hybrid_search(self, domain: str, query: str, **kwargs): + self.calls.append({"domain": domain, "query": query, **kwargs}) + return [ + SearchHit( + id="hit-1", + text="alpha beta", + score=1.0, + metadata={"file_id": "file-1", "page_start": 1, "page_end": 1}, + ) + ] + + +def test_retrieval_only_uses_configured_rerank_top_k(monkeypatch, tmp_path) -> None: + ds_path = tmp_path / "ds.yaml" + ds_path.write_text( + yaml.safe_dump( + { + "questions": [ + {"id": "q1", "question": "alpha", "expected_files": ["file-1"], "tags": []} + ] + } + ) + ) + out_path = tmp_path / "out.json" + + store = FakeStore() + + monkeypatch.setattr(eval_mod, "get_store", lambda: store) + monkeypatch.setattr( + eval_mod.pipeline, + "pipeline_config", + lambda domain: {"retrieve": {"top_k_dense": 15, "top_k_sparse": 9, "rerank_top_k": 4}}, + ) + + async def fake_get_file(file_id: str): + return {"id": file_id, "filename": f"{file_id}.pdf"} + + monkeypatch.setattr(eval_mod.repo, "get_file", fake_get_file) + + async def fake_rerank(query: str, hits, top_k: int): + return hits[:top_k] + + monkeypatch.setattr(eval_mod, "cross_rerank", fake_rerank) + + rc = asyncio.run( + eval_mod._run( + Namespace( + project="default", + domain="sec", + dataset=str(ds_path), + output=str(out_path), + ) + ) + ) + + assert rc == 0 + assert store.calls[0]["rerank_top_k"] == 4 diff --git a/tests/test_query_synthesis.py b/tests/test_query_synthesis.py index 72ba858..d9f37bd 100644 --- a/tests/test_query_synthesis.py +++ b/tests/test_query_synthesis.py @@ -11,6 +11,7 @@ _extract_cited_indices, _format_numbered_sources, _format_sources, + _verification_sources, ) from kb.query.intent import QueryIntent from kb.query.structured import _safe_field_key, maybe_structured_answer @@ -110,6 +111,29 @@ async def fake_mentions_for(entity_ids: list[str], project: str = "default"): assert out["entities"][0]["id"] == "e1" +def test_verification_sources_include_graph_then_retrieval() -> None: + out = _verification_sources( + { + "mentions": [ + { + "file_id": "graph-1", + "page_start": 1, + "page_end": 2, + "excerpt": "graph evidence", + } + ] + }, + [ + { + "text": "retrieval evidence", + "metadata": {"file_id": "retr-1", "page_start": 3, "page_end": 3}, + } + ], + ) + + assert [s["metadata"]["file_id"] for s in out] == ["graph-1", "retr-1"] + + def test_trailing_confidence_json_pattern() -> None: sample = 'The answer is X. [1] {"confidence": 0.83, "confidence_reason": "well-supported"}' m = re.search(r"\{[^{}]*confidence[^{}]*\}\s*$", sample)