diff --git a/app/frontend/components/Chat.tsx b/app/frontend/components/Chat.tsx index 30f1f6a7..2d0770c5 100644 --- a/app/frontend/components/Chat.tsx +++ b/app/frontend/components/Chat.tsx @@ -20,6 +20,8 @@ import { import { FilterIcon } from "../lib/icons"; import type { RAGState } from "../lib/sessionStorage"; import { + AGUI_STATE_KEY, + agentStateOf, createSession, getActiveSessionId, getLatestCitations, @@ -32,10 +34,7 @@ import DbInfo from "./DbInfo"; import DocumentFilter from "./DocumentFilter"; import SessionManager from "./SessionManager"; -// Must match RAGCapability.state_namespace. -const AGUI_STATE_KEY = "rag"; - -// AG-UI state is namespaced under AGUI_STATE_KEY +// AG-UI state is namespaced under AGUI_STATE_KEY (see sessionStorage). interface AgentState { [AGUI_STATE_KEY]?: RAGState; } @@ -319,11 +318,10 @@ function ChatContentInner({ useEffect(() => { if (agent.messages.length > 0) return; const session = getSession(sessionId); - // Seed state for the capability; the backend replaces it after each run. - agent.setState({ - ...agent.state, - [AGUI_STATE_KEY]: normalizeRAGState(session?.ragState), - }); + // Seed state for the capabilities; the backend replaces it after each run. + // The whole namespace map goes back, not just the fields this UI reads: + // compaction and the citation policy read what earlier questions recorded. + agent.setState({ ...agent.state, ...agentStateOf(session ?? undefined) }); if (session && session.messages.length > 0) { // biome-ignore lint/suspicious/noExplicitAny: AG-UI Message type is a broad union agent.setMessages(session.messages as any[]); @@ -336,13 +334,10 @@ function ChatContentInner({ // biome-ignore lint/correctness/useExhaustiveDependencies: JSON.stringify tracks content changes useEffect(() => { if (sessionId && agent.messages.length > 0) { - const currentRagState = normalizeRAGState( - (agent.state as AgentState)?.[AGUI_STATE_KEY], - ); updateSessionMessages( sessionId, serializeMessages(agent.messages), - currentRagState, + (agent.state ?? {}) as Record, ); } }, [JSON.stringify(agent.messages), ragState, sessionId]); diff --git a/app/frontend/lib/sessionStorage.ts b/app/frontend/lib/sessionStorage.ts index 69d8ac7c..63642444 100644 --- a/app/frontend/lib/sessionStorage.ts +++ b/app/frontend/lib/sessionStorage.ts @@ -34,11 +34,32 @@ export interface StoredSession { id: string; title: string; messages: StoredMessage[]; - ragState: RAGState; + // The whole AG-UI state. The rag namespace is not the only one a capability + // writes: the citation policy records violations beside it. + agentState: AgentState; + // Sessions stored before agentState existed. + ragState?: RAGState; createdAt: string; updatedAt: string; } +export const AGUI_STATE_KEY = "rag"; + +export type AgentState = Record; + +// Reads the rag namespace out of a stored session, whichever way it was stored. +export function ragStateOf(session?: StoredSession): RAGState { + const namespaced = session?.agentState?.[AGUI_STATE_KEY] as + | Partial + | undefined; + return normalizeRAGState(namespaced ?? session?.ragState); +} + +// The state to seed an agent with when a session is resumed. +export function agentStateOf(session?: StoredSession): AgentState { + return session?.agentState ?? { [AGUI_STATE_KEY]: ragStateOf(session) }; +} + const SESSIONS_KEY = "haiku.rag.sessions"; const ACTIVE_SESSION_KEY = "haiku.rag.activeSession"; @@ -86,7 +107,7 @@ export function createSession(): StoredSession { id: crypto.randomUUID(), title: "New Session", messages: [], - ragState: normalizeRAGState(), + agentState: { [AGUI_STATE_KEY]: normalizeRAGState() }, createdAt: now, updatedAt: now, }; @@ -111,7 +132,7 @@ export function saveSession(session: StoredSession): void { export function updateSessionMessages( id: string, messages: StoredMessage[], - ragState: RAGState, + agentState: AgentState, ): void { const sessions = getAllSessions(); const idx = sessions.findIndex((s) => s.id === id); @@ -119,7 +140,7 @@ export function updateSessionMessages( const session = sessions[idx]; session.messages = messages; - session.ragState = ragState; + session.agentState = agentState; session.updatedAt = new Date().toISOString(); // Derive title from first user message diff --git a/docs/capabilities/compaction.md b/docs/capabilities/compaction.md index 46c6a038..e0baeb25 100644 --- a/docs/capabilities/compaction.md +++ b/docs/capabilities/compaction.md @@ -21,6 +21,12 @@ agent = Agent( It exposes no tools and takes no configuration. Registering it is the only switch: leave it out and the transcript reaches the model untouched. +The host must carry the capability state between runs, alongside the message +history: the capsule is built from what earlier questions recorded there. Given +only a message history, every run starts from an empty record, and compaction +refuses rather than replace evidence it cannot retain. See +[Compose an agent](index.md#compose-an-agent) for the shape. + ## What it does On each request, evidence from earlier questions is replaced by the evidence those diff --git a/docs/capabilities/index.md b/docs/capabilities/index.md index 15636fff..f4b0e105 100644 --- a/docs/capabilities/index.md +++ b/docs/capabilities/index.md @@ -16,11 +16,22 @@ The two evidence capabilities are deferred by default. An agent initially sees o Pick one evidence capability, and add both optional capabilities to it: ```python +from dataclasses import dataclass, field +from typing import Any + from pydantic_ai import Agent +from pydantic_ai.messages import ModelMessage + from haiku.rag.capabilities.compaction import create_capability as compaction from haiku.rag.capabilities.policy import create_capability as citation_policy from haiku.rag.capabilities.rag import create_capability as rag + +@dataclass +class Deps: + state: dict[str, Any] = field(default_factory=dict) + + agent = Agent( "openai:gpt-5", capabilities=[ @@ -28,12 +39,27 @@ agent = Agent( compaction(), citation_policy(), ], + deps_type=Deps, ) -result = await agent.run("What does the knowledge base say about X?") +# One Deps and one history for the conversation: the capabilities read both. +deps = Deps() +history: list[ModelMessage] = [] + +result = await agent.run("What does the knowledge base say about X?", deps=deps, message_history=history) +history = list(result.all_messages()) print(result.output) ``` +!!! warning "Both optional capabilities need the host to carry state" + + They read what earlier questions retrieved and cited from the capability's + state, so the host must expose a `state` dict on its agent dependencies and + hand the same dict back on every run of a conversation, alongside the message + history. With only the message history, every run starts from an empty record: + compaction refuses rather than replace evidence it cannot retain, and the + citation policy cannot enforce a follow-up about evidence cited earlier. + Swap `rag` for `analysis` for an analysis agent. Both optional capabilities work the same way with either one, and neither exposes tools or takes configuration. diff --git a/docs/capabilities/policy.md b/docs/capabilities/policy.md index de0b56df..19e26728 100644 --- a/docs/capabilities/policy.md +++ b/docs/capabilities/policy.md @@ -16,9 +16,14 @@ agent = Agent( ) ``` -It exposes no tools and takes no configuration. Exactly one policy capability makes the -decision, however many evidence capabilities are registered, so two of them cannot each -demand a citation for one answer. +It exposes no tools and takes no configuration. Exactly one policy capability makes +the decision, however many evidence capabilities are registered, so two of them +cannot each demand a citation for one answer. + +The host must carry the capability state between runs, alongside the message +history. Enforcement reads what the conversation has already cited, so without it +a follow-up about evidence cited earlier goes unenforced. See +[Compose an agent](index.md#compose-an-agent) for the shape. ## Declaring nothing is a valid answer diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index 85052457..b9f0ff60 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py @@ -135,6 +135,13 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): request_count: int = field(default=0, repr=False) grace_requests_used: int = field(default=0, repr=False) epoch: int = field(default=0, repr=False) + state_carried: bool = field(default=False, repr=False) + """Whether the host handed back a record a previous question had stamped. + + False on a first question, and equally on every question of a host that does + not carry state between runs. Capabilities that need the record to mean + anything across questions read it to refuse rather than act on nothing. + """ async def for_run(self, ctx: RunContext[Any]) -> "RAGCapabilityBase[StateT]": """Start a run's own copy, and settle which question it is answering. @@ -155,6 +162,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): state = self.state_type.model_validate(raw_state or {}) record = cast(CapabilityEvidenceRecord, cast(Any, state).evidence) continuing = record.in_progress + state_carried = record.question is not None if not continuing and _awaits_the_model(ctx.messages): raise RuntimeError( f"The {self.state_namespace} capability is resuming a question with " @@ -176,6 +184,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): request_count=0, grace_requests_used=0, epoch=0, + state_carried=state_carried, ) run_capability._sync_state() return run_capability diff --git a/haiku_rag_slim/haiku/rag/capabilities/compaction.py b/haiku_rag_slim/haiku/rag/capabilities/compaction.py index 85e5170e..6b1decbb 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/compaction.py +++ b/haiku_rag_slim/haiku/rag/capabilities/compaction.py @@ -321,14 +321,23 @@ class EvidenceCompactionCapability(AbstractCapability[Any]): """ evidence = discover_evidence(ctx) boundary = question_in_progress(evidence) + owned_tools = frozenset().union(*(found.tool_names for found in evidence)) + if boundary > 0 and _newest_owned_return( + request_context.messages, boundary, owned_tools + ): + if not any(found.state_carried for found in evidence): + raise RuntimeError( + "Evidence compaction found evidence from an earlier question but " + "no record of what it cited, so replacing it would retain " + "nothing. The host must carry the capability state between runs, " + "alongside the message history, for this capability to work." + ) if boundary > 0: await self._build_once(ctx, evidence) request_context.messages = compact_history( request_context.messages, boundary=boundary, - owned_tools=frozenset().union( - *(found.tool_names for found in evidence) - ), + owned_tools=owned_tools, capsule_text=self.capsule.text, capsule_images=self.images, ) diff --git a/haiku_rag_slim/haiku/rag/capabilities/evidence.py b/haiku_rag_slim/haiku/rag/capabilities/evidence.py index 75a61cc4..49f8f385 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/evidence.py +++ b/haiku_rag_slim/haiku/rag/capabilities/evidence.py @@ -22,6 +22,7 @@ class DiscoveredEvidence: citations: Mapping[str, Citation] tool_names: frozenset[str] cite_available: bool + state_carried: bool def discover_evidence(ctx: RunContext[Any]) -> list[DiscoveredEvidence]: @@ -40,6 +41,7 @@ def discover_evidence(ctx: RunContext[Any]) -> list[DiscoveredEvidence]: citations=cast(Any, capability.state).citation_index, tool_names=frozenset(capability.evidence_tool_names()), cite_available=capability.cite_available, + state_carried=capability.state_carried, ) for capability in ctx.capabilities.values() if isinstance(capability, RAGCapabilityBase) diff --git a/tests/capabilities/test_evidence_capsule.py b/tests/capabilities/test_evidence_capsule.py index 981db4c3..056accf3 100644 --- a/tests/capabilities/test_evidence_capsule.py +++ b/tests/capabilities/test_evidence_capsule.py @@ -68,6 +68,7 @@ def discovered( return DiscoveredEvidence( capability=capability, record=record, + state_carried=True, citations={ chunk_id: citation( chunk_id, diff --git a/tests/capabilities/test_evidence_wire.py b/tests/capabilities/test_evidence_wire.py index 6ad32aa4..7e588ca1 100644 --- a/tests/capabilities/test_evidence_wire.py +++ b/tests/capabilities/test_evidence_wire.py @@ -287,6 +287,17 @@ def rag_and_compactor(temp_db_path): ) +def settled_deps(question: int = 0) -> Deps: + """State as a host carrying it has it: an earlier question, answered.""" + return Deps( + state={ + "rag": RAGState( + evidence=CapabilityEvidenceRecord(question=question, in_progress=False) + ).model_dump(mode="json") + } + ) + + def in_flight_history() -> list[Any]: """A question already asked and searched, still awaiting its answer.""" return [ @@ -322,7 +333,9 @@ async def test_without_the_compactor_the_history_is_untouched(temp_db_path): agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag]) settled = [*in_flight_history(), ModelResponse(parts=[TextPart("first answer")])] - await agent.run("a different question", deps=Deps(), message_history=settled) + await agent.run( + "a different question", deps=settled_deps(), message_history=settled + ) assert returns_of(wire[-1]) == ["EVIDENCE FOR THE LIVE TURN"] @@ -341,7 +354,9 @@ async def test_with_the_compactor_a_new_question_compacts_the_previous_one( agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag, compactor]) settled = [*in_flight_history(), ModelResponse(parts=[TextPart("first answer")])] - await agent.run("a different question", deps=Deps(), message_history=settled) + await agent.run( + "a different question", deps=settled_deps(), message_history=settled + ) assert returns_of(wire[-1]) == [RECEIPT] @@ -509,7 +524,8 @@ async def test_the_capsule_is_built_once_per_request_and_again_for_the_next( ctx = RunContext( deps=deps, model=TestModel(), usage=RunUsage(), run_id="run-1", run_step=1 ) - run_rag = await rag.for_run(ctx) + # A host carrying state, which is what compaction requires of one. + run_rag = replace(await rag.for_run(ctx), state_carried=True) run_compactor = await compactor.for_run(ctx) cast(Any, run_rag.state).evidence.begin_question(4) ctx = replace(ctx, capabilities={"rag": run_rag, "compaction": run_compactor}) @@ -727,3 +743,49 @@ def test_the_capsule_is_attached_beside_the_newest_return_of_that_request(): assert images_of(compacted) == [fresh] assert returns_of(compacted) == [RECEIPT, "CAPSULE"] + + +@pytest.mark.asyncio +async def test_compaction_refuses_to_strip_evidence_it_cannot_replace(temp_db_path): + """A host that does not carry state has no record to build a capsule from. + + Compacting anyway replaces the earlier evidence with receipts and retains + nothing, so the model loses what it cited and the loss is invisible: the + citations the host already displayed are still there. + """ + rag, compactor = rag_and_compactor(temp_db_path) + + async def model(_messages, _info): # pragma: no cover - never reached + return ModelResponse(parts=[TextPart("answer")]) + + agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag, compactor]) + history = answered_question("an earlier question", evidence="EVIDENCE TO LOSE") + + with pytest.raises(RuntimeError, match="carry the capability state"): + await agent.run("a follow-up", deps=Deps(), message_history=history) + + +@pytest.mark.asyncio +async def test_compaction_proceeds_for_a_host_that_carries_state(temp_db_path): + """The same history, with the record the earlier question left behind.""" + rag, compactor = rag_and_compactor(temp_db_path) + wire: list[list[Any]] = [] + + async def model(messages, _info): + wire.append(list(messages)) + return ModelResponse(parts=[TextPart("answer")]) + + agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag, compactor]) + history = answered_question("an earlier question", evidence="EVIDENCE TO LOSE") + + carried = Deps( + state={ + "rag": RAGState( + evidence=CapabilityEvidenceRecord(question=0, in_progress=False) + ).model_dump(mode="json") + } + ) + + await agent.run("a follow-up", deps=carried, message_history=history) + + assert returns_of(wire[-1]) == [RECEIPT]