diff --git a/haiku_rag_slim/haiku/rag/capabilities/compaction.py b/haiku_rag_slim/haiku/rag/capabilities/compaction.py index 6b1decbb..05b57667 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/compaction.py +++ b/haiku_rag_slim/haiku/rag/capabilities/compaction.py @@ -264,6 +264,32 @@ def compact_history( return compacted +def _require_a_record_of_what_was_cited( + evidence: Sequence[DiscoveredEvidence], + messages: list[ModelMessage], + boundary: int, +) -> None: + """Refuse to compact a capability's evidence when its record was not carried. + + Judged per capability, and only for one whose own evidence is actually at + stake: another capability's carried record says nothing about this one's, and a + capability the model never used has nothing to lose. Without the record there + is no capsule to put in the evidence's place, so compacting would drop it and + leave the citations the host already displayed as the only trace. + """ + for found in evidence: + if found.state_carried: + continue + if _newest_owned_return(messages, boundary, found.tool_names) is None: + continue + raise RuntimeError( + f"Evidence compaction found {found.capability} evidence from an earlier " + "question but no record of what it cited, so replacing it would retain " + "nothing. The host must carry the capability state between runs, " + f"alongside the message history: {found.capability} state was missing." + ) + + def _newest_owned_return( messages: list[ModelMessage], boundary: int, owned_tools: frozenset[str] ) -> tuple[int, int] | None: @@ -322,16 +348,10 @@ class EvidenceCompactionCapability(AbstractCapability[Any]): evidence = discover_evidence(ctx) boundary = question_in_progress(evidence) owned_tools = frozenset().union(*(found.tool_names for found in evidence)) - if boundary > 0 and _newest_owned_return( - request_context.messages, boundary, owned_tools - ): - if not any(found.state_carried for found in evidence): - raise RuntimeError( - "Evidence compaction found evidence from an earlier question but " - "no record of what it cited, so replacing it would retain " - "nothing. The host must carry the capability state between runs, " - "alongside the message history, for this capability to work." - ) + if boundary > 0: + _require_a_record_of_what_was_cited( + evidence, request_context.messages, boundary + ) if boundary > 0: await self._build_once(ctx, evidence) request_context.messages = compact_history( diff --git a/tests/capabilities/test_evidence_wire.py b/tests/capabilities/test_evidence_wire.py index 7e588ca1..ef6a1cef 100644 --- a/tests/capabilities/test_evidence_wire.py +++ b/tests/capabilities/test_evidence_wire.py @@ -19,6 +19,7 @@ from pydantic_ai.models.function import FunctionModel from pydantic_ai.models.test import TestModel from pydantic_ai.usage import RunUsage +from haiku.rag.capabilities.analysis import create_capability as create_analysis from haiku.rag.capabilities.compaction import ( RECEIPT, Capsule, @@ -789,3 +790,90 @@ async def test_compaction_proceeds_for_a_host_that_carries_state(temp_db_path): await agent.run("a follow-up", deps=carried, message_history=history) assert returns_of(wire[-1]) == [RECEIPT] + + +@pytest.mark.asyncio +async def test_compaction_refuses_when_one_capability_of_two_lost_its_record( + temp_db_path, +): + """One carried record does not vouch for the other capability's evidence. + + A host retaining only the RAG namespace leaves the analysis record empty, and + its earlier evidence would be replaced by receipts retaining nothing while the + RAG record made the loss look accounted for. + """ + 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 + ) + compactor = create_compaction() + + async def model(_messages, _info): # pragma: no cover - never reached + return ModelResponse(parts=[TextPart("answer")]) + + agent = Agent( + FunctionModel(model), + deps_type=Deps, + capabilities=[rag, analysis, compactor], + ) + history: list[Any] = [ + ModelRequest(parts=[UserPromptPart("an earlier question")]), + ModelResponse( + parts=[ToolCallPart("analysis_search", {"query": "q"}, "call-1")] + ), + ModelRequest( + parts=[ToolReturnPart("analysis_search", "ANALYSIS EVIDENCE", "call-1")] + ), + ModelResponse(parts=[TextPart("an answer")]), + ] + # Only the RAG namespace comes back, as a host whitelisting fields would send. + rag_only = Deps( + state={ + "rag": RAGState( + evidence=CapabilityEvidenceRecord(question=0, in_progress=False) + ).model_dump(mode="json") + } + ) + + with pytest.raises(RuntimeError, match="analysis"): + await agent.run("a follow-up", deps=rag_only, message_history=history) + + +@pytest.mark.asyncio +async def test_compaction_proceeds_when_the_capability_without_a_record_has_no_evidence( + temp_db_path, +): + """A capability the earlier question never used has nothing to lose. + + Refusing whenever any record is missing would stop a host that registers both + capabilities and only ever uses one, which is the composition the docs + recommend against but hosts still have. + """ + 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 + ) + compactor = create_compaction() + wire: list[list[Any]] = [] + + async def model(messages, _info): + wire.append(list(messages)) + return ModelResponse(parts=[TextPart("answer")]) + + agent = Agent( + FunctionModel(model), + deps_type=Deps, + capabilities=[rag, analysis, compactor], + ) + history = answered_question("an earlier question", evidence="RAG EVIDENCE") + rag_only = Deps( + state={ + "rag": RAGState( + evidence=CapabilityEvidenceRecord(question=0, in_progress=False) + ).model_dump(mode="json") + } + ) + + await agent.run("a follow-up", deps=rag_only, message_history=history) + + assert returns_of(wire[-1]) == [RECEIPT]