Refuse to compact evidence the host kept no record of
Both optional capabilities read what earlier questions retrieved and cited from the capability's state, so a host that carries only the message history hands every run an empty record. Compaction then replaced the earlier evidence with receipts and retained nothing, and the loss was invisible: the citations the host already displayed were still there. It now refuses when it finds evidence from an earlier question and no record of what that question cited. `state_carried` reaches the optional capabilities through discovery, so the refusal distinguishes a host that never carries state from a question that simply cited nothing. The documentation taught the pattern that breaks: the compose example is now stateful and the requirement is stated where each capability is introduced. The app's browser storage was doing exactly this, keeping only the fields the UI reads. It now persists the whole namespace map, so the citation policy's violations survive a reload as well as the evidence record.
This commit is contained in:
parent
53bbf52697
commit
e1dd8517f9
10 changed files with 163 additions and 27 deletions
|
|
@ -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<string, unknown>,
|
||||
);
|
||||
}
|
||||
}, [JSON.stringify(agent.messages), ragState, sessionId]);
|
||||
|
|
|
|||
|
|
@ -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<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";
|
||||
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -68,6 +68,7 @@ def discovered(
|
|||
return DiscoveredEvidence(
|
||||
capability=capability,
|
||||
record=record,
|
||||
state_carried=True,
|
||||
citations={
|
||||
chunk_id: citation(
|
||||
chunk_id,
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
Loading…
Reference in a new issue