diff --git a/CHANGELOG.md b/CHANGELOG.md index cb1d2e1f..e94f70d5 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 @@ -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/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/app/frontend/components/Chat.tsx b/app/frontend/components/Chat.tsx index 6f1ed583..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,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, ); } }, [JSON.stringify(agent.messages), ragState, sessionId]); diff --git a/app/frontend/lib/sessionStorage.ts b/app/frontend/lib/sessionStorage.ts index 02caae19..63642444 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 { @@ -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; + +// 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"; export function normalizeRAGState(state?: Partial): 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 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..e0baeb25 --- /dev/null +++ b/docs/capabilities/compaction.md @@ -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. diff --git a/docs/capabilities/index.md b/docs/capabilities/index.md index c7b582e5..f4b0e105 100644 --- a/docs/capabilities/index.md +++ b/docs/capabilities/index.md @@ -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 diff --git a/docs/capabilities/policy.md b/docs/capabilities/policy.md new file mode 100644 index 00000000..19e26728 --- /dev/null +++ b/docs/capabilities/policy.md @@ -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. 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/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, ) diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index fe13d323..b9f0ff60 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: @@ -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)." diff --git a/haiku_rag_slim/haiku/rag/capabilities/compaction.py b/haiku_rag_slim/haiku/rag/capabilities/compaction.py index 85e5170e..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: @@ -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, ) 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/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..3124104d 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, @@ -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 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_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 3915fe17..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, @@ -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] 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",