From 7bdd11db39212a2872c8c87682c48bd85253f860 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 13 Aug 2026 14:08:09 +0300 Subject: [PATCH 1/7] Update caps documentation --- docs/capabilities/analysis.md | 14 +++-- docs/capabilities/compaction.md | 43 +++++++++++++ docs/capabilities/index.md | 105 ++++++++------------------------ docs/capabilities/policy.md | 49 +++++++++++++++ docs/capabilities/rag.md | 19 ++++-- docs/chat.md | 11 ++-- zensical.toml | 2 + 7 files changed, 151 insertions(+), 92 deletions(-) create mode 100644 docs/capabilities/compaction.md create mode 100644 docs/capabilities/policy.md diff --git a/docs/capabilities/analysis.md b/docs/capabilities/analysis.md index 5effc127..b3a58138 100644 --- a/docs/capabilities/analysis.md +++ b/docs/capabilities/analysis.md @@ -18,18 +18,24 @@ When `qa.max_searches` or `analysis.max_executions` runs out, the exhausted tool The sandbox exposes documents under `/documents/{document_id}/` with `metadata.json`, `content.txt`, `items.jsonl`, and `toc.json`. -## Compose with RAG +## Compose an agent + +Register it on its own, not alongside `RAGCapability`: it already searches and cites, +and the two together give the model duplicate tools and separate budgets. See +[Capabilities](index.md#compose-an-agent). ```python from pydantic_ai import Agent from haiku.rag.capabilities.analysis import create_capability as analysis -from haiku.rag.capabilities.rag import create_capability as rag +from haiku.rag.capabilities.compaction import create_capability as compaction +from haiku.rag.capabilities.policy import create_capability as citation_policy agent = Agent( "openai:gpt-5", capabilities=[ - rag(db_path="my.lancedb"), analysis(db_path="my.lancedb"), + compaction(), + citation_policy(), ], ) ``` @@ -48,6 +54,6 @@ async with HaikuRAG("my.lancedb") as client: When dependencies expose a state dictionary, `AnalysisState` is stored under `"analysis"`. It contains the document filter, code execution log, searches, citations, and the `evidence` record of what was retrieved and cited per question. Searches and executions are cleared when a new question starts, and a resumed question keeps them; the filter, citation index and evidence record persist. -This capability does not alter the message history either. Register the [compaction capability](index.md#multi-turn-conversations) to compact earlier questions. +This capability does not alter the message history either. Register the [compaction capability](compaction.md) to compact earlier questions. The capability lazily opens both LanceDB and the sandbox only after it is loaded and a tool requires them. Resources close at the end of the agent run. diff --git a/docs/capabilities/compaction.md b/docs/capabilities/compaction.md new file mode 100644 index 00000000..46c6a038 --- /dev/null +++ b/docs/capabilities/compaction.md @@ -0,0 +1,43 @@ +# Evidence compaction capability + +`EvidenceCompactionCapability` keeps a multi-turn conversation from carrying every +search result it ever produced. Every question adds its evidence to the history, so +requests grow turn after turn, which degrades answers and can exceed a provider's +limits. + +Register it alongside an evidence capability: + +```python +from pydantic_ai import Agent +from haiku.rag.capabilities.compaction import create_capability as compaction +from haiku.rag.capabilities.rag import create_capability as rag + +agent = Agent( + "openai:gpt-5", + capabilities=[rag(db_path="my.lancedb"), compaction()], +) +``` + +It exposes no tools and takes no configuration. Registering it is the only switch: +leave it out and the transcript reaches the model untouched. + +## What it does + +On each request, evidence from earlier questions is replaced by the evidence those +questions actually cited. Cited text and cited page images are kept in full, grouped by +the question that cited them, and stay citable by the same chunk ids. Every other +earlier evidence return becomes a short receipt. The current question is untouched. + +Compaction rewrites the request, never the stored history, so `all_messages()` still +holds everything the run gathered. + +This reduces what a request carries. It does not bound it: retained evidence still +grows with the conversation. A host that needs more aggressive pruning can compact its +own requests further, on the wire only. + +## Resuming a question + +Resuming a question (deferred tool results, an interruption, a suspension) requires the +host to carry the capability state from the run being resumed, alongside the message +history. Without it the identity of the question in progress is unknowable, and the run +fails rather than silently treating it as a new question. diff --git a/docs/capabilities/index.md b/docs/capabilities/index.md index c7b582e5..15636fff 100644 --- a/docs/capabilities/index.md +++ b/docs/capabilities/index.md @@ -6,105 +6,50 @@ haiku.rag provides native [Pydantic AI capabilities](https://ai.pydantic.dev/cap |---|---| | [`RAGCapability`](rag.md) | Grounded document search and citations. | | [`AnalysisCapability`](analysis.md) | Corpus computation and structural analysis with sandboxed Python. | -| `EvidenceCompactionCapability` | Optional. Shrinking a conversation's history to the evidence that was cited. | -| `CitationPolicyCapability` | Optional. Requiring every answer to declare what grounds it. | +| [`EvidenceCompactionCapability`](compaction.md) | Optional. Shrinking a conversation's history to the evidence that was cited. | +| [`CitationPolicyCapability`](policy.md) | Optional. Requiring every answer to declare what grounds it. | The two evidence capabilities are deferred by default. An agent initially sees only their descriptions and the standard `load_capability` tool. Instructions and tools enter the model context only when the model loads a capability. ## Compose an agent +Pick one evidence capability, and add both optional capabilities to it: + ```python from pydantic_ai import Agent -from haiku.rag.capabilities.rag import create_capability - -rag = create_capability(db_path="my.lancedb") -agent = Agent("openai:gpt-5", capabilities=[rag]) - -result = await agent.run("What does the knowledge base say about X?") -print(result.output) -``` - -Attach both capabilities when an agent should choose between retrieval and computation: - -```python -from haiku.rag.capabilities.analysis import create_capability as analysis -from haiku.rag.capabilities.rag import create_capability as rag - -agent = Agent( - "openai:gpt-5", - capabilities=[rag(db_path="my.lancedb"), analysis(db_path="my.lancedb")], -) -``` - -## Multi-turn conversations - -Every question adds its search results to the history, so requests grow turn after -turn, and can degrade answers or exceed a provider's limits as they do. Register the -compaction capability to replace earlier questions' evidence with the evidence that -was actually cited: - -```python from haiku.rag.capabilities.compaction import create_capability as compaction -from haiku.rag.capabilities.rag import create_capability as rag - -agent = Agent( - "openai:gpt-5", - capabilities=[rag(db_path="my.lancedb"), compaction()], -) -``` - -Cited text and cited page images are kept in full, grouped by the question that -cited them, and stay citable by the same chunk ids. Everything else earlier becomes a -short receipt. Registering the capability is the only switch: leave it out and the -transcript reaches the model untouched. There is nothing to configure. - -Compaction rewrites the request, never the stored history, so `all_messages()` still -holds everything the run gathered. Retained evidence still grows with the -conversation — this reduces what a request carries, it does not bound it. A host that -needs more aggressive pruning can compact its own requests further, on the wire only. - -Resuming a question (deferred tool results, an interruption, a suspension) requires -the host to carry the capability state from the run being resumed, alongside the -message history. Without it the identity of the question in progress is unknowable -and the run fails rather than silently treating it as a new question. - -## Requiring citations - -Citing is always available and always recorded, but nothing requires it. Register the -citation policy capability to make every answer declare its grounding: - -```python from haiku.rag.capabilities.policy import create_capability as citation_policy from haiku.rag.capabilities.rag import create_capability as rag agent = Agent( "openai:gpt-5", - capabilities=[rag(db_path="my.lancedb"), citation_policy()], + capabilities=[ + rag(db_path="my.lancedb"), + compaction(), + citation_policy(), + ], ) + +result = await agent.run("What does the knowledge base say about X?") +print(result.output) ``` -An empty citation is a valid declaration: a model that finds nothing relevant calls -the cite tool with an empty list, which records the answer as *ungrounded* — distinct -from an answer that declared nothing at all. That distinction is what makes requiring -a declaration possible without forcing the model to invent grounding. +Swap `rag` for `analysis` for an analysis agent. Both optional capabilities work the +same way with either one, and neither exposes tools or takes configuration. -When a question ends undeclared, the model is asked once to record what grounded the -answer it already gave. It is not asked to change the answer. If the cite tool is no -longer available by then, or the question finishes undeclared anyway, it is recorded as -a violation in `CitationPolicyState` under `"citation_policy"`, since pointing a model -at a tool that is gone costs it retries. +!!! note "Register one evidence capability, not both" -What gets enforced is every answer in a conversation that has something to declare: -either this question retrieved evidence, or the conversation has already cited -something, which stays available to later answers. So a follow-up about evidence cited -earlier is enforced even though it searched nothing — that case is the reason the -capability exists. It also means that once anything has been cited, later turns are -enforced too, a greeting included; the model satisfies the policy by citing an empty -list, at the cost of one extra request. A conversation with neither a current-question -evidence outcome nor any earlier citation is not enforced. + `RAGCapability` and `AnalysisCapability` overlap. Both search the same corpus and + both register citations, so an agent holding both must choose between two + near-identical search tools, and its citations land in whichever capability it + happened to call. Each also carries its own request limit and its own search + budget, so registering both doubles what a question may spend. -Exactly one policy capability makes the decision, however many evidence capabilities -are registered, so two of them cannot each demand a citation for one answer. + Choose by what the questions need. `RAGCapability` answers questions from retrieved + passages. `AnalysisCapability` adds a Python sandbox and a document filesystem, for + questions that compute over many documents or read their structure, and it can + search too. If you need computation, register the analysis capability alone rather + than adding it to the RAG one. ## State diff --git a/docs/capabilities/policy.md b/docs/capabilities/policy.md new file mode 100644 index 00000000..de0b56df --- /dev/null +++ b/docs/capabilities/policy.md @@ -0,0 +1,49 @@ +# Citation policy capability + +`CitationPolicyCapability` requires every answer to declare what grounds it. Citing is +always available and always recorded without it, but nothing makes the model do it. + +Register it alongside an evidence capability: + +```python +from pydantic_ai import Agent +from haiku.rag.capabilities.policy import create_capability as citation_policy +from haiku.rag.capabilities.rag import create_capability as rag + +agent = Agent( + "openai:gpt-5", + capabilities=[rag(db_path="my.lancedb"), citation_policy()], +) +``` + +It exposes no tools and takes no configuration. Exactly one policy capability makes the +decision, however many evidence capabilities are registered, so two of them cannot each +demand a citation for one answer. + +## Declaring nothing is a valid answer + +A model that finds nothing relevant calls the cite tool with an empty list. That records +the answer as *ungrounded*, which is distinct from an answer that declared nothing at +all (*missing*). The distinction is what makes a declaration requirable without forcing +the model to invent grounding. + +## What happens when a question ends undeclared + +The model is asked once to record what grounded the answer it already gave. It is not +asked to change the answer. If the cite tool is no longer available by then, or the +question finishes undeclared anyway, the question is recorded in +`CitationPolicyState.violations` under the `"citation_policy"` state key. Pointing a +model at a tool that is gone costs it retries, so the capability records the failure +instead. + +## Which answers are enforced + +Every answer in a conversation that has something to declare: either this question +retrieved evidence, or the conversation has already cited something, which stays +available to later answers. A follow-up about evidence cited earlier is enforced even +though it searched nothing, which is the case the capability exists for. + +Once anything has been cited, later turns are enforced too, a greeting included. The +model satisfies the policy by citing an empty list, at the cost of one extra request. A +conversation with neither a current-question evidence outcome nor any earlier citation +is not enforced. diff --git a/docs/capabilities/rag.md b/docs/capabilities/rag.md index 35c3fbf2..ba3c3267 100644 --- a/docs/capabilities/rag.md +++ b/docs/capabilities/rag.md @@ -13,12 +13,23 @@ The distinct `rag_` prefix lets this capability coexist with analysis and other ## Create and compose +Register it on its own rather than alongside `AnalysisCapability`, which searches and +cites as well. See [Capabilities](index.md#compose-an-agent). + ```python from pydantic_ai import Agent -from haiku.rag.capabilities.rag import create_capability +from haiku.rag.capabilities.compaction import create_capability as compaction +from haiku.rag.capabilities.policy import create_capability as citation_policy +from haiku.rag.capabilities.rag import create_capability as rag -rag = create_capability(db_path="my.lancedb") -agent = Agent("openai:gpt-5", capabilities=[rag]) +agent = Agent( + "openai:gpt-5", + capabilities=[ + rag(db_path="my.lancedb"), + compaction(), + citation_policy(), + ], +) result = await agent.run("What safety equipment does the manual require?") print(result.output) @@ -49,7 +60,7 @@ State is ordinary application state; the capability does not depend on AG-UI. An ## Context management -This capability does not alter the message history. To stop long conversations resending old retrieved content, register the [compaction capability](index.md#multi-turn-conversations) alongside it. +This capability does not alter the message history. To stop long conversations resending old retrieved content, register the [compaction capability](compaction.md) alongside it. ## Domain context and vision diff --git a/docs/chat.md b/docs/chat.md index 091a3eb1..f5efdcad 100644 --- a/docs/chat.md +++ b/docs/chat.md @@ -61,13 +61,16 @@ Retrieval stays text-based; the images are sent to the model alongside your mess The default capability is `rag`. Enable `analysis` when the question needs computation, aggregation, comparison across documents, or section-scoped reading that a single search can't deliver: ```bash -# both capabilities (the agent routes between them) -haiku-rag chat -c rag -c analysis - -# analysis only +# analysis instead of rag haiku-rag chat -c analysis + +# both, which gives the model duplicate search and cite tools +haiku-rag chat -c rag -c analysis ``` +Prefer one. `analysis` searches and cites as well as computing, so pairing it with `rag` +duplicates tools and budgets. See [Capabilities](capabilities/index.md#compose-an-agent). + The `analysis` capability mounts every document as a virtual filesystem at `/documents/{id}/` (with `metadata.json`, `content.txt`, `items.jsonl`, and `toc.json`) and runs Python in a sandboxed interpreter with `search` and `list_documents` as awaitable functions. It's the right choice for questions like: - "How many of these documents mention X?" diff --git a/zensical.toml b/zensical.toml index 9ce0ec76..fbf59e7f 100644 --- a/zensical.toml +++ b/zensical.toml @@ -23,6 +23,8 @@ nav = [ "capabilities/index.md", { "RAG capability" = "capabilities/rag.md" }, { "Analysis capability" = "capabilities/analysis.md" }, + { "Evidence compaction" = "capabilities/compaction.md" }, + { "Citation policy" = "capabilities/policy.md" }, ] }, { Configure = [ "configuration/index.md", From 3ec72363bc6853b88b26ff9dc3c29e6906b894f8 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 13 Aug 2026 15:03:24 +0300 Subject: [PATCH 2/7] Let the record say whether a question is still being answered Resumption was inferred from the transcript, and no shape says it. A missing prompt is how UI adapters ask their first question as much as how pydantic-ai resumes one, so every AG-UI host failed on its first message. Reading the tail instead moved the error rather than fixing it: a settled structured answer ends with an output tool's return, indistinguishable from results delivered to a question still in progress, so following questions inherited the first one's identity. `CapabilityEvidenceRecord.in_progress` is now the authority. `begin_question` sets it, `after_run` clears it unless the run is only pausing for deferred work, and a run that raised never reaches `after_run`, which is what leaves an interrupted question claimable. The history is consulted only to catch a host that dropped the state of a question the model is unmistakably still owed. --- CHANGELOG.md | 4 +- .../haiku/rag/capabilities/_base.py | 51 ++--- .../haiku/rag/capabilities/ledger.py | 11 ++ tests/capabilities/test_capabilities.py | 177 +++++++++++++++++- tests/capabilities/test_citation_policy.py | 4 +- tests/capabilities/test_evidence_wire.py | 4 +- 6 files changed, 216 insertions(+), 35 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index cb1d2e1f..8d8bc2b2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,7 +6,7 @@ - `CitationPolicyCapability` (`haiku.rag.capabilities.policy.create_capability`): registering it requires every answer to declare its grounding, in any conversation that has something to declare — this question retrieved evidence, or something was cited earlier. A question that ends undeclared is sent back once to record what grounded the answer already given, and is recorded in `CitationPolicyState.violations` if it finishes undeclared regardless. A conversation with neither a current-question evidence outcome nor any earlier citation is not enforced. - `haiku.rag.capabilities.evidence.discover_evidence()` and `DiscoveredEvidence`, moved out of `compaction` so both optional capabilities share them. `RAGCapabilityBase.cite_available`. - `EvidenceCompactionCapability` (`haiku.rag.capabilities.compaction.create_capability`): registering it replaces earlier questions' evidence on the model request with the evidence that was cited, grouped by the question that cited it, cited page images re-attached, other earlier evidence returns reduced to a receipt. Requests only; `all_messages()` is untouched. No configuration. -- `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. +- `RAGState.evidence` / `AnalysisState.evidence` (`CapabilityEvidenceRecord`): which evidence a capability retrieved and cited, per question, keyed by message-count question identities and epochs, and whether that question is still being answered. `haiku.rag.capabilities.ledger.citation_status(records, question=...)` derives `missing` / `grounded` / `ungrounded` across capabilities. - `RAGCapabilityBase.evidence_tool_names()` and `get_picture_bytes()`. - `haiku.rag.tools.search.decode_picture()`. @@ -14,7 +14,7 @@ - `rag_cite` / `analysis_cite` accept an empty `chunk_ids`, recording the answer as ungrounded rather than failing the call, and the instructions no longer exempt a refusal or a corpus-level computation from citing. - `RAGCapability` and `AnalysisCapability` no longer rewrite the model request. Register `create_capability()` from `haiku.rag.capabilities.compaction` alongside them to keep earlier questions compacted. -- 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. +- Resuming a run (deferred tool results, an unfinished history tail) raises `RuntimeError` unless the host carries the capability state from the run being resumed. ### Removed diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index fe13d323..ab844ab8 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py @@ -6,13 +6,19 @@ from pathlib import Path from typing import Any, cast from pydantic import BaseModel -from pydantic_ai import ModelRetry, RunContext, ToolFailed +from pydantic_ai import ( + DeferredToolRequests, + ModelRetry, + RunContext, + ToolFailed, +) from pydantic_ai.capabilities import AbstractCapability from pydantic_ai.messages import ( InstructionPart, ModelMessage, ModelRequest, ModelResponse, + RetryPromptPart, ToolCallPart, ToolReturn, ) @@ -82,29 +88,21 @@ def _clear_invocation_state(state: BaseModel) -> None: value.clear() -def _is_resumption(prompt: Any, messages: list[ModelMessage]) -> bool: - """Whether this run continues a question rather than asking a new one. +def _awaits_the_model(messages: list[ModelMessage]) -> bool: + """Whether the history unmistakably leaves the model something to answer. - Two signals, either of which is enough, because getting this wrong hands the - model a notice where its own evidence should be: - - - no prompt: how pydantic-ai resumes for interruptions and suspensions. - - an unfinished tail: the history ends with a request the model has not - answered, or with a response whose tool calls have no returns yet. Deferred - tool results may arrive *with* a prompt, so the prompt alone is not enough. - - A settled history ends with the previous answer, so a genuinely new question - is not mistaken for a continuation. The framework's own first-new-message - index would be better than either signal, but it is not public here. + Used to validate what the record already says, never to decide it. Only two + shapes are unambiguous: a response whose tool calls have no returns, and a + retry the model has not answered. A trailing tool return is not one of them, + being both how a settled structured answer ends and how results reach a + question still in progress. """ - if prompt is None: - return True if not messages: return False last = messages[-1] - if isinstance(last, ModelRequest): - return True - return any(isinstance(part, ToolCallPart) for part in last.parts) + if isinstance(last, ModelResponse): + return any(isinstance(part, ToolCallPart) for part in last.parts) + return any(isinstance(part, RetryPromptPart) for part in last.parts) def _called_own_tool(messages: list[ModelMessage], tool_names: frozenset[str]) -> bool: @@ -154,11 +152,10 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): 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) - 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: + continuing = record.in_progress + if not continuing and _awaits_the_model(ctx.messages): raise RuntimeError( f"The {self.state_namespace} capability is resuming a question with " "no stored question identity. Capabilities cannot be added, removed " @@ -321,6 +318,16 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): async def after_run( self, ctx: RunContext[Any], *, result: AgentRunResult[Any] ) -> AgentRunResult[Any]: + """Close the question, unless the run is only pausing for deferred results. + + A run that raised never arrives here, which is what leaves an interrupted + question in progress for the resumption to claim. + """ + if self.state is not None and not isinstance( + result.output, DeferredToolRequests + ): + self._evidence_record().end_question() + self._sync_state() await self._close() return result diff --git a/haiku_rag_slim/haiku/rag/capabilities/ledger.py b/haiku_rag_slim/haiku/rag/capabilities/ledger.py index 778ae7ec..4599f672 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/ledger.py +++ b/haiku_rag_slim/haiku/rag/capabilities/ledger.py @@ -51,6 +51,11 @@ class CapabilityEvidenceRecord(BaseModel): overwritten by whichever of them synced its state last; merging happens in the transient views built by ``citation_status`` and the optional capabilities. + ``in_progress`` is whether the question this record names is still being + answered. It is the only authority on that: a transcript ending in a tool + return is a settled structured answer and a question awaiting its model + equally, so the shape of the history cannot decide it. + ``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 @@ -61,6 +66,7 @@ class CapabilityEvidenceRecord(BaseModel): occurrences: dict[str, EvidenceOccurrence] = Field(default_factory=dict) question: int | None = None + in_progress: bool = False latest_evidence_epoch: int = 0 declaration: CitationDeclaration | None = None @@ -110,9 +116,14 @@ class CapabilityEvidenceRecord(BaseModel): "one question from the next and are compared as recency." ) self.question = identity + self.in_progress = True self.latest_evidence_epoch = 0 self.declaration = None + def end_question(self) -> None: + """Mark the question answered, so the next run knows it is a new one.""" + self.in_progress = False + def note_evidence(self, epoch: int) -> None: """Record that the model has seen an evidence outcome. diff --git a/tests/capabilities/test_capabilities.py b/tests/capabilities/test_capabilities.py index cb58b2c7..f73fb7b6 100644 --- a/tests/capabilities/test_capabilities.py +++ b/tests/capabilities/test_capabilities.py @@ -4,7 +4,16 @@ from typing import Any, cast from unittest.mock import AsyncMock, patch import pytest -from pydantic_ai import Agent, DeferredToolResults, ModelRetry, RunContext, ToolFailed +from pydantic import BaseModel +from pydantic_ai import ( + Agent, + CallDeferred, + DeferredToolRequests, + DeferredToolResults, + ModelRetry, + RunContext, + ToolFailed, +) from pydantic_ai.messages import ( ModelRequest, ModelResponse, @@ -920,9 +929,9 @@ async def test_a_resumption_keeps_the_identity_of_the_question_in_progress( agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag]) deps = Deps( state={ - "rag": RAGState(evidence=CapabilityEvidenceRecord(question=7)).model_dump( - mode="json" - ) + "rag": RAGState( + evidence=CapabilityEvidenceRecord(question=7, in_progress=True) + ).model_dump(mode="json") } ) history = [ @@ -1159,10 +1168,12 @@ async def test_citing_without_searching_grounds_the_question(temp_db_path): @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. + """A seeded record says nothing about a question, so the history has to. - 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. + Seeding one is what a host does when it has no state to send. Its flag is + unset, so a history that unmistakably awaits the model means the host dropped + the state of a question in progress, and answering as question zero would + silently relabel it. """ rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False) @@ -1174,7 +1185,10 @@ async def test_a_host_seeded_record_does_not_pass_for_a_resumption(temp_db_path) with pytest.raises(RuntimeError, match="no stored question identity"): await agent.run( "carry on", - message_history=_in_flight_history(), + message_history=[ + *_in_flight_history(), + ModelResponse(parts=[ToolCallPart("external_tool", {}, "call-2")]), + ], deps=Deps(state={"rag": RAGState().model_dump(mode="json")}), ) @@ -1208,6 +1222,9 @@ async def test_a_resumption_keeps_the_evidence_the_question_already_gathered( interrupted = await agent.run("what does the supervisor do?", deps=deps) identity = _record(deps, "rag").question assert identity is not None + # A run that ends awaiting external work leaves the question in progress, + # which is what the resumption claims. See the deferred-request test. + deps.state["rag"]["evidence"]["in_progress"] = True await agent.run( deferred_tool_results=DeferredToolResults( calls={"call-2": "external result"} @@ -1280,3 +1297,147 @@ async def test_citing_nothing_after_citing_something_keeps_it_grounded(temp_db_p record = capability.state.evidence assert citation_status([record], question=0) == "grounded" + + +@pytest.mark.asyncio +async def test_a_promptless_run_on_a_settled_history_is_a_new_question(temp_db_path): + """AG-UI hosts never pass a prompt: the client's message is the history. + + Pydantic AI's UI adapter builds `message_history` from the frontend messages + and calls the agent without a prompt, so the run has no prompt *and* the + history ends with the user's own request. Reading either as a continuation + fails every AG-UI host on its first message. + """ + 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().model_dump(mode="json")}) + history: list[Any] = [ + ModelRequest(parts=[UserPromptPart("what does the manual say about masks?")]) + ] + + await agent.run(message_history=history, deps=deps) + + assert _record(deps, "rag").question == len(history) + + +@pytest.mark.asyncio +async def test_a_promptless_run_on_an_unfinished_tail_is_still_a_continuation( + temp_db_path, +): + """A suspended run resumes without a prompt, and must keep its question.""" + 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=3, in_progress=True) + ).model_dump(mode="json") + } + ) + + await agent.run(message_history=_in_flight_history(), deps=deps) + + assert _record(deps, "rag").question == 3 + + +@pytest.mark.asyncio +async def test_a_structured_answer_does_not_leave_the_question_in_progress( + temp_db_path, +): + """A settled run ends with a tool return, which says nothing about progress. + + Pydantic AI answers a structured `output_type` by calling an output tool, so the + history ends with a request carrying that tool's return. Reading the transcript + shape alone, that is indistinguishable from tool results delivered to a question + still being answered, and every following question inherited the first one's + identity. + """ + + class Answer(BaseModel): + text: str + + rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False) + + async def model(_messages, info): + return ModelResponse( + parts=[ToolCallPart(info.output_tools[0].name, {"text": "answer"})] + ) + + agent = Agent( + FunctionModel(model), + deps_type=Deps, + capabilities=[rag], + output_type=Answer, + ) + deps = Deps() + + first = await agent.run("first question", deps=deps) + first_identity = _record(deps, "rag").question + assert _record(deps, "rag").in_progress is False + + await agent.run( + "second question", deps=deps, message_history=list(first.all_messages()) + ) + + second_identity = _record(deps, "rag").question + assert first_identity == 0 + assert second_identity is not None and second_identity > 0 + + +@pytest.mark.asyncio +async def test_a_run_pausing_for_deferred_work_leaves_the_question_in_progress( + temp_db_path, +): + """The question is unfinished, so its resumption must find it claimable. + + A deferred tool call ends the run with `DeferredToolRequests` rather than an + answer. Closing the question here would let the resumption relabel it. + """ + rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False) + + async def model(_messages, _info): + return ModelResponse(parts=[ToolCallPart("external_tool", {})]) + + agent = Agent( + FunctionModel(model), + deps_type=Deps, + capabilities=[rag], + output_type=[str, DeferredToolRequests], + ) + + @agent.tool_plain + def external_tool() -> str: + raise CallDeferred + + deps = Deps() + + result = await agent.run("a question needing external work", deps=deps) + + assert isinstance(result.output, DeferredToolRequests) + assert _record(deps, "rag").in_progress is True + + +@pytest.mark.asyncio +async def test_an_answered_question_is_no_longer_in_progress(temp_db_path): + """The flag is what tells the next run it is asking something new.""" + 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("a question", deps=deps) + + record = _record(deps, "rag") + assert record.question == 0 + assert record.in_progress is False diff --git a/tests/capabilities/test_citation_policy.py b/tests/capabilities/test_citation_policy.py index b800f8c6..bcc3fde9 100644 --- a/tests/capabilities/test_citation_policy.py +++ b/tests/capabilities/test_citation_policy.py @@ -365,7 +365,9 @@ async def test_a_resumed_question_is_not_redirected_twice(temp_db_path): with patch.object(RAGCapability, "_search", stub_search): first = await agent.run("what does the supervisor do?", deps=deps) - # The same question again, continued rather than asked anew. + # The same question again, continued rather than asked anew: a run that + # ends awaiting external work leaves the question in progress. + deps.state["rag"]["evidence"]["in_progress"] = True await agent.run( deps=deps, message_history=[ diff --git a/tests/capabilities/test_evidence_wire.py b/tests/capabilities/test_evidence_wire.py index 3915fe17..6ad32aa4 100644 --- a/tests/capabilities/test_evidence_wire.py +++ b/tests/capabilities/test_evidence_wire.py @@ -299,11 +299,11 @@ def in_flight_history() -> list[Any]: def resuming_deps(question: int = 0) -> Deps: - """State as a resumption always finds it: the question already identified.""" + """State as a resumption always finds it: the question identified and unfinished.""" return Deps( state={ "rag": RAGState( - evidence=CapabilityEvidenceRecord(question=question) + evidence=CapabilityEvidenceRecord(question=question, in_progress=True) ).model_dump(mode="json") } ) From 21e261f60896843d49350311fc21ed216a531986 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 13 Aug 2026 15:03:51 +0300 Subject: [PATCH 3/7] Stop a partly resolved citation asking to be retried A cite call naming one good id and one mangled one registered the good one and then asked the model to cite again. A model that mangles ids obeys, mangles again, and the run dies on output retries with the answer lost: observed at 22 consecutive cite calls against gemma4-26b, with the citation policy registered to press for a declaration. The branch is only reachable once something has registered, so the answer already has grounding and there is nothing to ask for. --- CHANGELOG.md | 1 + haiku_rag_slim/haiku/rag/capabilities/_base.py | 7 +++++-- tests/capabilities/test_capabilities.py | 4 +++- 3 files changed, 9 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8d8bc2b2..e94f70d5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -22,6 +22,7 @@ ### Fixed +- `rag_cite` and `analysis_cite` no longer ask for another call when a call resolved some ids and not others. - `cross-encoder` reranking no longer ties the scores of strongly-relevant candidates, which left their order to the sort. Scores remain 0-1. - `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`. diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index ab844ab8..85052457 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py @@ -487,11 +487,14 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): resolved = {citation.chunk_id for citation in citations} unresolved = [cid for cid in missing if cid not in resolved] if unresolved: + # States the outcome without asking for another call. Reaching here + # means something registered, so the answer already has grounding: a + # model that keeps mangling ids would obey an invitation to retry + # until the run dies on output retries. return ( f"Registered {len(citations)} citation(s); " f"ignored {len(unresolved)} unresolvable id(s): " - f"{unresolved}. Copy chunk_ids verbatim from search " - "results and cite again." + f"{unresolved}, which were not verbatim from search results." ) return f"Registered {len(citations)} citation(s)." diff --git a/tests/capabilities/test_capabilities.py b/tests/capabilities/test_capabilities.py index f73fb7b6..3124104d 100644 --- a/tests/capabilities/test_capabilities.py +++ b/tests/capabilities/test_capabilities.py @@ -362,7 +362,9 @@ async def test_cite_reports_unresolved_ids_on_partial_success(temp_db_path): assert "Registered 1 citation(s)" in result assert "6.43" in result assert "6.51.2" in result - assert "verbatim" in result + # Never ask for another call here: something did register, and a model that + # keeps mangling ids obeys the ask until the run dies on output retries. + assert "again" not in result assert capability.state.citations == ["chunk-1"] From 771ac9c96c7679db1972663e55274f5cacd3bb40 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 13 Aug 2026 15:04:05 +0300 Subject: [PATCH 4/7] Register the optional capabilities where agents are composed The README feature list and the overview stopped at the analysis capability. Both examples and the app backend composed agents without the capabilities the documentation recommends alongside an evidence capability. custom_agent.py ran each input as an independent agent run, so it needed a state dict and a carried history before compaction could mean anything there: without state the evidence record is empty, and earlier evidence would reduce to receipts retaining nothing. --- README.md | 2 ++ app/backend/main.py | 8 ++++++-- docs/overview.md | 2 +- examples/custom_agent.py | 33 +++++++++++++++++++++++++++------ examples/custom_agent_agui.py | 6 +++++- 5 files changed, 41 insertions(+), 10 deletions(-) diff --git a/README.md b/README.md index 8680d460..869b3bfd 100644 --- a/README.md +++ b/README.md @@ -15,6 +15,8 @@ Agentic RAG built on [LanceDB](https://lancedb.com/), [Pydantic AI](https://ai.p - **Vision QA** — Vision-capable models receive figure bytes alongside chunk text; attach your own images to questions in `ask`, `analyze`, MCP, and the chat TUI - **Reranking** — local cross-encoders, Cohere, Zero Entropy, or vLLM - **Analysis capability** — Complex analytical tasks via sandboxed Python code execution (aggregation, computation, multi-document analysis) +- **Evidence compaction** — Optional capability that replaces earlier questions' search results on the request with the evidence they cited, so long conversations stop resending everything they retrieved +- **Citation policy** — Optional capability that requires every answer to declare what grounds it, including declaring that nothing does - **Conversational RAG** — Chat TUI and web application for multi-turn conversations with session memory - **Document structure** — Stores full [DoclingDocument](https://docling-project.github.io/docling/concepts/docling_document/), enabling structure-aware context expansion - **Multiple providers** — Embeddings: Ollama, OpenAI, VoyageAI, Cohere, LM Studio, vLLM (multimodal via `multimodal: true` on vLLM/VoyageAI/Cohere). QA: any model supported by Pydantic AI diff --git a/app/backend/main.py b/app/backend/main.py index 252e938c..109fb659 100644 --- a/app/backend/main.py +++ b/app/backend/main.py @@ -21,6 +21,9 @@ from starlette.routing import Route from haiku.rag.capabilities.compaction import ( create_capability as create_compaction, ) +from haiku.rag.capabilities.policy import ( + create_capability as create_citation_policy, +) from haiku.rag.capabilities.rag import AGENT_PREAMBLE, RAGState, create_capability from haiku.rag.client import HaikuRAG from haiku.rag.config import load_yaml_config @@ -85,8 +88,9 @@ agent = Agent( get_model(Config.qa.model, Config), instructions=AGENT_PREAMBLE, # Conversations here are multi-turn, so earlier questions are reduced to the - # evidence they cited rather than carried whole. - capabilities=[capability, create_compaction()], + # evidence they cited rather than carried whole, and every answer declares + # what grounds it so the UI can show citations for all of them. + capabilities=[capability, create_compaction(), create_citation_policy()], deps_type=AppDeps, ) diff --git a/docs/overview.md b/docs/overview.md index 83170fb9..8655a6aa 100644 --- a/docs/overview.md +++ b/docs/overview.md @@ -27,7 +27,7 @@ The chat TUI is one way to interact with the database. `haiku-rag ask` and `haik **Search.** Hybrid retrieval (vector + full-text with reciprocal rank fusion), optional cross-encoder reranking, structure-aware context expansion. Image-as-query and cross-modal retrieval when configured with a multimodal embedder. -**Answer.** RAG capability with citations including page numbers, section headings, and visual grounding. Vision-capable models receive figure bytes alongside chunk text. Analysis capability with a sandboxed Python interpreter for aggregation and computation across documents. +**Answer.** RAG capability with citations including page numbers, section headings, and visual grounding. Vision-capable models receive figure bytes alongside chunk text. Analysis capability with a sandboxed Python interpreter for aggregation and computation across documents. Optional capabilities compact a long conversation down to the evidence it cited, and require every answer to declare its grounding. **Integrate.** Use it from Python, the CLI, the [MCP server](mcp.md), or through composable native Pydantic AI [capabilities](capabilities/index.md). diff --git a/examples/custom_agent.py b/examples/custom_agent.py index 99ecadf8..0bf2adb7 100644 --- a/examples/custom_agent.py +++ b/examples/custom_agent.py @@ -1,6 +1,7 @@ """Custom agent using the native haiku.rag RAG capability. -Demonstrates composing a native Pydantic AI capability into an agent. +Demonstrates composing native Pydantic AI capabilities into an agent, and what a +multi-turn conversation needs to carry between runs. Requirements: - An Ollama instance running locally (default embedder) @@ -13,21 +14,40 @@ Usage: import asyncio import sys +from dataclasses import dataclass, field from pathlib import Path +from typing import Any from pydantic_ai import Agent +from pydantic_ai.messages import ModelMessage -from haiku.rag.capabilities.rag import create_capability +from haiku.rag.capabilities.compaction import create_capability as compaction +from haiku.rag.capabilities.policy import create_capability as citation_policy +from haiku.rag.capabilities.rag import create_capability as rag + + +@dataclass +class Deps: + state: dict[str, Any] = field(default_factory=dict) async def main(db_path: str) -> None: - capability = create_capability(db_path=Path(db_path), defer_loading=False) - agent = Agent( "anthropic:claude-haiku-4-5-20251001", - capabilities=[capability], + capabilities=[ + rag(db_path=Path(db_path), defer_loading=False), + compaction(), + citation_policy(), + ], + deps_type=Deps, ) + # One state dict and one history for the whole session. The capabilities read + # both: the state holds what was retrieved and cited, and the message counts + # are how they tell one question from the next. + deps = Deps() + messages: list[ModelMessage] = [] + print("Custom agent ready. Ctrl+C to exit.\n") while True: try: @@ -38,7 +58,8 @@ async def main(db_path: str) -> None: if not user_input: continue - result = await agent.run(user_input) + result = await agent.run(user_input, deps=deps, message_history=messages) + messages = list(result.all_messages()) print(f"\nAgent: {result.output}\n") diff --git a/examples/custom_agent_agui.py b/examples/custom_agent_agui.py index 850c2f7f..78f7b08e 100644 --- a/examples/custom_agent_agui.py +++ b/examples/custom_agent_agui.py @@ -26,6 +26,8 @@ from starlette.requests import Request from starlette.responses import JSONResponse, Response, StreamingResponse from starlette.routing import Route +from haiku.rag.capabilities.compaction import create_capability as compaction +from haiku.rag.capabilities.policy import create_capability as citation_policy from haiku.rag.capabilities.rag import RAGState, create_capability db_path = os.environ.get("DB_PATH") @@ -45,7 +47,9 @@ class AppDeps: agent = Agent( "anthropic:claude-haiku-4-5-20251001", - capabilities=[capability], + # The client returns the state snapshot with every run, so earlier questions are + # reduced to the evidence they cited and every answer declares its grounding. + capabilities=[capability, compaction(), citation_policy()], deps_type=AppDeps, ) From 53bbf52697e7529e5b75b6f83089dfcb55272181 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 13 Aug 2026 15:20:02 +0300 Subject: [PATCH 5/7] Round-trip the capability's whole state through the browser The session store rebuilt the rag namespace from four known keys, so the evidence record never survived a turn, let alone a reload from localStorage. Compaction then ran with an empty ledger: earlier evidence was replaced by receipts retaining nothing, while the citations already in citation_index kept the UI looking correct. The UI still names the fields it reads, but everything else in the namespace passes through untouched, and seeding the namespace no longer replaces sibling namespaces. --- app/frontend/components/Chat.tsx | 1 + app/frontend/lib/sessionStorage.ts | 7 ++++++- 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/app/frontend/components/Chat.tsx b/app/frontend/components/Chat.tsx index 6f1ed583..30f1f6a7 100644 --- a/app/frontend/components/Chat.tsx +++ b/app/frontend/components/Chat.tsx @@ -321,6 +321,7 @@ function ChatContentInner({ const session = getSession(sessionId); // Seed state for the capability; the backend replaces it after each run. agent.setState({ + ...agent.state, [AGUI_STATE_KEY]: normalizeRAGState(session?.ragState), }); if (session && session.messages.length > 0) { diff --git a/app/frontend/lib/sessionStorage.ts b/app/frontend/lib/sessionStorage.ts index 02caae19..69d8ac7c 100644 --- a/app/frontend/lib/sessionStorage.ts +++ b/app/frontend/lib/sessionStorage.ts @@ -11,12 +11,16 @@ export interface Citation { doc_item_refs?: string[]; } -// Matches RAGState from the backend capability. +// Matches RAGState from the backend capability. The fields named here are the +// ones this UI reads; the capability owns the rest of its namespace, including +// the evidence record that compaction builds its capsule from, so the state has +// to round-trip whole rather than be rebuilt from known keys. export interface RAGState { citation_index: Record; citations: string[]; document_filter: string | null; searches: Record; + [key: string]: unknown; } export interface StoredMessage { @@ -40,6 +44,7 @@ const ACTIVE_SESSION_KEY = "haiku.rag.activeSession"; export function normalizeRAGState(state?: Partial): RAGState { return { + ...state, citation_index: state?.citation_index ?? {}, citations: state?.citations ?? [], document_filter: state?.document_filter ?? null, From e1dd8517f9f05fc0554a321efc7d1fdfe998542d Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 13 Aug 2026 15:46:23 +0300 Subject: [PATCH 6/7] Refuse to compact evidence the host kept no record of Both optional capabilities read what earlier questions retrieved and cited from the capability's state, so a host that carries only the message history hands every run an empty record. Compaction then replaced the earlier evidence with receipts and retained nothing, and the loss was invisible: the citations the host already displayed were still there. It now refuses when it finds evidence from an earlier question and no record of what that question cited. `state_carried` reaches the optional capabilities through discovery, so the refusal distinguishes a host that never carries state from a question that simply cited nothing. The documentation taught the pattern that breaks: the compose example is now stateful and the requirement is stated where each capability is introduced. The app's browser storage was doing exactly this, keeping only the fields the UI reads. It now persists the whole namespace map, so the citation policy's violations survive a reload as well as the evidence record. --- app/frontend/components/Chat.tsx | 21 +++--- app/frontend/lib/sessionStorage.ts | 29 ++++++-- docs/capabilities/compaction.md | 6 ++ docs/capabilities/index.md | 28 +++++++- docs/capabilities/policy.md | 11 ++- .../haiku/rag/capabilities/_base.py | 9 +++ .../haiku/rag/capabilities/compaction.py | 15 +++- .../haiku/rag/capabilities/evidence.py | 2 + tests/capabilities/test_evidence_capsule.py | 1 + tests/capabilities/test_evidence_wire.py | 68 ++++++++++++++++++- 10 files changed, 163 insertions(+), 27 deletions(-) diff --git a/app/frontend/components/Chat.tsx b/app/frontend/components/Chat.tsx index 30f1f6a7..2d0770c5 100644 --- a/app/frontend/components/Chat.tsx +++ b/app/frontend/components/Chat.tsx @@ -20,6 +20,8 @@ import { import { FilterIcon } from "../lib/icons"; import type { RAGState } from "../lib/sessionStorage"; import { + AGUI_STATE_KEY, + agentStateOf, createSession, getActiveSessionId, getLatestCitations, @@ -32,10 +34,7 @@ import DbInfo from "./DbInfo"; import DocumentFilter from "./DocumentFilter"; import SessionManager from "./SessionManager"; -// Must match RAGCapability.state_namespace. -const AGUI_STATE_KEY = "rag"; - -// AG-UI state is namespaced under AGUI_STATE_KEY +// AG-UI state is namespaced under AGUI_STATE_KEY (see sessionStorage). interface AgentState { [AGUI_STATE_KEY]?: RAGState; } @@ -319,11 +318,10 @@ function ChatContentInner({ useEffect(() => { if (agent.messages.length > 0) return; const session = getSession(sessionId); - // Seed state for the capability; the backend replaces it after each run. - agent.setState({ - ...agent.state, - [AGUI_STATE_KEY]: normalizeRAGState(session?.ragState), - }); + // Seed state for the capabilities; the backend replaces it after each run. + // The whole namespace map goes back, not just the fields this UI reads: + // compaction and the citation policy read what earlier questions recorded. + agent.setState({ ...agent.state, ...agentStateOf(session ?? undefined) }); if (session && session.messages.length > 0) { // biome-ignore lint/suspicious/noExplicitAny: AG-UI Message type is a broad union agent.setMessages(session.messages as any[]); @@ -336,13 +334,10 @@ function ChatContentInner({ // biome-ignore lint/correctness/useExhaustiveDependencies: JSON.stringify tracks content changes useEffect(() => { if (sessionId && agent.messages.length > 0) { - const currentRagState = normalizeRAGState( - (agent.state as AgentState)?.[AGUI_STATE_KEY], - ); updateSessionMessages( sessionId, serializeMessages(agent.messages), - currentRagState, + (agent.state ?? {}) as Record, ); } }, [JSON.stringify(agent.messages), ragState, sessionId]); diff --git a/app/frontend/lib/sessionStorage.ts b/app/frontend/lib/sessionStorage.ts index 69d8ac7c..63642444 100644 --- a/app/frontend/lib/sessionStorage.ts +++ b/app/frontend/lib/sessionStorage.ts @@ -34,11 +34,32 @@ export interface StoredSession { id: string; title: string; messages: StoredMessage[]; - ragState: RAGState; + // The whole AG-UI state. The rag namespace is not the only one a capability + // writes: the citation policy records violations beside it. + agentState: AgentState; + // Sessions stored before agentState existed. + ragState?: RAGState; createdAt: string; updatedAt: string; } +export const AGUI_STATE_KEY = "rag"; + +export type AgentState = Record; + +// Reads the rag namespace out of a stored session, whichever way it was stored. +export function ragStateOf(session?: StoredSession): RAGState { + const namespaced = session?.agentState?.[AGUI_STATE_KEY] as + | Partial + | undefined; + return normalizeRAGState(namespaced ?? session?.ragState); +} + +// The state to seed an agent with when a session is resumed. +export function agentStateOf(session?: StoredSession): AgentState { + return session?.agentState ?? { [AGUI_STATE_KEY]: ragStateOf(session) }; +} + const SESSIONS_KEY = "haiku.rag.sessions"; const ACTIVE_SESSION_KEY = "haiku.rag.activeSession"; @@ -86,7 +107,7 @@ export function createSession(): StoredSession { id: crypto.randomUUID(), title: "New Session", messages: [], - ragState: normalizeRAGState(), + agentState: { [AGUI_STATE_KEY]: normalizeRAGState() }, createdAt: now, updatedAt: now, }; @@ -111,7 +132,7 @@ export function saveSession(session: StoredSession): void { export function updateSessionMessages( id: string, messages: StoredMessage[], - ragState: RAGState, + agentState: AgentState, ): void { const sessions = getAllSessions(); const idx = sessions.findIndex((s) => s.id === id); @@ -119,7 +140,7 @@ export function updateSessionMessages( const session = sessions[idx]; session.messages = messages; - session.ragState = ragState; + session.agentState = agentState; session.updatedAt = new Date().toISOString(); // Derive title from first user message diff --git a/docs/capabilities/compaction.md b/docs/capabilities/compaction.md index 46c6a038..e0baeb25 100644 --- a/docs/capabilities/compaction.md +++ b/docs/capabilities/compaction.md @@ -21,6 +21,12 @@ agent = Agent( It exposes no tools and takes no configuration. Registering it is the only switch: leave it out and the transcript reaches the model untouched. +The host must carry the capability state between runs, alongside the message +history: the capsule is built from what earlier questions recorded there. Given +only a message history, every run starts from an empty record, and compaction +refuses rather than replace evidence it cannot retain. See +[Compose an agent](index.md#compose-an-agent) for the shape. + ## What it does On each request, evidence from earlier questions is replaced by the evidence those diff --git a/docs/capabilities/index.md b/docs/capabilities/index.md index 15636fff..f4b0e105 100644 --- a/docs/capabilities/index.md +++ b/docs/capabilities/index.md @@ -16,11 +16,22 @@ The two evidence capabilities are deferred by default. An agent initially sees o Pick one evidence capability, and add both optional capabilities to it: ```python +from dataclasses import dataclass, field +from typing import Any + from pydantic_ai import Agent +from pydantic_ai.messages import ModelMessage + from haiku.rag.capabilities.compaction import create_capability as compaction from haiku.rag.capabilities.policy import create_capability as citation_policy from haiku.rag.capabilities.rag import create_capability as rag + +@dataclass +class Deps: + state: dict[str, Any] = field(default_factory=dict) + + agent = Agent( "openai:gpt-5", capabilities=[ @@ -28,12 +39,27 @@ agent = Agent( compaction(), citation_policy(), ], + deps_type=Deps, ) -result = await agent.run("What does the knowledge base say about X?") +# One Deps and one history for the conversation: the capabilities read both. +deps = Deps() +history: list[ModelMessage] = [] + +result = await agent.run("What does the knowledge base say about X?", deps=deps, message_history=history) +history = list(result.all_messages()) print(result.output) ``` +!!! warning "Both optional capabilities need the host to carry state" + + They read what earlier questions retrieved and cited from the capability's + state, so the host must expose a `state` dict on its agent dependencies and + hand the same dict back on every run of a conversation, alongside the message + history. With only the message history, every run starts from an empty record: + compaction refuses rather than replace evidence it cannot retain, and the + citation policy cannot enforce a follow-up about evidence cited earlier. + Swap `rag` for `analysis` for an analysis agent. Both optional capabilities work the same way with either one, and neither exposes tools or takes configuration. diff --git a/docs/capabilities/policy.md b/docs/capabilities/policy.md index de0b56df..19e26728 100644 --- a/docs/capabilities/policy.md +++ b/docs/capabilities/policy.md @@ -16,9 +16,14 @@ agent = Agent( ) ``` -It exposes no tools and takes no configuration. Exactly one policy capability makes the -decision, however many evidence capabilities are registered, so two of them cannot each -demand a citation for one answer. +It exposes no tools and takes no configuration. Exactly one policy capability makes +the decision, however many evidence capabilities are registered, so two of them +cannot each demand a citation for one answer. + +The host must carry the capability state between runs, alongside the message +history. Enforcement reads what the conversation has already cited, so without it +a follow-up about evidence cited earlier goes unenforced. See +[Compose an agent](index.md#compose-an-agent) for the shape. ## Declaring nothing is a valid answer diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index 85052457..b9f0ff60 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py @@ -135,6 +135,13 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): request_count: int = field(default=0, repr=False) grace_requests_used: int = field(default=0, repr=False) epoch: int = field(default=0, repr=False) + state_carried: bool = field(default=False, repr=False) + """Whether the host handed back a record a previous question had stamped. + + False on a first question, and equally on every question of a host that does + not carry state between runs. Capabilities that need the record to mean + anything across questions read it to refuse rather than act on nothing. + """ async def for_run(self, ctx: RunContext[Any]) -> "RAGCapabilityBase[StateT]": """Start a run's own copy, and settle which question it is answering. @@ -155,6 +162,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): state = self.state_type.model_validate(raw_state or {}) record = cast(CapabilityEvidenceRecord, cast(Any, state).evidence) continuing = record.in_progress + state_carried = record.question is not None if not continuing and _awaits_the_model(ctx.messages): raise RuntimeError( f"The {self.state_namespace} capability is resuming a question with " @@ -176,6 +184,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): request_count=0, grace_requests_used=0, epoch=0, + state_carried=state_carried, ) run_capability._sync_state() return run_capability diff --git a/haiku_rag_slim/haiku/rag/capabilities/compaction.py b/haiku_rag_slim/haiku/rag/capabilities/compaction.py index 85e5170e..6b1decbb 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/compaction.py +++ b/haiku_rag_slim/haiku/rag/capabilities/compaction.py @@ -321,14 +321,23 @@ 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: await self._build_once(ctx, evidence) request_context.messages = compact_history( request_context.messages, boundary=boundary, - owned_tools=frozenset().union( - *(found.tool_names for found in evidence) - ), + owned_tools=owned_tools, capsule_text=self.capsule.text, capsule_images=self.images, ) diff --git a/haiku_rag_slim/haiku/rag/capabilities/evidence.py b/haiku_rag_slim/haiku/rag/capabilities/evidence.py index 75a61cc4..49f8f385 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/evidence.py +++ b/haiku_rag_slim/haiku/rag/capabilities/evidence.py @@ -22,6 +22,7 @@ class DiscoveredEvidence: citations: Mapping[str, Citation] tool_names: frozenset[str] cite_available: bool + state_carried: bool def discover_evidence(ctx: RunContext[Any]) -> list[DiscoveredEvidence]: @@ -40,6 +41,7 @@ def discover_evidence(ctx: RunContext[Any]) -> list[DiscoveredEvidence]: citations=cast(Any, capability.state).citation_index, tool_names=frozenset(capability.evidence_tool_names()), cite_available=capability.cite_available, + state_carried=capability.state_carried, ) for capability in ctx.capabilities.values() if isinstance(capability, RAGCapabilityBase) diff --git a/tests/capabilities/test_evidence_capsule.py b/tests/capabilities/test_evidence_capsule.py index 981db4c3..056accf3 100644 --- a/tests/capabilities/test_evidence_capsule.py +++ b/tests/capabilities/test_evidence_capsule.py @@ -68,6 +68,7 @@ def discovered( return DiscoveredEvidence( capability=capability, record=record, + state_carried=True, citations={ chunk_id: citation( chunk_id, diff --git a/tests/capabilities/test_evidence_wire.py b/tests/capabilities/test_evidence_wire.py index 6ad32aa4..7e588ca1 100644 --- a/tests/capabilities/test_evidence_wire.py +++ b/tests/capabilities/test_evidence_wire.py @@ -287,6 +287,17 @@ def rag_and_compactor(temp_db_path): ) +def settled_deps(question: int = 0) -> Deps: + """State as a host carrying it has it: an earlier question, answered.""" + return Deps( + state={ + "rag": RAGState( + evidence=CapabilityEvidenceRecord(question=question, in_progress=False) + ).model_dump(mode="json") + } + ) + + def in_flight_history() -> list[Any]: """A question already asked and searched, still awaiting its answer.""" return [ @@ -322,7 +333,9 @@ async def test_without_the_compactor_the_history_is_untouched(temp_db_path): agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag]) settled = [*in_flight_history(), ModelResponse(parts=[TextPart("first answer")])] - await agent.run("a different question", deps=Deps(), message_history=settled) + await agent.run( + "a different question", deps=settled_deps(), message_history=settled + ) assert returns_of(wire[-1]) == ["EVIDENCE FOR THE LIVE TURN"] @@ -341,7 +354,9 @@ async def test_with_the_compactor_a_new_question_compacts_the_previous_one( agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag, compactor]) settled = [*in_flight_history(), ModelResponse(parts=[TextPart("first answer")])] - await agent.run("a different question", deps=Deps(), message_history=settled) + await agent.run( + "a different question", deps=settled_deps(), message_history=settled + ) assert returns_of(wire[-1]) == [RECEIPT] @@ -509,7 +524,8 @@ async def test_the_capsule_is_built_once_per_request_and_again_for_the_next( ctx = RunContext( deps=deps, model=TestModel(), usage=RunUsage(), run_id="run-1", run_step=1 ) - run_rag = await rag.for_run(ctx) + # A host carrying state, which is what compaction requires of one. + run_rag = replace(await rag.for_run(ctx), state_carried=True) run_compactor = await compactor.for_run(ctx) cast(Any, run_rag.state).evidence.begin_question(4) ctx = replace(ctx, capabilities={"rag": run_rag, "compaction": run_compactor}) @@ -727,3 +743,49 @@ def test_the_capsule_is_attached_beside_the_newest_return_of_that_request(): assert images_of(compacted) == [fresh] assert returns_of(compacted) == [RECEIPT, "CAPSULE"] + + +@pytest.mark.asyncio +async def test_compaction_refuses_to_strip_evidence_it_cannot_replace(temp_db_path): + """A host that does not carry state has no record to build a capsule from. + + Compacting anyway replaces the earlier evidence with receipts and retains + nothing, so the model loses what it cited and the loss is invisible: the + citations the host already displayed are still there. + """ + rag, compactor = rag_and_compactor(temp_db_path) + + async def model(_messages, _info): # pragma: no cover - never reached + return ModelResponse(parts=[TextPart("answer")]) + + agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag, compactor]) + history = answered_question("an earlier question", evidence="EVIDENCE TO LOSE") + + with pytest.raises(RuntimeError, match="carry the capability state"): + await agent.run("a follow-up", deps=Deps(), message_history=history) + + +@pytest.mark.asyncio +async def test_compaction_proceeds_for_a_host_that_carries_state(temp_db_path): + """The same history, with the record the earlier question left behind.""" + rag, compactor = rag_and_compactor(temp_db_path) + 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, compactor]) + history = answered_question("an earlier question", evidence="EVIDENCE TO LOSE") + + carried = Deps( + state={ + "rag": RAGState( + evidence=CapabilityEvidenceRecord(question=0, in_progress=False) + ).model_dump(mode="json") + } + ) + + await agent.run("a follow-up", deps=carried, message_history=history) + + assert returns_of(wire[-1]) == [RECEIPT] From d1e3b828f1bbc4508379383ff5f826942e5d713e Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 13 Aug 2026 15:57:33 +0300 Subject: [PATCH 7/7] Judge each capability's evidence by its own record The guard accepted one carried record as covering the whole request, so a host retaining only the RAG namespace had its earlier analysis evidence replaced by receipts and lost, with the RAG record making the loss look accounted for. Each capability is now judged on its own, and only when its own evidence is at stake: another capability's record says nothing about this one's, and requiring every record would stop a host that registers both capabilities and only ever uses one. --- .../haiku/rag/capabilities/compaction.py | 40 ++++++--- tests/capabilities/test_evidence_wire.py | 88 +++++++++++++++++++ 2 files changed, 118 insertions(+), 10 deletions(-) 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]