From 85594a1fd198a75e2c1d1384463ef7e354ac01bf Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 10 Aug 2026 15:09:53 +0300 Subject: [PATCH] Record what each evidence capability retrieved and cited `CapabilityEvidenceRecord` holds the relationships a transcript cannot express: which chunks a capability retrieved, which it cited, in which questions, and at which point in the conversation. RAG and analysis each own one in their own state namespace. Nothing is co-written: the host's state is JSON storage, so a shared record would be overwritten by whichever capability synced last, and merging happens in transient per-request views instead. Both clocks are derived from the conversation rather than counted locally, so every participant computes the same values without sharing a counter. Question identity is the message count when the question arrived; epoch is the message count at an outcome. Epochs are therefore globally comparable, which is what lets `citation_status` require a declaration to follow the newest evidence of every capability, and what makes equal epochs mean one request. A declaration is written only after `resolve_citations` succeeds, so a call naming only unresolvable ids is not a citation. Status is derived, never stored, so refs and status cannot contradict. Resuming a question requires the host to carry the capability state from the run being resumed. Without it the identity of the question in progress is unknowable, and adopting the current message count would relabel that question as a new one and judge every declaration in it against the wrong identity. Nothing reads the records yet and no wire behaviour changes. --- CHANGELOG.md | 8 + .../haiku/rag/capabilities/_base.py | 64 +++- .../haiku/rag/capabilities/analysis.py | 3 + .../haiku/rag/capabilities/ledger.py | 155 ++++++++ haiku_rag_slim/haiku/rag/capabilities/rag.py | 2 + tests/capabilities/test_capabilities.py | 346 +++++++++++++++++- tests/capabilities/test_evidence_ledger.py | 150 ++++++++ 7 files changed, 722 insertions(+), 6 deletions(-) create mode 100644 haiku_rag_slim/haiku/rag/capabilities/ledger.py create mode 100644 tests/capabilities/test_evidence_ledger.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 3944f227..9bd1b8b2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,14 @@ # Changelog ## [Unreleased] +### Added + +- `RAGState.evidence` / `AnalysisState.evidence` (`CapabilityEvidenceRecord`): which evidence a capability retrieved and cited, per question, keyed by message-count question identities and epochs. `haiku.rag.capabilities.ledger.citation_status(records, question=...)` derives `missing` / `grounded` / `ungrounded` across capabilities. + +### Changed + +- Resuming a run (no prompt, deferred tool results, an unfinished history tail) raises `RuntimeError` unless the host carries the capability state from the run being resumed. + ### Fixed - `cross-encoder` reranking no longer ties the scores of strongly-relevant candidates, which left their order to the sort. Scores remain 0-1. diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index 5803865c..626a5652 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py @@ -23,6 +23,7 @@ from pydantic_ai.tools import ToolDefinition from pydantic_ai.toolsets import AgentToolset from haiku.rag.capabilities._tools import CodeExecutionEntry, search_corpus +from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord, EvidenceRef from haiku.rag.client import HaikuRAG from haiku.rag.config.models import AppConfig from haiku.rag.store.models.chunk import SearchResult @@ -178,6 +179,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): request_count: int = field(default=0, repr=False) grace_requests_used: int = field(default=0, repr=False) turn_start: int = field(default=0, repr=False) + epoch: int = field(default=0, repr=False) async def for_run(self, ctx: RunContext[Any]) -> "RAGCapabilityBase[StateT]": """Start a run's own copy, and decide what counts as an earlier question. @@ -187,12 +189,31 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): messages and replace its evidence with the earlier-question notice, leaving the model to answer with the evidence taken away. Failing this way costs a larger request; failing the other way costs the answer. + + A new question takes the message count as its identity, which every + participant derives identically from the same history. A resumption keeps + the identity already recorded: the question is the one in progress, and + adopting the current count would relabel it as a new one and judge its + declarations against the wrong question. A resumption with no recorded + identity is a state this design does not produce, so it is reported rather + than guessed at — unless there is no history at all, where an absent prompt + means an instructions-only first question and nothing is in progress. """ outer = getattr(ctx.deps, "state", None) outer_state = outer if isinstance(outer, dict) else None raw_state = outer_state.get(self.state_namespace) if outer_state else None + resuming = _is_resumption(ctx.prompt, ctx.messages) + if resuming and ctx.messages and not (raw_state or {}).get("evidence"): + raise RuntimeError( + f"The {self.state_namespace} capability is resuming a question with " + "no stored question identity. Capabilities cannot be added, removed " + "or migrated while a question is unfinished, and the run's state " + "must be carried between its runs." + ) state = self.state_type.model_validate(raw_state or {}) _clear_invocation_state(state) + if not resuming: + cast(Any, state).evidence.question = len(ctx.messages) run_capability = replace( self, state=state, @@ -203,9 +224,8 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): search_count=0, request_count=0, grace_requests_used=0, - turn_start=( - 0 if _is_resumption(ctx.prompt, ctx.messages) else len(ctx.messages) - ), + epoch=0, + turn_start=0 if resuming else len(ctx.messages), ) run_capability._sync_state() return run_capability @@ -238,6 +258,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): async def before_model_request( self, ctx: RunContext[Any], request_context: ModelRequestContext ) -> ModelRequestContext: + self.epoch = len(ctx.messages) if instruction := self._budget_notice(): current_request = request_context.messages[-1] if isinstance(current_request, ModelRequest): @@ -390,6 +411,41 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): finally: self._sync_state() + def _evidence_record(self) -> CapabilityEvidenceRecord: + assert self.state is not None + return cast(CapabilityEvidenceRecord, cast(Any, self.state).evidence) + + def _note_evidence(self) -> None: + """Record an outcome the model can ground an answer on. + + Includes an empty search result and a failed execution that still printed + output: negative evidence grounds a refusal. Excludes a spent budget, which + yields nothing to ground anything on. + """ + self._evidence_record().note_evidence(self.epoch) + + def _declare(self, citations: list[Citation]) -> None: + """Record what the model cited, once the ids have resolved. + + Declaring earlier would let a call naming only unresolvable ids read as a + grounded answer. + """ + state = cast(Any, self.state) + retrieved = { + result.chunk_id + for results in state.searches.values() + for result in results + if result.chunk_id + } + self._evidence_record().declare( + [ + EvidenceRef(capability=self.state_namespace, chunk_id=c.chunk_id) + for c in citations + ], + epoch=self.epoch, + retrieved_now=retrieved, + ) + async def _search(self, query: str, limit: int | None) -> str | ToolReturn: assert self.state is not None self.search_count += 1 @@ -407,6 +463,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): ) state = cast(Any, self.state) state.searches[query] = results + self._note_evidence() if self.vision and (parts := build_image_content_from_results(results)): return ToolReturn(return_value=formatted, content=parts) return formatted @@ -454,6 +511,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): "Copy chunk_ids verbatim from search results." ) self._register_citations(citations) + self._declare(citations) resolved = {citation.chunk_id for citation in citations} unresolved = [cid for cid in missing if cid not in resolved] if unresolved: diff --git a/haiku_rag_slim/haiku/rag/capabilities/analysis.py b/haiku_rag_slim/haiku/rag/capabilities/analysis.py index 95fa155e..a9495e7f 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/analysis.py +++ b/haiku_rag_slim/haiku/rag/capabilities/analysis.py @@ -13,6 +13,7 @@ from haiku.rag.capabilities._base import ( RAGCapabilityBase, resolve_db_path, ) +from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord from haiku.rag.config.models import AppConfig from haiku.rag.sandbox import AnalysisContext, Sandbox from haiku.rag.store.models.chunk import SearchResult @@ -29,6 +30,7 @@ class AnalysisState(BaseModel): executions: list[CodeExecutionEntry] = Field(default_factory=list) citation_index: dict[str, Citation] = Field(default_factory=dict) citations: list[str] = Field(default_factory=list) + evidence: CapabilityEvidenceRecord = Field(default_factory=CapabilityEvidenceRecord) searches: dict[str, list[SearchResult]] = Field(default_factory=dict) @@ -107,6 +109,7 @@ class AnalysisCapability(RAGCapabilityBase[AnalysisState]): ) sandbox = await self._ensure_sandbox() result = await sandbox.execute(code) + self._note_evidence() if sandbox._search_results: existing = self.state.searches.get("_sandbox", []) seen = {item.chunk_id for item in existing} diff --git a/haiku_rag_slim/haiku/rag/capabilities/ledger.py b/haiku_rag_slim/haiku/rag/capabilities/ledger.py new file mode 100644 index 00000000..83868aad --- /dev/null +++ b/haiku_rag_slim/haiku/rag/capabilities/ledger.py @@ -0,0 +1,155 @@ +from collections.abc import Iterable +from typing import Literal + +from pydantic import BaseModel, Field + +CitationStatus = Literal["missing", "grounded", "ungrounded"] + + +class EvidenceRef(BaseModel): + """One piece of evidence, identified by its owner as well as its chunk. + + A chunk id alone is not an identity: the same id can be reported by more than + one capability, and ownership is what tells compaction whose output it may + touch. + """ + + capability: str + chunk_id: str + + +class EvidenceOccurrence(BaseModel): + """Which questions retrieved a piece of evidence, and which cited it.""" + + capability: str + chunk_id: str + retrieved_in_questions: list[int] = Field(default_factory=list) + cited_in_questions: list[int] = Field(default_factory=list) + + +class CitationDeclaration(BaseModel): + """What a question declared as its grounding, and when. + + Bound to a question *and* an epoch: the epoch outlives a question, so matching + it alone would let a question that gathered no evidence inherit the previous + declaration and read as compliant having declared nothing. + """ + + question: int + epoch: int + refs: list[EvidenceRef] = Field(default_factory=list) + + +class CapabilityEvidenceRecord(BaseModel): + """What one evidence capability wrote, in its own state namespace. + + Holds no content: ``Citation`` in ``citation_index`` is the canonical record + and already persists content, document id and picture refs. This is the index + over it that compaction needs and the transcript cannot provide. + + Single-writer by construction. A record shared between capabilities would be + overwritten by whichever of them synced its state last; merging happens in the + transient views built by ``citation_status`` and the optional capabilities. + + ``question`` is the number of messages that existed when the question arrived, + and ``epoch`` the number when an outcome occurred. Both are derived from the + conversation rather than counted locally, so every participant computes the + same values without sharing a counter. + """ + + occurrences: dict[str, EvidenceOccurrence] = Field(default_factory=dict) + question: int = 0 + latest_evidence_epoch: int = 0 + declaration: CitationDeclaration | None = None + + def note_evidence(self, epoch: int) -> None: + """Record that the model has seen an evidence outcome. + + Called for anything an answer could rest on, including a search that + returned nothing and a failed execution that still printed output — a + fruitless search grounds a refusal. Not called for a failure that yields + no evidence at all, such as an exhausted budget. + """ + self.latest_evidence_epoch = max(self.latest_evidence_epoch, epoch) + + def declare( + self, + refs: list[EvidenceRef], + *, + epoch: int, + retrieved_now: set[str] | None = None, + ) -> None: + """Record validated citations for the current question. + + Repeated calls at the same epoch merge, so citing again cannot narrow what + was already declared: an empty call after a grounded one leaves it + grounded. A call at a later epoch declares afresh, because evidence the + model saw in between may be what it is now citing. + """ + current = self.declaration + if current is not None and (current.question, current.epoch) == ( + self.question, + epoch, + ): + known = {(ref.capability, ref.chunk_id) for ref in current.refs} + current.refs.extend( + ref for ref in refs if (ref.capability, ref.chunk_id) not in known + ) + else: + self.declaration = CitationDeclaration( + question=self.question, epoch=epoch, refs=list(refs) + ) + + for ref in refs: + occurrence = self.occurrences.setdefault( + ref.chunk_id, + EvidenceOccurrence(capability=ref.capability, chunk_id=ref.chunk_id), + ) + if self.question not in occurrence.cited_in_questions: + occurrence.cited_in_questions.append(self.question) + if ( + retrieved_now + and ref.chunk_id in retrieved_now + and self.question not in occurrence.retrieved_in_questions + ): + occurrence.retrieved_in_questions.append(self.question) + + +def citation_status( + records: Iterable[CapabilityEvidenceRecord], *, question: int +) -> CitationStatus: + """Derived, never stored, so refs and status cannot contradict. + + A declaration is current only for the question it was made in, and only if it + followed the newest evidence outcome of *every* capability: a question where + one capability cited and another then searched without citing is not grounded. + Strictly later, since a citation made in the same request as an outcome cannot + have read it. + + A grounding *violation* is not one of these: that is an enforcement outcome + recorded by the policy capability, not something the model declared. + """ + records = list(records) + horizon = max((record.latest_evidence_epoch for record in records), default=0) + current = [ + record.declaration + for record in records + if record.declaration is not None + and record.declaration.question == question + and record.declaration.epoch > horizon + ] + if not current: + return "missing" + return ( + "grounded" if any(declaration.refs for declaration in current) else "ungrounded" + ) + + +__all__ = [ + "CapabilityEvidenceRecord", + "CitationDeclaration", + "CitationStatus", + "EvidenceOccurrence", + "EvidenceRef", + "citation_status", +] diff --git a/haiku_rag_slim/haiku/rag/capabilities/rag.py b/haiku_rag_slim/haiku/rag/capabilities/rag.py index 08931954..79216a87 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/rag.py +++ b/haiku_rag_slim/haiku/rag/capabilities/rag.py @@ -12,6 +12,7 @@ from haiku.rag.capabilities._base import ( RAGCapabilityBase, resolve_db_path, ) +from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord from haiku.rag.config.models import AppConfig from haiku.rag.store.models.chunk import SearchResult from haiku.rag.store.models.citation import Citation @@ -33,6 +34,7 @@ _instructions_path = Path(__file__).parent / "instructions" / "rag.md" class RAGState(BaseModel): citation_index: dict[str, Citation] = Field(default_factory=dict) citations: list[str] = Field(default_factory=list) + evidence: CapabilityEvidenceRecord = Field(default_factory=CapabilityEvidenceRecord) document_filter: str | None = None searches: dict[str, list[SearchResult]] = Field(default_factory=dict) diff --git a/tests/capabilities/test_capabilities.py b/tests/capabilities/test_capabilities.py index 7b37b11f..98e6cc43 100644 --- a/tests/capabilities/test_capabilities.py +++ b/tests/capabilities/test_capabilities.py @@ -26,6 +26,10 @@ from haiku.rag.capabilities._base import ( ) from haiku.rag.capabilities.analysis import AnalysisCapability, AnalysisState from haiku.rag.capabilities.analysis import create_capability as create_analysis +from haiku.rag.capabilities.ledger import ( + CapabilityEvidenceRecord, + citation_status, +) from haiku.rag.capabilities.rag import AGENT_PREAMBLE, RAGCapability, RAGState from haiku.rag.capabilities.rag import create_capability as create_rag from haiku.rag.config.models import AppConfig, PromptsConfig @@ -463,6 +467,44 @@ async def test_analysis_execution_limit_fails_the_tool(temp_db_path): await capability._execute_code("print('done')") +@pytest.mark.asyncio +async def test_a_spent_execution_budget_is_not_evidence(temp_db_path): + """Nothing was produced to ground an answer on, so nothing is recorded.""" + config = AppConfig() + config.analysis.max_executions = 0 + capability = create_analysis(db_path=temp_db_path, config=config) + capability.state = AnalysisState() + capability.epoch = 5 + + with pytest.raises(ToolFailed): + await capability._execute_code("print('done')") + + assert capability.state.evidence.latest_evidence_epoch == 0 + + +@pytest.mark.parametrize("success", [True, False]) +@pytest.mark.asyncio +async def test_a_code_execution_is_evidence_even_when_it_fails(temp_db_path, success): + """Output the model can read grounds an answer, whether the code raised or not.""" + capability = create_analysis(db_path=temp_db_path, config=AppConfig()) + capability.state = AnalysisState() + capability.epoch = 5 + sandbox = AsyncMock(spec=Sandbox) + sandbox._search_results = [] + sandbox.execute.return_value = SandboxResult( + stdout="42", stderr="" if success else "boom", success=success + ) + capability.sandbox = sandbox + + if success: + await capability._execute_code("print(42)") + else: + with pytest.raises(ToolFailed): + await capability._execute_code("print(42)") + + assert capability.state.evidence.latest_evidence_epoch == 5 + + @pytest.mark.asyncio async def test_spent_search_budget_is_announced_but_keeps_the_tool(rag_db): """A spent budget is announced; the tool stays declared to avoid a dead run. @@ -965,6 +1007,22 @@ async def test_compaction_never_reaches_the_stored_message_history(temp_db_path) assert PRIOR_TURN_NOTICE not in returns +def _resuming_deps() -> Deps: + """State as a resumption always finds it: the question already identified. + + A run that resumes has been through ``for_run`` before, so the identity of the + question in progress is stored. Fabricating the history without it is a state + the design does not produce, and is rejected rather than guessed at. + """ + return Deps( + state={ + "rag": RAGState(evidence=CapabilityEvidenceRecord(question=0)).model_dump( + mode="json" + ) + } + ) + + def _in_flight_history() -> list[Any]: """A question already asked and searched, still awaiting its answer.""" return [ @@ -1033,7 +1091,9 @@ async def test_a_prompt_on_an_unanswered_tail_is_treated_as_a_continuation( agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[capability]) await agent.run( - "a different question", deps=Deps(), message_history=_in_flight_history() + "a different question", + deps=_resuming_deps(), + message_history=_in_flight_history(), ) assert _wire_returns(wire[-1]) == ["EVIDENCE FOR THE LIVE TURN"] @@ -1066,7 +1126,7 @@ async def test_a_resume_carrying_a_prompt_keeps_the_active_evidence(temp_db_path "carry on", deferred_tool_results=DeferredToolResults(calls={"call-2": "external result"}), message_history=history, - deps=Deps(), + deps=_resuming_deps(), ) assert "EVIDENCE FOR THE LIVE TURN" in _wire_returns(wire[-1]) @@ -1105,7 +1165,9 @@ async def test_a_resumed_run_keeps_the_active_questions_evidence( agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[capability]) - await agent.run(deps=Deps(), message_history=_in_flight_history(), **resume_kwargs) + await agent.run( + deps=_resuming_deps(), message_history=_in_flight_history(), **resume_kwargs + ) assert _wire_returns(wire[-1]) == ["EVIDENCE FOR THE LIVE TURN"] @@ -1161,3 +1223,281 @@ async def test_prior_turn_search_is_compacted_but_its_cite_receipt_is_not(temp_d } assert prior["rag_search"] == PRIOR_TURN_NOTICE assert prior["rag_cite"] != PRIOR_TURN_NOTICE + + +def _record(deps: Deps, namespace: str) -> CapabilityEvidenceRecord: + return CapabilityEvidenceRecord.model_validate(deps.state[namespace]["evidence"]) + + +async def _stub_search(self, query: str, _limit: int | None) -> str: + """Record a result the way the real search does, so citing resolves.""" + cast(Any, self.state).searches[query] = [ + SearchResult(content="evidence", score=1.0, chunk_id="chunk-1") + ] + self._note_evidence() + return "EVIDENCE" + + +@pytest.mark.asyncio +async def test_a_question_takes_its_own_identity_and_both_capabilities_agree( + temp_db_path, +): + """Identity is derived from the conversation, so no counter is shared.""" + rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False) + analysis = create_analysis( + db_path=temp_db_path, config=AppConfig(), defer_loading=False + ) + + async def model(_messages, _info): + return ModelResponse(parts=[TextPart("answer")]) + + agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag, analysis]) + deps = Deps() + + first = await agent.run("first question", deps=deps) + first_identity = _record(deps, "rag").question + await agent.run("second question", deps=deps, message_history=first.all_messages()) + + assert first_identity == 0 + assert _record(deps, "rag").question > first_identity + assert _record(deps, "analysis").question == _record(deps, "rag").question + + +@pytest.mark.asyncio +async def test_a_resumption_keeps_the_identity_of_the_question_in_progress( + temp_db_path, +): + rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False) + + async def model(_messages, _info): + return ModelResponse(parts=[TextPart("answer")]) + + agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag]) + deps = Deps( + state={ + "rag": RAGState(evidence=CapabilityEvidenceRecord(question=7)).model_dump( + mode="json" + ) + } + ) + history = [ + *_in_flight_history(), + ModelResponse(parts=[ToolCallPart("external_tool", {}, "call-2")]), + ] + + await agent.run( + "carry on", + deferred_tool_results=DeferredToolResults(calls={"call-2": "external result"}), + message_history=history, + deps=deps, + ) + + assert _record(deps, "rag").question == 7 + + +@pytest.mark.asyncio +async def test_resuming_without_a_stored_identity_fails_instead_of_guessing( + temp_db_path, +): + """Adopting the message count would relabel a question already in progress. + + Every declaration and epoch comparison in it would then be judged against the + wrong question, silently. This state is not one the design produces, so it is + reported rather than repaired. + """ + rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False) + + async def model(_messages, _info): # pragma: no cover - never reached + return ModelResponse(parts=[TextPart("answer")]) + + agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag]) + + with pytest.raises(RuntimeError, match="no stored question identity"): + await agent.run( + "carry on", + deferred_tool_results=DeferredToolResults( + calls={"call-2": "external result"} + ), + message_history=[ + *_in_flight_history(), + ModelResponse(parts=[ToolCallPart("external_tool", {}, "call-2")]), + ], + deps=Deps(), + ) + + +@pytest.mark.asyncio +async def test_citing_after_searching_grounds_the_question(temp_db_path): + """The whole rule, end to end, with no compactor and no policy capability.""" + rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False) + calls = iter( + [ + [ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")], + [ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-2")], + [TextPart("answer")], + ] + ) + + async def model(_messages, _info): + return ModelResponse(parts=next(calls)) + + agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag]) + deps = Deps() + + with patch.object(RAGCapability, "_search", _stub_search): + await agent.run("what does the supervisor do?", deps=deps) + + record = _record(deps, "rag") + assert record.declaration is not None + assert [ref.chunk_id for ref in record.declaration.refs] == ["chunk-1"] + assert record.occurrences["chunk-1"].retrieved_in_questions == [record.question] + assert citation_status([record], question=record.question) == "grounded" + + +@pytest.mark.asyncio +async def test_searching_after_citing_leaves_the_question_uncited(temp_db_path): + rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False) + calls = iter( + [ + [ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")], + [ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-2")], + [ToolCallPart("rag_search", {"query": "again"}, "call-3")], + [TextPart("answer")], + ] + ) + + async def model(_messages, _info): + return ModelResponse(parts=next(calls)) + + agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag]) + deps = Deps() + + with patch.object(RAGCapability, "_search", _stub_search): + await agent.run("what does the supervisor do?", deps=deps) + + record = _record(deps, "rag") + assert record.declaration is not None + assert citation_status([record], question=record.question) == "missing" + + +@pytest.mark.asyncio +async def test_a_citation_in_the_same_request_as_its_search_is_not_current( + temp_db_path, +): + """Two calls in one response share an epoch, and citing must follow seeing.""" + rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False) + calls = iter( + [ + [ + ToolCallPart("rag_search", {"query": "supervisor"}, "call-1"), + ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-2"), + ], + [TextPart("answer")], + ] + ) + + async def model(_messages, _info): + return ModelResponse(parts=next(calls)) + + agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag]) + deps = Deps() + + with patch.object(RAGCapability, "_search", _stub_search): + await agent.run("what does the supervisor do?", deps=deps) + + record = _record(deps, "rag") + assert record.declaration is not None + assert record.declaration.epoch == record.latest_evidence_epoch + assert citation_status([record], question=record.question) == "missing" + + +@pytest.mark.asyncio +async def test_evidence_cited_in_two_questions_keeps_both_in_the_record(temp_db_path): + """Occurrences outlive the question that wrote them, through the state dict.""" + rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False) + calls = iter( + [ + [ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")], + [ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-2")], + [TextPart("first answer")], + [ToolCallPart("rag_search", {"query": "supervisor again"}, "call-3")], + [ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-4")], + [TextPart("second answer")], + ] + ) + + async def model(_messages, _info): + return ModelResponse(parts=next(calls)) + + agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag]) + deps = Deps() + + with patch.object(RAGCapability, "_search", _stub_search): + first = await agent.run("who supervises?", deps=deps) + first_question = _record(deps, "rag").question + await agent.run( + "and who supervises them?", deps=deps, message_history=first.all_messages() + ) + + record = _record(deps, "rag") + assert record.occurrences["chunk-1"].cited_in_questions == [ + first_question, + record.question, + ] + + +@pytest.mark.asyncio +async def test_a_run_with_no_prompt_and_no_history_starts_a_question(temp_db_path): + """An instructions-only run is a first question, not a resumption. + + There is no question in progress to keep an identity for, so nothing is + missing and the run proceeds with a fresh one. + """ + rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False) + + async def model(_messages, _info): + return ModelResponse(parts=[TextPart("answer")]) + + agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag]) + deps = Deps() + + await agent.run(deps=deps) + + assert _record(deps, "rag").question == 0 + + +@pytest.mark.asyncio +async def test_citing_without_searching_grounds_the_question(temp_db_path): + """A direct chunk-id citation stands on its own, with no evidence outcome. + + Epochs count messages and so start above zero, which is what lets a + declaration made in the first request still beat an empty evidence horizon. + """ + rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False) + calls = iter( + [ + [ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-1")], + [TextPart("answer")], + ] + ) + + async def model(_messages, _info): + return ModelResponse(parts=next(calls)) + + client = AsyncMock() + client.get_chunk_by_id.return_value = Chunk( + id="chunk-1", document_id="doc-1", content="evidence" + ) + client.get_document_by_id.return_value = None + agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag]) + deps = Deps() + + with patch.object(RAGCapability, "_ensure_rag", AsyncMock(return_value=client)): + await agent.run("cite chunk-1", deps=deps) + + record = _record(deps, "rag") + assert record.latest_evidence_epoch == 0 + assert record.declaration is not None + assert record.declaration.epoch > 0 + assert record.occurrences["chunk-1"].retrieved_in_questions == [] + assert citation_status([record], question=record.question) == "grounded" diff --git a/tests/capabilities/test_evidence_ledger.py b/tests/capabilities/test_evidence_ledger.py new file mode 100644 index 00000000..d9c264c9 --- /dev/null +++ b/tests/capabilities/test_evidence_ledger.py @@ -0,0 +1,150 @@ +from haiku.rag.capabilities.ledger import ( + CapabilityEvidenceRecord, + CitationDeclaration, + EvidenceRef, + citation_status, +) + + +def rag_ref(chunk_id: str = "c1") -> EvidenceRef: + return EvidenceRef(capability="rag", chunk_id=chunk_id) + + +def test_a_record_survives_the_state_round_trip(): + """Capability state is persisted as JSON, so the schema must survive it. + + A dict keyed by ``(capability, chunk_id)`` does not: the key serialises to + ``"rag,c1"`` and fails revalidation as a tuple. + """ + record = CapabilityEvidenceRecord(question=4) + record.note_evidence(5) + record.declare([rag_ref()], epoch=7, retrieved_now={"c1"}) + + restored = CapabilityEvidenceRecord.model_validate(record.model_dump(mode="json")) + + assert restored == record + assert citation_status([restored], question=4) == "grounded" + assert restored.occurrences["c1"].cited_in_questions == [4] + assert restored.occurrences["c1"].retrieved_in_questions == [4] + + +def test_no_declaration_reads_as_missing(): + assert citation_status([CapabilityEvidenceRecord()], question=0) == "missing" + assert citation_status([], question=0) == "missing" + + +def test_refs_make_it_grounded_and_no_refs_make_it_ungrounded(): + grounded = CapabilityEvidenceRecord() + grounded.declare([rag_ref()], epoch=1) + + ungrounded = CapabilityEvidenceRecord() + ungrounded.declare([], epoch=1) + + assert citation_status([grounded], question=0) == "grounded" + assert citation_status([ungrounded], question=0) == "ungrounded" + + +def test_an_earlier_questions_declaration_is_never_current(): + """Epochs outlive a question, so the epoch alone would inherit it.""" + record = CapabilityEvidenceRecord(question=2) + record.declare([rag_ref()], epoch=3) + assert citation_status([record], question=2) == "grounded" + + record.question = 8 + + assert record.declaration is not None + assert citation_status([record], question=8) == "missing" + + +def test_a_citation_in_the_same_request_as_the_evidence_is_not_current(): + """Citing must follow seeing: equal epochs mean one request.""" + record = CapabilityEvidenceRecord() + record.note_evidence(5) + record.declare([rag_ref()], epoch=5) + + assert citation_status([record], question=0) == "missing" + + record.declare([rag_ref()], epoch=7) + + assert citation_status([record], question=0) == "grounded" + + +def test_evidence_from_another_capability_after_citing_makes_it_uncited(): + """Currency spans capabilities, which only works because epochs are global.""" + cited = CapabilityEvidenceRecord() + cited.note_evidence(3) + cited.declare([rag_ref()], epoch=5) + searched_after = CapabilityEvidenceRecord() + searched_after.note_evidence(7) + + assert citation_status([cited], question=0) == "grounded" + assert citation_status([cited, searched_after], question=0) == "missing" + + +def test_declarations_at_the_same_epoch_merge(): + record = CapabilityEvidenceRecord() + record.declare([rag_ref("c1")], epoch=3) + record.declare([rag_ref("c2")], epoch=3) + + assert record.declaration is not None + assert [ref.chunk_id for ref in record.declaration.refs] == ["c1", "c2"] + + +def test_repeating_a_ref_at_the_same_epoch_does_not_duplicate_it(): + record = CapabilityEvidenceRecord() + record.declare([rag_ref()], epoch=3) + record.declare([rag_ref()], epoch=3) + + assert record.declaration is not None + assert len(record.declaration.refs) == 1 + + +def test_neither_cite_order_downgrades_a_grounded_declaration(): + grounded_then_empty = CapabilityEvidenceRecord() + grounded_then_empty.declare([rag_ref()], epoch=3) + grounded_then_empty.declare([], epoch=3) + + empty_then_grounded = CapabilityEvidenceRecord() + empty_then_grounded.declare([], epoch=3) + empty_then_grounded.declare([rag_ref()], epoch=3) + + assert citation_status([grounded_then_empty], question=0) == "grounded" + assert citation_status([empty_then_grounded], question=0) == "grounded" + + +def test_the_same_chunk_id_under_two_capabilities_stays_separate(): + rag = CapabilityEvidenceRecord() + rag.declare([EvidenceRef(capability="rag", chunk_id="shared")], epoch=3) + analysis = CapabilityEvidenceRecord() + analysis.declare([EvidenceRef(capability="analysis", chunk_id="shared")], epoch=3) + + assert rag.occurrences["shared"].capability == "rag" + assert analysis.occurrences["shared"].capability == "analysis" + + +def test_an_evidence_epoch_never_moves_backwards(): + record = CapabilityEvidenceRecord() + record.note_evidence(9) + record.note_evidence(4) + + assert record.latest_evidence_epoch == 9 + + +def test_citing_the_same_chunk_in_two_questions_records_both(): + record = CapabilityEvidenceRecord(question=2) + record.declare([rag_ref()], epoch=3, retrieved_now={"c1"}) + record.question = 8 + record.declare([rag_ref()], epoch=9) + + occurrence = record.occurrences["c1"] + assert occurrence.cited_in_questions == [2, 8] + assert occurrence.retrieved_in_questions == [2] + + +def test_a_declaration_records_the_question_and_epoch_it_was_made_at(): + record = CapabilityEvidenceRecord(question=6) + record.declare([rag_ref()], epoch=11) + + assert record.declaration == CitationDeclaration( + question=6, epoch=11, refs=[rag_ref()] + )