From 42d923fe4bb16555c594ca329ebd8a6bfe297dad Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 11 Aug 2026 10:08:10 +0300 Subject: [PATCH] Make the ledger's clocks defensible against the host A question identity is unset until a run establishes it. Testing whether the record exists cannot stand in for that: a host with no state to send seeds a default record, and a default record is truthy, so a resumption missing its real state would have proceeded as question zero. Every way of recording a message count now refuses one behind what is already stored, through one shared check: a new identity, an evidence outcome and a declaration, each against the newest identity, evidence epoch and declaration epoch. Identities and epochs are only comparable while the conversation grows, and `before_model_request` results are assigned back onto history, so that is a constraint on the host rather than a guarantee of the framework. Left unchecked, an evidence outcome moving backwards freezes every later declaration as stale, and a declaration moving backwards replaces a newer one with an older one and revives the answer it grounded. The searches, citations and executions of a question in progress survive a resumption. Clearing them cost the results the model was still answering from: a citation afterwards recorded no provenance and could not resolve against the expanded result it had seen, falling through to a database lookup. A code execution counts as evidence when it succeeded or printed something. A raised error with an empty stdout grounds nothing. --- CHANGELOG.md | 1 + .../haiku/rag/capabilities/_base.py | 24 +++- .../haiku/rag/capabilities/analysis.py | 3 +- .../haiku/rag/capabilities/ledger.py | 42 +++++- tests/capabilities/test_capabilities.py | 129 +++++++++++++++--- tests/capabilities/test_evidence_ledger.py | 93 ++++++++++--- 6 files changed, 243 insertions(+), 49 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9bd1b8b2..b9f86243 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,6 +15,7 @@ - `haiku-rag` and `haiku-ingester` CLI startup no longer imports `lancedb`, `pyarrow` and `pydantic_ai`. - `haiku.rag.store` no longer re-exports `Store`; import it from `haiku.rag.store.engine`. - `docling-local` reuses one docling `DocumentConverter` per set of conversion options instead of building one per document, so local layout, table and OCR models are no longer loaded per document. Conversions through a shared converter are serialized. +- A resumed run keeps the searches, citations and executions of the question in progress instead of clearing them. - Each page image attached to a search result is preceded by a line giving its position and the chunk id it came from. `build_binary_parts_from_results` is now `build_image_content_from_results` and returns those labels interleaved with the pictures. - Prior-question tool output is trimmed from the model request only; `all_messages()` retains what the run gathered. - Evidence retrieved for the current question is no longer trimmed mid-question, whether or not it carries page images. diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index 626a5652..9752d2b0 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py @@ -70,6 +70,13 @@ def resolve_db_path(db_path: Path | None, config: AppConfig) -> Path: def _clear_invocation_state(state: BaseModel) -> None: + """Drop the working evidence of the previous question. + + Only ever called when a new question starts. A resumption keeps it: the + results belong to the question still being answered, and dropping them leaves + a later citation unable to resolve against the expanded result the model saw, + recording no provenance for it. + """ for field_name in ("citations", "searches", "executions"): value = getattr(state, field_name, None) if hasattr(value, "clear"): @@ -196,24 +203,27 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): 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. + than guessed at. With no history at all there is nothing in progress: an + absent prompt is then an instructions-only first question, which takes an + identity like any other. """ 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"): + continuing = resuming and bool(ctx.messages) + state = self.state_type.model_validate(raw_state or {}) + record = cast(CapabilityEvidenceRecord, cast(Any, state).evidence) + if continuing and record.question is None: 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) + if not continuing: + _clear_invocation_state(state) + record.begin_question(len(ctx.messages)) run_capability = replace( self, state=state, diff --git a/haiku_rag_slim/haiku/rag/capabilities/analysis.py b/haiku_rag_slim/haiku/rag/capabilities/analysis.py index a9495e7f..90bb8147 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/analysis.py +++ b/haiku_rag_slim/haiku/rag/capabilities/analysis.py @@ -109,7 +109,8 @@ class AnalysisCapability(RAGCapabilityBase[AnalysisState]): ) sandbox = await self._ensure_sandbox() result = await sandbox.execute(code) - self._note_evidence() + if result.success or result.stdout: + 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 index 83868aad..4d837f79 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/ledger.py +++ b/haiku_rag_slim/haiku/rag/capabilities/ledger.py @@ -54,14 +54,43 @@ class CapabilityEvidenceRecord(BaseModel): ``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. + same values without sharing a counter. ``question`` is unset until a run + establishes it, so a record a host merely created is distinguishable from one + that has been through a question. """ occurrences: dict[str, EvidenceOccurrence] = Field(default_factory=dict) - question: int = 0 + question: int | None = None latest_evidence_epoch: int = 0 declaration: CitationDeclaration | None = None + def _reject_regression(self, count: int, what: str) -> None: + """Refuse a message count below one already recorded. + + Identities and epochs are both message counts, and every comparison + between them assumes the conversation only grows. One capability + truncating or reordering the history breaks that, and each way of + recording it has to refuse the same way: an unchecked evidence outcome + freezes every later declaration as stale, while an unchecked declaration + replaces a newer one with an older one and revives the answer it grounded. + """ + recorded = max( + self.question or 0, + self.latest_evidence_epoch, + self.declaration.epoch if self.declaration else 0, + ) + if count < recorded: + raise ValueError( + f"{what} at message count {count} is behind {recorded}, which is " + "already recorded: message history must be append-only for " + "question identities and epochs to hold." + ) + + def begin_question(self, identity: int) -> None: + """Take the identity of a question that has just arrived.""" + self._reject_regression(identity, "A question") + self.question = identity + def note_evidence(self, epoch: int) -> None: """Record that the model has seen an evidence outcome. @@ -70,7 +99,8 @@ class CapabilityEvidenceRecord(BaseModel): 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) + self._reject_regression(epoch, "Evidence") + self.latest_evidence_epoch = epoch def declare( self, @@ -86,6 +116,12 @@ class CapabilityEvidenceRecord(BaseModel): grounded. A call at a later epoch declares afresh, because evidence the model saw in between may be what it is now citing. """ + if self.question is None: + raise ValueError( + "Citations cannot be declared before a run establishes the " + "question identity." + ) + self._reject_regression(epoch, "A declaration") current = self.declaration if current is not None and (current.question, current.epoch) == ( self.question, diff --git a/tests/capabilities/test_capabilities.py b/tests/capabilities/test_capabilities.py index 98e6cc43..c0ab7f3f 100644 --- a/tests/capabilities/test_capabilities.py +++ b/tests/capabilities/test_capabilities.py @@ -314,7 +314,7 @@ async def test_search_and_empty_citation_limits(temp_db_path): @pytest.mark.asyncio async def test_cite_resolves_direct_chunk_ids_and_reuses_document_lookup(temp_db_path): capability = create_rag(db_path=temp_db_path, config=AppConfig()) - capability.state = RAGState() + capability.state = RAGState(evidence=CapabilityEvidenceRecord(question=0)) client = AsyncMock() client.get_chunk_by_id.side_effect = [ Chunk(id="chunk-1", document_id="doc-1", content="first"), @@ -340,7 +340,7 @@ async def test_cite_resolves_direct_chunk_ids_and_reuses_document_lookup(temp_db @pytest.mark.asyncio async def test_cite_reports_unresolved_ids_on_partial_success(temp_db_path): capability = create_rag(db_path=temp_db_path, config=AppConfig()) - capability.state = RAGState() + capability.state = RAGState(evidence=CapabilityEvidenceRecord(question=0)) client = AsyncMock() client.get_chunk_by_id.side_effect = [ Chunk(id="chunk-1", document_id="doc-1", content="first"), @@ -370,6 +370,7 @@ async def test_cite_repairs_chunk_ids_damaged_in_transcription(temp_db_path): unrelated = "9c2cd07e-5a3f-45a6-968d-cbd6f06ab57b" capability = create_rag(db_path=temp_db_path, config=AppConfig()) capability.state = RAGState( + evidence=CapabilityEvidenceRecord(question=0), searches={ "q": [ SearchResult( @@ -380,7 +381,7 @@ async def test_cite_repairs_chunk_ids_damaged_in_transcription(temp_db_path): document_uri="test://document", ) ] - } + }, ) client = AsyncMock() client.get_chunk_by_id.return_value = None @@ -482,17 +483,31 @@ async def test_a_spent_execution_budget_is_not_evidence(temp_db_path): assert capability.state.evidence.latest_evidence_epoch == 0 -@pytest.mark.parametrize("success", [True, False]) +@pytest.mark.parametrize( + ("success", "stdout", "expected_epoch"), + [ + pytest.param(True, "42", 5, id="succeeded"), + pytest.param(True, "", 5, id="succeeded without output"), + pytest.param(False, "42", 5, id="failed after printing"), + pytest.param(False, "", 0, id="failed without printing"), + ], +) @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.""" +async def test_only_a_code_execution_the_model_can_read_is_evidence( + temp_db_path, success, stdout, expected_epoch +): + """A raised error with nothing printed grounds nothing, so it is not evidence. + + A failure that printed first does ground an answer, and so does a successful + run whose outcome is that it printed nothing. + """ capability = create_analysis(db_path=temp_db_path, config=AppConfig()) - capability.state = AnalysisState() + capability.state = AnalysisState(evidence=CapabilityEvidenceRecord(question=0)) capability.epoch = 5 sandbox = AsyncMock(spec=Sandbox) sandbox._search_results = [] sandbox.execute.return_value = SandboxResult( - stdout="42", stderr="" if success else "boom", success=success + stdout=stdout, stderr="" if success else "boom", success=success ) capability.sandbox = sandbox @@ -502,7 +517,7 @@ async def test_a_code_execution_is_evidence_even_when_it_fails(temp_db_path, suc with pytest.raises(ToolFailed): await capability._execute_code("print(42)") - assert capability.state.evidence.latest_evidence_epoch == 5 + assert capability.state.evidence.latest_evidence_epoch == expected_epoch @pytest.mark.asyncio @@ -790,7 +805,9 @@ async def test_native_agent_composition_initializes_host_state(temp_db_path): result = await agent.run("Hello", deps=deps) assert result.output == "success (no tool calls)" - assert deps.state["rag"] == RAGState().model_dump(mode="json") + assert deps.state["rag"] == RAGState( + evidence=CapabilityEvidenceRecord(question=0) + ).model_dump(mode="json") @pytest.mark.asyncio @@ -1257,10 +1274,11 @@ async def test_a_question_takes_its_own_identity_and_both_capabilities_agree( 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()) + second_identity = _record(deps, "rag").question assert first_identity == 0 - assert _record(deps, "rag").question > first_identity - assert _record(deps, "analysis").question == _record(deps, "rag").question + assert second_identity is not None and second_identity > 0 + assert _record(deps, "analysis").question == second_identity @pytest.mark.asyncio @@ -1348,10 +1366,12 @@ async def test_citing_after_searching_grounds_the_question(temp_db_path): await agent.run("what does the supervisor do?", deps=deps) record = _record(deps, "rag") + question = record.question + assert question is not None 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" + assert record.occurrences["chunk-1"].retrieved_in_questions == [question] + assert citation_status([record], question=question) == "grounded" @pytest.mark.asyncio @@ -1376,8 +1396,10 @@ async def test_searching_after_citing_leaves_the_question_uncited(temp_db_path): await agent.run("what does the supervisor do?", deps=deps) record = _record(deps, "rag") + question = record.question + assert question is not None assert record.declaration is not None - assert citation_status([record], question=record.question) == "missing" + assert citation_status([record], question=question) == "missing" @pytest.mark.asyncio @@ -1406,9 +1428,11 @@ async def test_a_citation_in_the_same_request_as_its_search_is_not_current( await agent.run("what does the supervisor do?", deps=deps) record = _record(deps, "rag") + question = record.question + assert question is not None assert record.declaration is not None assert record.declaration.epoch == record.latest_evidence_epoch - assert citation_status([record], question=record.question) == "missing" + assert citation_status([record], question=question) == "missing" @pytest.mark.asyncio @@ -1444,6 +1468,7 @@ async def test_evidence_cited_in_two_questions_keeps_both_in_the_record(temp_db_ first_question, record.question, ] + assert record.question != first_question @pytest.mark.asyncio @@ -1496,8 +1521,78 @@ async def test_citing_without_searching_grounds_the_question(temp_db_path): await agent.run("cite chunk-1", deps=deps) record = _record(deps, "rag") + question = record.question + assert question is not None 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" + assert citation_status([record], question=question) == "grounded" + + +@pytest.mark.asyncio +async def test_a_host_seeded_record_does_not_pass_for_a_resumption(temp_db_path): + """A default record is truthy, so its presence cannot stand in for identity. + + Seeding one is what a host does when it has no state to send, and taking it + at face value would silently answer as question zero. + """ + 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", + message_history=_in_flight_history(), + deps=Deps(state={"rag": RAGState().model_dump(mode="json")}), + ) + + +@pytest.mark.asyncio +async def test_a_resumption_keeps_the_evidence_the_question_already_gathered( + temp_db_path, +): + """Clearing it would lose the results the model is still answering from. + + A citation after the resumption then records no provenance, and cannot resolve + against the expanded search result the model actually saw. + """ + rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False) + calls = iter( + [ + [ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")], + [TextPart("partial answer")], + [ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "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): + interrupted = await agent.run("what does the supervisor do?", deps=deps) + identity = _record(deps, "rag").question + assert identity is not None + await agent.run( + deferred_tool_results=DeferredToolResults( + calls={"call-2": "external result"} + ), + message_history=[ + *interrupted.all_messages(), + ModelResponse(parts=[ToolCallPart("external_tool", {}, "call-2")]), + ], + deps=deps, + ) + + record = _record(deps, "rag") + assert record.question == identity + assert record.occurrences["chunk-1"].retrieved_in_questions == [identity] + assert citation_status([record], question=identity) == "grounded" diff --git a/tests/capabilities/test_evidence_ledger.py b/tests/capabilities/test_evidence_ledger.py index d9c264c9..fa849c64 100644 --- a/tests/capabilities/test_evidence_ledger.py +++ b/tests/capabilities/test_evidence_ledger.py @@ -1,3 +1,5 @@ +import pytest + from haiku.rag.capabilities.ledger import ( CapabilityEvidenceRecord, CitationDeclaration, @@ -34,10 +36,10 @@ def test_no_declaration_reads_as_missing(): def test_refs_make_it_grounded_and_no_refs_make_it_ungrounded(): - grounded = CapabilityEvidenceRecord() + grounded = CapabilityEvidenceRecord(question=0) grounded.declare([rag_ref()], epoch=1) - ungrounded = CapabilityEvidenceRecord() + ungrounded = CapabilityEvidenceRecord(question=0) ungrounded.declare([], epoch=1) assert citation_status([grounded], question=0) == "grounded" @@ -50,7 +52,7 @@ def test_an_earlier_questions_declaration_is_never_current(): record.declare([rag_ref()], epoch=3) assert citation_status([record], question=2) == "grounded" - record.question = 8 + record.begin_question(8) assert record.declaration is not None assert citation_status([record], question=8) == "missing" @@ -58,7 +60,7 @@ def test_an_earlier_questions_declaration_is_never_current(): 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 = CapabilityEvidenceRecord(question=0) record.note_evidence(5) record.declare([rag_ref()], epoch=5) @@ -71,10 +73,10 @@ def test_a_citation_in_the_same_request_as_the_evidence_is_not_current(): def test_evidence_from_another_capability_after_citing_makes_it_uncited(): """Currency spans capabilities, which only works because epochs are global.""" - cited = CapabilityEvidenceRecord() + cited = CapabilityEvidenceRecord(question=0) cited.note_evidence(3) cited.declare([rag_ref()], epoch=5) - searched_after = CapabilityEvidenceRecord() + searched_after = CapabilityEvidenceRecord(question=0) searched_after.note_evidence(7) assert citation_status([cited], question=0) == "grounded" @@ -82,7 +84,7 @@ def test_evidence_from_another_capability_after_citing_makes_it_uncited(): def test_declarations_at_the_same_epoch_merge(): - record = CapabilityEvidenceRecord() + record = CapabilityEvidenceRecord(question=0) record.declare([rag_ref("c1")], epoch=3) record.declare([rag_ref("c2")], epoch=3) @@ -91,7 +93,7 @@ def test_declarations_at_the_same_epoch_merge(): def test_repeating_a_ref_at_the_same_epoch_does_not_duplicate_it(): - record = CapabilityEvidenceRecord() + record = CapabilityEvidenceRecord(question=0) record.declare([rag_ref()], epoch=3) record.declare([rag_ref()], epoch=3) @@ -100,11 +102,11 @@ def test_repeating_a_ref_at_the_same_epoch_does_not_duplicate_it(): def test_neither_cite_order_downgrades_a_grounded_declaration(): - grounded_then_empty = CapabilityEvidenceRecord() + grounded_then_empty = CapabilityEvidenceRecord(question=0) grounded_then_empty.declare([rag_ref()], epoch=3) grounded_then_empty.declare([], epoch=3) - empty_then_grounded = CapabilityEvidenceRecord() + empty_then_grounded = CapabilityEvidenceRecord(question=0) empty_then_grounded.declare([], epoch=3) empty_then_grounded.declare([rag_ref()], epoch=3) @@ -113,27 +115,19 @@ def test_neither_cite_order_downgrades_a_grounded_declaration(): def test_the_same_chunk_id_under_two_capabilities_stays_separate(): - rag = CapabilityEvidenceRecord() + rag = CapabilityEvidenceRecord(question=0) rag.declare([EvidenceRef(capability="rag", chunk_id="shared")], epoch=3) - analysis = CapabilityEvidenceRecord() + analysis = CapabilityEvidenceRecord(question=0) 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.begin_question(8) record.declare([rag_ref()], epoch=9) occurrence = record.occurrences["c1"] @@ -148,3 +142,60 @@ def test_a_declaration_records_the_question_and_epoch_it_was_made_at(): assert record.declaration == CitationDeclaration( question=6, epoch=11, refs=[rag_ref()] ) + + +def test_a_fresh_record_has_no_question_identity(): + """The identity is established by the run, and its absence must be detectable. + + A default record is truthy, so its mere presence cannot stand in for having + been through ``for_run``: a host that seeds one would otherwise pass the + resumption check with a fabricated identity of zero. + """ + assert CapabilityEvidenceRecord().question is None + + +def test_evidence_cannot_move_backwards_in_the_conversation(): + """Epochs are message counts, and currency depends on them only growing. + + Silently keeping the newer value would leave every later declaration stale + for the rest of the conversation, permanently and invisibly. + """ + record = CapabilityEvidenceRecord(question=0) + record.note_evidence(9) + + with pytest.raises(ValueError, match="append-only"): + record.note_evidence(4) + + +def test_citing_before_a_run_establishes_the_question_is_refused(): + with pytest.raises(ValueError, match="question identity"): + CapabilityEvidenceRecord().declare([rag_ref()], epoch=3) + + +def test_a_declaration_cannot_move_backwards(): + """Otherwise a stale citation replaces a newer one and revives the answer.""" + record = CapabilityEvidenceRecord(question=0) + record.declare([rag_ref("newer")], epoch=5) + + with pytest.raises(ValueError, match="append-only"): + record.declare([rag_ref("older")], epoch=3) + + assert record.declaration is not None + assert [ref.chunk_id for ref in record.declaration.refs] == ["newer"] + + +def test_evidence_cannot_predate_a_recorded_declaration(): + """The declaration's epoch is a recorded message count like any other.""" + record = CapabilityEvidenceRecord(question=0) + record.declare([rag_ref()], epoch=5) + + with pytest.raises(ValueError, match="append-only"): + record.note_evidence(3) + + +def test_a_question_cannot_start_before_what_is_already_recorded(): + record = CapabilityEvidenceRecord(question=0) + record.note_evidence(9) + + with pytest.raises(ValueError, match="append-only"): + record.begin_question(4)