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()] + )