danvics/haiku.rag
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge pull request #541 from ggozad/fix/new-question-detection

Browse source 
Decide a question's lifecycle from state, not from the transcript
This commit is contained in:
Yiorgis Gozadinos 2026-08-13 16:09:27 +03:00 committed by GitHub
commit c5ecf718d9
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
23 changed files with 689 additions and 162 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
CHANGELOG.md
View file

@ -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
@ -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`.

2
README.md
View file

@ -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

8
app/backend/main.py
View file

@ -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,
)

20
app/frontend/components/Chat.tsx
View file

@ -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,10 +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({
[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[]);
@ -335,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<string, unknown>,
);
}
}, [JSON.stringify(agent.messages), ragState, sessionId]);

36
app/frontend/lib/sessionStorage.ts
View file

@ -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<string, Citation>;
citations: string[];
document_filter: string | null;
searches: Record<string, unknown[]>;
[key: string]: unknown;
}
export interface StoredMessage {
@ -30,16 +34,38 @@ 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<string, unknown>;
// 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<RAGState>
| 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";
export function normalizeRAGState(state?: Partial<RAGState>): RAGState {
return {
...state,
citation_index: state?.citation_index ?? {},
citations: state?.citations ?? [],
document_filter: state?.document_filter ?? null,
@ -81,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,
};
@ -106,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);
@ -114,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

14
docs/capabilities/analysis.md
View file

@ -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.

49
docs/capabilities/compaction.md Normal file
View file

@ -0,0 +1,49 @@
# 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.
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
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.

129
docs/capabilities/index.md
View file

@ -6,105 +6,76 @@ 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 dataclasses import dataclass, field
from typing import Any
from pydantic_ai import Agent
from haiku.rag.capabilities.rag import create_capability
from pydantic_ai.messages import ModelMessage
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
@dataclass
class Deps:
state: dict[str, Any] = field(default_factory=dict)
agent = Agent(
"openai:gpt-5",
capabilities=[rag(db_path="my.lancedb"), citation_policy()],
capabilities=[
rag(db_path="my.lancedb"),
compaction(),
citation_policy(),
],
deps_type=Deps,
)
# 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)
```
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.
!!! warning "Both optional capabilities need the host to carry state"
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.
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.
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.
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.
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.
!!! note "Register one evidence capability, not both"
`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.
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

54
docs/capabilities/policy.md Normal file
View file

@ -0,0 +1,54 @@
# 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.
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
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.

19
docs/capabilities/rag.md
View file

@ -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

11
docs/chat.md
View file

@ -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?"

2
docs/overview.md
View file

@ -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).

33
examples/custom_agent.py
View file

@ -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")

6
examples/custom_agent_agui.py
View file

@ -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,
)

67
haiku_rag_slim/haiku/rag/capabilities/_base.py
View file

@ -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:
@ -137,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.
@ -154,11 +159,11 @@ 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
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 "
"no stored question identity. Capabilities cannot be added, removed "
@ -179,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
@ -321,6 +327,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
@ -480,11 +496,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)."

35
haiku_rag_slim/haiku/rag/capabilities/compaction.py
View file

@ -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:
@ -321,14 +347,17 @@ 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:
_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(
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,
)

2
haiku_rag_slim/haiku/rag/capabilities/evidence.py
View file

@ -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)

11
haiku_rag_slim/haiku/rag/capabilities/ledger.py
View file

@ -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.

181
tests/capabilities/test_capabilities.py
View file

@ -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,
@ -353,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"]
@ -920,9 +931,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 +1170,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 +1187,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 +1224,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 +1299,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

4
tests/capabilities/test_citation_policy.py
View file

@ -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=[

1
tests/capabilities/test_evidence_capsule.py
View file

@ -68,6 +68,7 @@ def discovered(
return DiscoveredEvidence(
capability=capability,
record=record,
state_carried=True,
citations={
chunk_id: citation(
chunk_id,

160
tests/capabilities/test_evidence_wire.py
View file

@ -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,
@ -287,6 +288,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 [
@ -299,11 +311,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")
}
)
@ -322,7 +334,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 +355,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 +525,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 +744,136 @@ 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]
@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]

2
zensical.toml
View file

@ -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",