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

Move the wire rewrite into the compaction capability

Browse source 
`_compact_old_tool_returns`, `PRIOR_TURN_NOTICE` and `turn_start` leave
`RAGCapabilityBase`, along with its `wrap_model_request` hook. The evidence
capabilities now retrieve and validate, and nothing else. Registering the compaction
capability is what rewrites a request; leaving it out sends the transcript untouched,
which was never a choice a host could make before.

The boundary is the recorded question identity rather than message shape, so a
resumption compacts what lies below the question in progress instead of switching
compaction off for the whole run. The newest earlier evidence return carries the
capsule and every other becomes a receipt, so one capsule exists by construction and
every return stays paired with its call.

Pictures of cited evidence are fetched through the capability that retrieved them and
re-attached beside the capsule with fresh labels. Ownership of a picture on the wire
requires the machine tag we write and an image directly after it, since neither
position nor prose is proof: several tools' results can arrive in one request, and a
user quoting our wording above their own picture had it removed. A picture that cannot
be fetched or decoded is emitted with neither its image nor its label.

The chat TUI and the example backend register the compactor, being multi-turn.
`client.ask`, `client.analyze` and the MCP tools do not: a single-shot question has
nothing earlier to compact.
This commit is contained in:
Yiorgis Gozadinos 2026-08-11 13:47:30 +03:00
parent 0aa6d79f88
commit 2cd568847e
No known key found for this signature in database
9 changed files with 969 additions and 460 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
CHANGELOG.md
View file

@ -3,12 +3,20 @@
### Added
- `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.
- `RAGCapabilityBase.evidence_tool_names()` and `get_picture_bytes()`.
- `haiku.rag.tools.search.decode_picture()`.
### Changed
- `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.
### Removed
- `PRIOR_TURN_NOTICE` and `_compact_old_tool_returns` from `haiku.rag.capabilities._base`, and `RAGCapabilityBase.turn_start`.
### Fixed
- `cross-encoder` reranking no longer ties the scores of strongly-relevant candidates, which left their order to the sort. Scores remain 0-1.
@ -17,10 +25,6 @@
- `docling-local` reuses one docling `DocumentConverter` per set of conversion options instead of building one per document, so local layout, table and OCR models are no longer loaded per document. Conversions through a shared converter are serialized.
- A resumed run keeps the searches, citations and executions of the question in progress instead of clearing them.
- Each page image attached to a search result is preceded by a line giving its position and the chunk id it came from. `build_binary_parts_from_results` is now `build_image_content_from_results` and returns those labels interleaved with the pictures.
- Prior-question tool output is trimmed from the model request only; `all_messages()` retains what the run gathered.
- Evidence retrieved for the current question is no longer trimmed mid-question, whether or not it carries page images.
- A run resumed without a prompt (deferred tool results, interruption, suspension) no longer has the active question's evidence trimmed as if it belonged to an earlier one.
- `rag_cite` / `analysis_cite` returns are no longer replaced by the prior-question notice.
## [0.73.0] - 2026-08-06

7
app/backend/main.py
View file

@ -18,6 +18,9 @@ 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 create_compaction,
)
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
@ -81,7 +84,9 @@ capability = create_capability(db_path=db_path, config=Config, defer_loading=Fal
agent = Agent(
get_model(Config.qa.model, Config),
instructions=AGENT_PREAMBLE,
capabilities=[capability],
# Conversations here are multi-turn, so earlier questions are reduced to the
# evidence they cited rather than carried whole.
capabilities=[capability, create_compaction()],
deps_type=AppDeps,
)

37
docs/capabilities/index.md
View file

@ -1,13 +1,14 @@
# Capabilities
haiku.rag provides two native [Pydantic AI capabilities](https://ai.pydantic.dev/capabilities/):
haiku.rag provides native [Pydantic AI capabilities](https://ai.pydantic.dev/capabilities/):
| Capability | Use it for |
|---|---|
| [`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. |
Both 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.
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
@ -34,6 +35,38 @@ agent = Agent(
)
```
## 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.
## State
Capabilities use a plain `state: dict[str, Any]` attribute on agent dependencies when one is available. RAG state lives under `"rag"`; analysis state lives under `"analysis"`. This keeps state independent of any transport or UI protocol.

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

@ -7,7 +7,7 @@ from typing import Any, cast
from pydantic import BaseModel
from pydantic_ai import ModelRetry, RunContext, ToolFailed
from pydantic_ai.capabilities import AbstractCapability, WrapModelRequestHandler
from pydantic_ai.capabilities import AbstractCapability
from pydantic_ai.messages import (
InstructionPart,
ModelMessage,
@ -15,7 +15,6 @@ from pydantic_ai.messages import (
ModelResponse,
ToolCallPart,
ToolReturn,
ToolReturnPart,
)
from pydantic_ai.models import ModelRequestContext
from pydantic_ai.run import AgentRunResult
@ -83,54 +82,6 @@ def _clear_invocation_state(state: BaseModel) -> None:
value.clear()
PRIOR_TURN_NOTICE = (
"[Evidence retrieved for an earlier question, no longer shown. It does not "
"count as cited for the current question.]"
)
def _compact_old_tool_returns(
messages: list[ModelMessage],
tool_names: frozenset[str],
*,
turn_start: int,
) -> list[ModelMessage]:
"""Remove bulky earlier-question evidence while retaining the current one.
Tool call and return parts remain paired; only the old return payload is
replaced.
Page images attached to a replaced return are deliberately left in place.
Dropping them alongside their text is tempting — they are the bulk, and
they accumulate — but a follow-up about a figure already shown ("what
colour is that box?") carries no terms that could retrieve it again, so
removing the image turns an answerable question into a refusal.
``turn_start`` is how many messages existed when the current question
arrived, so everything below it belongs to an earlier one. The run reports
it rather than this function deriving it from message shape: a
``UserPromptPart`` mid-question is as likely to be page images on a tool
return, or a notice a capability injected, and reading either as the next
question strips evidence the model is still answering from.
"""
if turn_start <= 0:
return messages
compacted = list(messages)
for index, message in enumerate(messages[:turn_start]):
if not isinstance(message, ModelRequest):
continue
parts = [
replace(part, content=PRIOR_TURN_NOTICE)
if isinstance(part, ToolReturnPart) and part.tool_name in tool_names
else part
for part in message.parts
]
if parts != message.parts:
compacted[index] = replace(message, parts=parts)
return compacted
def _is_resumption(prompt: Any, messages: list[ModelMessage]) -> bool:
"""Whether this run continues a question rather than asking a new one.
@ -185,17 +136,10 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
search_count: int = field(default=0, repr=False)
request_count: int = field(default=0, repr=False)
grace_requests_used: int = field(default=0, repr=False)
turn_start: int = field(default=0, repr=False)
epoch: int = field(default=0, repr=False)
async def for_run(self, ctx: RunContext[Any]) -> "RAGCapabilityBase[StateT]":
"""Start a run's own copy, and decide what counts as an earlier question.
On a resumption (see ``_is_resumption``) compaction is switched off for the
whole run: ``len(ctx.messages)`` would count the live question's own
messages and replace its evidence with the earlier-question notice, leaving
the model to answer with the evidence taken away. Failing this way costs a
larger request; failing the other way costs the answer.
"""Start a run's own copy, and settle which question it is answering.
A new question takes the message count as its identity, which every
participant derives identically from the same history. A resumption keeps
@ -235,7 +179,6 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
request_count=0,
grace_requests_used=0,
epoch=0,
turn_start=0 if resuming else len(ctx.messages),
)
run_capability._sync_state()
return run_capability
@ -245,26 +188,6 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
return f"{self.config.prompts.domain_preamble}\n\n{self.instruction_text}"
return self.instruction_text
async def wrap_model_request(
self,
ctx: RunContext[Any],
*,
request_context: ModelRequestContext,
handler: WrapModelRequestHandler,
) -> ModelResponse:
"""Trim earlier-question evidence off the wire only.
Deliberately not ``before_model_request``: that hook's result is
assigned back onto the run's message history, so trimming there would
destroy the host's record of what was retrieved.
"""
request_context.messages = _compact_old_tool_returns(
request_context.messages,
self.tool_names - {self._cite_tool_name},
turn_start=self.turn_start,
)
return await handler(request_context)
async def before_model_request(
self, ctx: RunContext[Any], request_context: ModelRequestContext
) -> ModelRequestContext:
@ -406,6 +329,18 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]):
self.rag = rag
return self.rag
async def get_picture_bytes(self, document_id: str, self_ref: str) -> bytes | None:
"""Fetch a picture of this capability's evidence, for whoever re-attaches it.
Public because compaction rehydrates cited pictures and this capability
already holds the connection they came from; bytes are never kept in state.
"""
async with self.rag_lock:
rag = await self._ensure_rag()
return await rag.document_item_repository.get_picture_bytes(
document_id, self_ref
)
async def _close(self) -> None:
if self.rag is not None:
await self.rag.__aexit__(None, None, None)

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

@ -1,13 +1,23 @@
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from dataclasses import dataclass, field, replace
from typing import Any, cast
from pydantic_ai import RunContext
from pydantic_ai.capabilities import AbstractCapability
from pydantic_ai.capabilities import AbstractCapability, WrapModelRequestHandler
from pydantic_ai.messages import (
BinaryContent,
ModelMessage,
ModelRequest,
ModelResponse,
ToolReturnPart,
UserPromptPart,
)
from pydantic_ai.models import ModelRequestContext
from haiku.rag.capabilities._base import RAGCapabilityBase
from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord
from haiku.rag.store.models.citation import Citation
from haiku.rag.tools.search import RETRIEVED_IMAGE_TAG, decode_picture
CAPABILITY_ID = "haiku-rag-evidence-compaction"
@ -38,7 +48,7 @@ def group_label(position: int) -> str:
def picture_label(chunk_id: str, self_ref: str) -> str:
return (
f"Page image retrieved from the knowledge base for cited evidence "
f"[{chunk_id}] ({self_ref}). Not provided by the user."
f"[{chunk_id}] ({self_ref}). Not provided by the user. {RETRIEVED_IMAGE_TAG}"
)
@ -181,6 +191,102 @@ def build_capsule(evidence: Sequence[DiscoveredEvidence]) -> Capsule:
return Capsule(text=ENTRY_SEPARATOR.join(lines), pictures=tuple(pictures))
def _strip_our_pictures(part: UserPromptPart) -> UserPromptPart | None:
"""Drop the pictures we attached, together with the labels describing them.
Ours is a label carrying the machine tag immediately followed by an image —
both halves required. Position alone is not ownership, since several tools'
results can arrive in one request; prose alone is not either, because a user can
write any phrase, and treating one as proof removed a user's own picture along
with their text. A label is only ever dropped with its picture: left behind it
would tell the model a figure is present when it is gone.
"""
if isinstance(part.content, str):
return part
items = list(part.content)
kept: list[Any] = []
index = 0
while index < len(items):
item = items[index]
following = items[index + 1] if index + 1 < len(items) else None
is_ours = (
isinstance(item, str)
and RETRIEVED_IMAGE_TAG in item
and isinstance(following, BinaryContent)
)
if is_ours:
index += 2
continue
kept.append(item)
index += 1
return replace(part, content=kept) if kept else None
def compact_history(
messages: list[ModelMessage],
*,
boundary: int,
owned_tools: frozenset[str],
capsule_text: str,
capsule_images: Sequence[str | BinaryContent] = (),
) -> list[ModelMessage]:
"""Replace earlier questions' evidence with the capsule, on a copy.
``boundary`` is how many messages existed when the current question arrived, so
everything below it belongs to an earlier one. It comes from the recorded
question identity rather than from message shape: mid-question a user-role part
is as likely to be page images or an injected notice, and reading either as the
next question strips evidence the model is still answering from.
The newest earlier return carries the capsule and every other becomes a receipt,
so exactly one capsule exists by construction. Returns are never removed, only
rewritten, which keeps each one paired with its call. Nothing outside this
capability's evidence tools is touched — not a cite acknowledgement, not another
capability's output, not a picture the user attached.
"""
if boundary <= 0:
return messages
carrier = _newest_owned_return(messages, boundary, owned_tools)
compacted = list(messages)
for index, message in enumerate(messages[:boundary]):
if not isinstance(message, ModelRequest):
continue
parts: list[Any] = []
for part in message.parts:
if isinstance(part, ToolReturnPart) and part.tool_name in owned_tools:
body = capsule_text or RECEIPT if index == carrier else RECEIPT
parts.append(replace(part, content=body))
elif isinstance(part, UserPromptPart):
if (kept := _strip_our_pictures(part)) is not None:
parts.append(kept)
else:
parts.append(part)
if index == carrier and capsule_images:
parts.append(UserPromptPart(content=list(capsule_images)))
if not parts:
# A request with no parts is not a message; whatever emptied it was not
# ours to remove after all.
continue
if parts != message.parts:
compacted[index] = replace(message, parts=parts)
return compacted
def _newest_owned_return(
messages: list[ModelMessage], boundary: int, owned_tools: frozenset[str]
) -> int | None:
"""Index of the last earlier request holding one of our evidence returns."""
for index in range(min(boundary, len(messages)) - 1, -1, -1):
message = messages[index]
if isinstance(message, ModelRequest) and any(
isinstance(part, ToolReturnPart) and part.tool_name in owned_tools
for part in message.parts
):
return index
return None
@dataclass
class EvidenceCompactionCapability(AbstractCapability[Any]):
"""Rewrites the history from what the evidence capabilities recorded.
@ -194,6 +300,94 @@ class EvidenceCompactionCapability(AbstractCapability[Any]):
would share this capability's id.
"""
built_for: tuple[str | None, int] | None = field(default=None, repr=False)
capsule: Capsule = field(default_factory=Capsule, repr=False)
images: tuple[str | BinaryContent, ...] = field(default=(), repr=False)
async def for_run(self, ctx: RunContext[Any]) -> "EvidenceCompactionCapability":
"""Give the run its own build cache, so concurrent runs cannot share one."""
return replace(self, built_for=None, capsule=Capsule(), images=())
async def wrap_model_request(
self,
ctx: RunContext[Any],
*,
request_context: ModelRequestContext,
handler: WrapModelRequestHandler,
) -> ModelResponse:
"""Rewrite the request, never the stored history.
Deliberately not ``before_model_request``: that hook's result is assigned
back onto the run's message history, which would destroy the host's record
of what was retrieved and break the message counts that question identities
and epochs are derived from.
"""
evidence = self.discover(ctx)
boundary = max((found.record.question or 0 for found in evidence), default=0)
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)
),
capsule_text=self.capsule.text,
capsule_images=self.images,
)
return await handler(request_context)
async def _build_once(
self, ctx: RunContext[Any], evidence: Sequence[DiscoveredEvidence]
) -> None:
"""Build the capsule once per model request, however often the hook runs.
Keyed on the run and its step rather than persisted: a stored key would
freeze one question's capsule across the next.
"""
key = (ctx.run_id, ctx.run_step)
if key == self.built_for:
return
self.capsule = build_capsule(evidence)
self.images = await self._rehydrate(ctx)
self.built_for = key
async def _rehydrate(self, ctx: RunContext[Any]) -> tuple[str | BinaryContent, ...]:
"""Fetch the cited pictures through the capability that retrieved them.
Bytes are never stored in state, and the owner already holds an open
connection. A picture that cannot be fetched, for any reason, or that will
not decode, is emitted with neither its image nor its label: a label can
never outlive what it describes, and a figure the model has already been
given in text is not worth failing a question over.
"""
owners = {
capability.state_namespace: capability
for capability in ctx.capabilities.values()
if isinstance(capability, RAGCapabilityBase)
}
content: list[str | BinaryContent] = []
for retained in self.capsule.pictures:
# Indexed, not looked up defensively: the capsule was built from these
# same capabilities in this same call, so a missing owner is a broken
# invariant rather than a picture to skip.
owner = owners[retained.capability]
try:
data = await owner.get_picture_bytes(
retained.document_id, retained.self_ref
)
except Exception:
# A read that fails costs this picture, not the answer.
continue
if data is None:
continue
picture = decode_picture(data, retained.self_ref)
if picture is None:
continue
content.append(retained.label)
content.append(picture)
return tuple(content)
def discover(self, ctx: RunContext[Any]) -> list[DiscoveredEvidence]:
"""Read what each evidence capability recorded, without writing anything.
@ -237,6 +431,7 @@ __all__ = [
"EvidenceCompactionCapability",
"RetainedPicture",
"build_capsule",
"compact_history",
"create_capability",
"group_label",
"picture_label",

5
haiku_rag_slim/haiku/rag/chat/app.py
View file

@ -25,6 +25,7 @@ from textual.worker import Worker
from haiku.rag.capabilities._base import RAGCapabilityBase
from haiku.rag.capabilities.analysis import AnalysisState
from haiku.rag.capabilities.compaction import create_capability as create_compaction
from haiku.rag.capabilities.rag import AGENT_PREAMBLE, RAGState
from haiku.rag.chat.widgets.chat_history import ChatHistory, CitationWidget
from haiku.rag.chat.widgets.image_select import ImageAdded
@ -153,7 +154,9 @@ class ChatApp(App):
self._model,
deps_type=ChatDeps,
instructions=AGENT_PREAMBLE,
capabilities=self._capabilities,
# A chat is multi-turn by definition, so earlier questions are reduced
# to the evidence they cited rather than carried whole.
capabilities=[*self._capabilities, create_compaction()],
)
self._state = {}
for capability in self._capabilities:

47
haiku_rag_slim/haiku/rag/tools/search.py
View file

@ -10,6 +10,30 @@ from haiku.rag.config.models import AppConfig
from haiku.rag.store.models import SearchResult
from haiku.rag.tools.context import RAGDeps
RETRIEVED_IMAGE_TAG = "[haiku.rag/retrieved-image]"
"""Tag every label we attach to a retrieved picture ends with.
Identifies our own pictures on the wire without inferring ownership from position,
which is wrong as soon as two tools' results arrive in one request. Deliberately not
a phrase: a user writing "retrieved from the knowledge base for my report" above
their own picture had it removed, along with their text.
"""
def decode_picture(data: bytes, self_ref: str) -> BinaryContent | None:
"""Wrap picture bytes for the wire, or return nothing if they will not decode.
The model adapter renders one vision placeholder per ``BinaryContent``, so
emitting one for an image the server cannot decode leaves the processor with an
off-by-one count.
"""
try:
with Image.open(BytesIO(data)) as image:
image.verify()
except Exception:
return None
return BinaryContent(data=data, media_type="image/png", identifier=self_ref)
def build_image_content_from_results(
results: list[SearchResult],
@ -33,7 +57,7 @@ def build_image_content_from_results(
``BinaryContent.identifier`` cannot do — it does not survive serialization
to the vision API.
"""
collected: list[tuple[str | None, str, bytes]] = []
collected: list[tuple[str | None, str, BinaryContent]] = []
seen: set[tuple[str | None, str]] = set()
for result in results:
if not result.image_data:
@ -42,26 +66,21 @@ def build_image_content_from_results(
key = (result.document_id, self_ref)
if key in seen:
continue
data = base64.b64decode(b64)
try:
with Image.open(BytesIO(data)) as img:
img.verify()
except Exception:
picture = decode_picture(base64.b64decode(b64), self_ref)
if picture is None:
continue
collected.append((result.chunk_id, self_ref, data))
collected.append((result.chunk_id, self_ref, picture))
seen.add(key)
content: list[str | BinaryContent] = []
total = len(collected)
for position, (chunk_id, self_ref, data) in enumerate(collected, 1):
for position, (chunk_id, self_ref, picture) in enumerate(collected, 1):
content.append(
f"Page image {position} of {total}, retrieved from the knowledge "
f"base for search result [{chunk_id}] ({self_ref}). "
"Not provided by the user."
)
content.append(
BinaryContent(data=data, media_type="image/png", identifier=self_ref)
f"Page image {position} of {total}, retrieved from the knowledge base "
f"for search result [{chunk_id}] ({self_ref}). "
f"Not provided by the user. {RETRIEVED_IMAGE_TAG}"
)
content.append(picture)
return content

372
tests/capabilities/test_capabilities.py
View file

@ -6,7 +6,6 @@ from unittest.mock import AsyncMock, patch
import pytest
from pydantic_ai import Agent, DeferredToolResults, ModelRetry, RunContext, ToolFailed
from pydantic_ai.messages import (
BinaryContent,
ModelRequest,
ModelResponse,
TextPart,
@ -20,9 +19,7 @@ from pydantic_ai.usage import RunUsage
from haiku.rag.capabilities._base import (
CITATION_GRACE_REQUESTS,
PRIOR_TURN_NOTICE,
_called_own_tool,
_compact_old_tool_returns,
)
from haiku.rag.capabilities.analysis import AnalysisCapability, AnalysisState
from haiku.rag.capabilities.analysis import create_capability as create_analysis
@ -848,182 +845,6 @@ async def test_deferred_capability_loads_native_tools(temp_db_path):
assert "rag_search" in loaded_payloads[0]
def test_prior_turn_tool_results_are_compacted_but_current_evidence_is_kept():
messages = [
ModelRequest(parts=[UserPromptPart("old question")]),
ModelResponse(parts=[ToolCallPart("rag_search", {}, "old-call")]),
ModelRequest(
parts=[ToolReturnPart("rag_search", "large old evidence", "old-call")]
),
ModelRequest(parts=[UserPromptPart("current question")]),
ModelResponse(parts=[ToolCallPart("rag_search", {}, "current-call")]),
ModelRequest(
parts=[ToolReturnPart("rag_search", "current evidence", "current-call")]
),
]
compacted = _compact_old_tool_returns(
messages, frozenset({"rag_search"}), turn_start=3
)
old_return = compacted[2].parts[0]
current_return = compacted[5].parts[0]
assert isinstance(old_return, ToolReturnPart)
assert old_return.content == PRIOR_TURN_NOTICE
assert isinstance(current_return, ToolReturnPart)
assert current_return.content == "current evidence"
def test_nothing_is_compacted_on_the_first_question():
messages = [
ModelResponse(parts=[ToolCallPart("rag_search", {}, "current-call")]),
ModelRequest(
parts=[ToolReturnPart("rag_search", "current evidence", "current-call")]
),
]
compacted = _compact_old_tool_returns(
messages, frozenset({"rag_search"}), turn_start=0
)
assert compacted is messages
current_return = compacted[1].parts[0]
assert isinstance(current_return, ToolReturnPart)
assert current_return.content == "current evidence"
PAGE_IMAGE = BinaryContent(data=b"\x89PNG" + b"\x00" * 64, media_type="image/png")
def _search_exchange(call_id: str, evidence: str, *, images: bool):
"""One search round-trip in the shape pydantic-ai produces.
Images on a ``ToolReturn`` arrive as a separate ``UserPromptPart`` appended
to the same ``ModelRequest`` as the ``ToolReturnPart``.
"""
parts: list[Any] = [ToolReturnPart("rag_search", evidence, call_id)]
if images:
parts.append(UserPromptPart(content=[PAGE_IMAGE]))
return [
ModelResponse(parts=[ToolCallPart("rag_search", {}, call_id)]),
ModelRequest(parts=parts),
]
@pytest.mark.parametrize(
("label", "trailing"),
[
("page images on a tool return", [UserPromptPart(content=[PAGE_IMAGE])]),
("a notice injected mid-run", [UserPromptPart("You answered without citing")]),
],
)
def test_current_turn_evidence_survives_later_user_prompt_parts(label, trailing):
"""Nothing that arrives mid-question may be read as the next question.
Page images on a tool return and a notice this capability injects both
appear as a ``UserPromptPart`` after the question, so deriving the turn from
message shape stripped evidence the model was still answering from.
"""
messages = [
ModelRequest(parts=[UserPromptPart("current question")]),
*_search_exchange("first-call", "first evidence", images=False),
ModelRequest(parts=trailing),
]
compacted = _compact_old_tool_returns(
messages, frozenset({"rag_search"}), turn_start=0
)
first_return = compacted[2].parts[0]
assert isinstance(first_return, ToolReturnPart)
assert first_return.content == "first evidence", label
def test_prior_turn_images_outlive_their_tool_return():
"""A follow-up about a figure cannot retrieve it again, so keep the image.
"What colour is that box?" has no terms the search can use, so dropping the
image with its text turns an answerable question into a refusal.
"""
messages = [
ModelRequest(parts=[UserPromptPart("old question")]),
*_search_exchange("old-call", "old evidence", images=True),
ModelResponse(parts=[TextPart("an answer")]),
ModelRequest(parts=[UserPromptPart("current question")]),
]
compacted = _compact_old_tool_returns(
messages, frozenset({"rag_search"}), turn_start=4
)
old_return = compacted[2].parts[0]
assert isinstance(old_return, ToolReturnPart)
assert old_return.content == PRIOR_TURN_NOTICE
assert any(
isinstance(part, UserPromptPart) and not isinstance(part.content, str)
for message in compacted
for part in message.parts
)
def test_user_attached_image_is_never_dropped():
messages = [
ModelRequest(parts=[UserPromptPart("old question")]),
*_search_exchange("old-call", "old evidence", images=False),
ModelResponse(parts=[TextPart("an answer")]),
ModelRequest(parts=[UserPromptPart(content=[PAGE_IMAGE, "what is this?"])]),
]
compacted = _compact_old_tool_returns(
messages, frozenset({"rag_search"}), turn_start=4
)
old_return = compacted[2].parts[0]
assert isinstance(old_return, ToolReturnPart)
assert old_return.content == PRIOR_TURN_NOTICE
attached = compacted[-1].parts[0]
assert isinstance(attached, UserPromptPart)
assert attached.content == [PAGE_IMAGE, "what is this?"]
async def test_compaction_never_reaches_the_stored_message_history(temp_db_path):
"""Trimming is for the wire; hosts keep the evidence they gathered."""
capability = create_rag(
db_path=temp_db_path, config=AppConfig(), defer_loading=False
)
turns = iter(
[
ModelResponse(parts=[ToolCallPart("rag_search", {"query": "q"}, "call-1")]),
ModelResponse(parts=[TextPart("first answer")]),
ModelResponse(parts=[TextPart("second answer")]),
]
)
async def model(_messages, _info):
return next(turns)
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[capability])
deps = Deps(state={"rag": RAGState().model_dump(mode="json")})
with patch.object(
RAGCapability, "_search", AsyncMock(return_value="REAL EVIDENCE")
):
first = await agent.run("old question", deps=deps)
second = await agent.run(
"current question", deps=deps, message_history=first.all_messages()
)
returns = [
str(part.content)
for message in second.all_messages()
if isinstance(message, ModelRequest)
for part in message.parts
if isinstance(part, ToolReturnPart)
]
assert "REAL EVIDENCE" in returns
assert PRIOR_TURN_NOTICE not in returns
def _resuming_deps() -> Deps:
"""State as a resumption always finds it: the question already identified.
@ -1051,105 +872,9 @@ def _in_flight_history() -> list[Any]:
]
def _wire_returns(sent: list[Any]) -> list[str]:
return [
str(part.content)
for message in sent
if isinstance(message, ModelRequest)
for part in message.parts
if isinstance(part, ToolReturnPart)
]
@pytest.mark.asyncio
async def test_a_new_question_compacts_the_previous_one(temp_db_path):
"""The baseline the resume cases are contrasted against.
A settled history — the previous question answered — is what a genuinely new
question follows. An unfinished tail is ambiguous instead, and treated as a
continuation.
"""
capability = create_rag(
db_path=temp_db_path, config=AppConfig(), defer_loading=False
)
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=[capability])
settled = [*_in_flight_history(), ModelResponse(parts=[TextPart("first answer")])]
await agent.run("a different question", deps=Deps(), message_history=settled)
assert _wire_returns(wire[-1]) == [PRIOR_TURN_NOTICE]
@pytest.mark.asyncio
async def test_a_prompt_on_an_unanswered_tail_is_treated_as_a_continuation(
temp_db_path,
):
"""Ambiguous shape, resolved the safe way.
A history ending in a request the model never answered, plus a new prompt,
could be a fresh question or a continuation. Compacting would cost the answer
if it is a continuation; not compacting only costs a larger request.
"""
capability = create_rag(
db_path=temp_db_path, config=AppConfig(), defer_loading=False
)
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=[capability])
await agent.run(
"a different question",
deps=_resuming_deps(),
message_history=_in_flight_history(),
)
assert _wire_returns(wire[-1]) == ["EVIDENCE FOR THE LIVE TURN"]
@pytest.mark.asyncio
async def test_a_resume_carrying_a_prompt_keeps_the_active_evidence(temp_db_path):
"""Deferred results may arrive with a prompt, and that is still a continuation.
``ctx.prompt`` is non-null here, so the absence of a prompt cannot be the only
signal: the history tail is unfinished — a response whose tool call has no
return yet — and the question it belongs to is still being answered.
"""
capability = create_rag(
db_path=temp_db_path, config=AppConfig(), defer_loading=False
)
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=[capability])
history = [
*_in_flight_history(),
ModelResponse(parts=[ToolCallPart("external_tool", {}, "call-2")]),
]
await agent.run(
"carry on",
deferred_tool_results=DeferredToolResults(calls={"call-2": "external result"}),
message_history=history,
deps=_resuming_deps(),
)
assert "EVIDENCE FOR THE LIVE TURN" in _wire_returns(wire[-1])
assert PRIOR_TURN_NOTICE not in _wire_returns(wire[-1])
@pytest.mark.asyncio
@pytest.mark.parametrize(
"resume_kwargs",
@ -1160,88 +885,7 @@ async def test_a_resume_carrying_a_prompt_keeps_the_active_evidence(temp_db_path
),
],
)
async def test_a_resumed_run_keeps_the_active_questions_evidence(
temp_db_path, resume_kwargs
):
"""A run without a prompt continues a question; nothing in it is prior.
``len(ctx.messages)`` cannot tell the two apart — on a resumption it counts
the live question's own messages and marks its evidence as earlier-question
evidence, leaving the model to answer with a notice where its search result
used to be. Deferred, interrupted and suspended resumes all differ in shape,
so the absence of a prompt is the signal rather than the message layout.
"""
capability = create_rag(
db_path=temp_db_path, config=AppConfig(), defer_loading=False
)
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=[capability])
await agent.run(
deps=_resuming_deps(), message_history=_in_flight_history(), **resume_kwargs
)
assert _wire_returns(wire[-1]) == ["EVIDENCE FOR THE LIVE TURN"]
@pytest.mark.asyncio
async def test_prior_turn_search_is_compacted_but_its_cite_receipt_is_not(temp_db_path):
"""Only evidence is compacted, and the boundary comes from the run.
A cite acknowledgement is a receipt, not evidence: replacing it lengthened
the request and erased the record that citations had been registered.
"""
capability = create_rag(
db_path=temp_db_path, config=AppConfig(), defer_loading=False
)
turns = iter(
[
ModelResponse(parts=[ToolCallPart("rag_search", {"query": "q"}, "call-1")]),
ModelResponse(
parts=[ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-2")]
),
ModelResponse(parts=[TextPart("first answer")]),
ModelResponse(parts=[TextPart("second answer")]),
]
)
wire: list[list[Any]] = []
async def model(messages, _info):
wire.append(list(messages))
return next(turns)
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[capability])
deps = Deps(state={"rag": RAGState().model_dump(mode="json")})
async def search(self, query: str, _limit: int | None) -> str:
"""Record a result the way the real search does, so citing resolves."""
cast(Any, self.state).searches[query] = [
SearchResult(content="evidence", score=1.0, chunk_id="chunk-1")
]
return "REAL EVIDENCE"
with patch.object(RAGCapability, "_search", search):
first = await agent.run("old question", deps=deps)
await agent.run(
"current question", deps=deps, message_history=first.all_messages()
)
prior = {
part.tool_name: str(part.content)
for message in wire[-1]
if isinstance(message, ModelRequest)
for part in message.parts
if isinstance(part, ToolReturnPart)
}
assert prior["rag_search"] == PRIOR_TURN_NOTICE
assert prior["rag_cite"] != PRIOR_TURN_NOTICE
def _record(deps: Deps, namespace: str) -> CapabilityEvidenceRecord:
return CapabilityEvidenceRecord.model_validate(deps.state[namespace]["evidence"])
@ -1596,3 +1240,19 @@ async def test_a_resumption_keeps_the_evidence_the_question_already_gathered(
assert record.question == identity
assert record.occurrences["chunk-1"].retrieved_in_questions == [identity]
assert citation_status([record], question=identity) == "grounded"
@pytest.mark.asyncio
async def test_a_capability_fetches_its_own_evidences_pictures(temp_db_path):
"""Compaction rehydrates through the owner, which already holds the connection."""
capability = create_rag(db_path=temp_db_path, config=AppConfig())
client = AsyncMock()
client.document_item_repository.get_picture_bytes.return_value = b"picture-bytes"
capability.rag = client
data = await capability.get_picture_bytes("doc-1", "#/pictures/0")
assert data == b"picture-bytes"
client.document_item_repository.get_picture_bytes.assert_awaited_once_with(
"doc-1", "#/pictures/0"
)

655
tests/capabilities/test_evidence_wire.py Normal file
View file

@ -0,0 +1,655 @@
import base64
from dataclasses import dataclass, field, replace
from typing import Any, cast
from unittest.mock import AsyncMock, patch
import pytest
from pydantic_ai import Agent, RunContext
from pydantic_ai.messages import (
BinaryContent,
ModelRequest,
ModelResponse,
TextPart,
ToolCallPart,
ToolReturnPart,
UserPromptPart,
)
from pydantic_ai.models import ModelRequestContext, ModelRequestParameters
from pydantic_ai.models.function import FunctionModel
from pydantic_ai.models.test import TestModel
from pydantic_ai.usage import RunUsage
from haiku.rag.capabilities.compaction import (
RECEIPT,
Capsule,
compact_history,
picture_label,
)
from haiku.rag.capabilities.compaction import create_capability as create_compaction
from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord
from haiku.rag.capabilities.rag import RAGCapability, RAGState
from haiku.rag.capabilities.rag import create_capability as create_rag
from haiku.rag.config.models import AppConfig
from haiku.rag.store.models.chunk import SearchResult
OWNED = frozenset({"rag_search"})
PNG = BinaryContent(data=b"fake-image-bytes", media_type="image/png")
def retrieved_image(chunk_id: str = "chunk-1", self_ref: str = "#/pictures/0"):
"""A page image on the wire, labelled the way a search result attaches it."""
return UserPromptPart(content=[picture_label(chunk_id, self_ref), PNG])
def answered_question(question: str, *, evidence: str, images: bool = False):
"""One settled question: prompt, search, result, answer."""
returned: list[Any] = [ToolReturnPart("rag_search", evidence, "call-1")]
if images:
returned.append(retrieved_image())
return [
ModelRequest(parts=[UserPromptPart(question)]),
ModelResponse(parts=[ToolCallPart("rag_search", {"query": "q"}, "call-1")]),
ModelRequest(parts=returned),
ModelResponse(parts=[TextPart("an answer")]),
]
def returns_of(messages) -> list[str]:
return [
str(part.content)
for message in messages
if isinstance(message, ModelRequest)
for part in message.parts
if isinstance(part, ToolReturnPart)
]
def images_of(messages) -> list[BinaryContent]:
return [
item
for message in messages
if isinstance(message, ModelRequest)
for part in message.parts
if isinstance(part, UserPromptPart) and not isinstance(part.content, str)
for item in part.content
if isinstance(item, BinaryContent)
]
def texts_of(messages) -> list[str]:
return [
item
for message in messages
if isinstance(message, ModelRequest)
for part in message.parts
if isinstance(part, UserPromptPart) and not isinstance(part.content, str)
for item in part.content
if isinstance(item, str)
]
def test_nothing_before_the_first_question_is_compacted():
messages = answered_question("first", evidence="EVIDENCE")
compacted = compact_history(
messages, boundary=0, owned_tools=OWNED, capsule_text="CAPSULE"
)
assert compacted == messages
def test_the_newest_earlier_return_carries_the_capsule():
messages = [
*answered_question("first", evidence="OLD EVIDENCE"),
ModelRequest(parts=[UserPromptPart("second")]),
]
compacted = compact_history(
messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE"
)
assert returns_of(compacted) == ["CAPSULE"]
assert "OLD EVIDENCE" not in returns_of(compacted)
def test_older_returns_become_receipts_and_only_the_newest_carries_the_capsule():
messages = [
*answered_question("first", evidence="OLDEST"),
*answered_question("second", evidence="NEWER"),
ModelRequest(parts=[UserPromptPart("third")]),
]
compacted = compact_history(
messages, boundary=8, owned_tools=OWNED, capsule_text="CAPSULE"
)
assert returns_of(compacted) == [RECEIPT, "CAPSULE"]
def test_the_current_question_keeps_its_own_evidence():
messages = [
*answered_question("first", evidence="OLD EVIDENCE"),
ModelRequest(parts=[UserPromptPart("second")]),
ModelResponse(parts=[ToolCallPart("rag_search", {"query": "q"}, "call-2")]),
ModelRequest(parts=[ToolReturnPart("rag_search", "LIVE EVIDENCE", "call-2")]),
]
compacted = compact_history(
messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE"
)
assert returns_of(compacted) == ["CAPSULE", "LIVE EVIDENCE"]
def test_another_capabilitys_return_is_left_alone():
messages = [
ModelRequest(parts=[UserPromptPart("first")]),
ModelResponse(parts=[ToolCallPart("other_tool", {}, "call-1")]),
ModelRequest(parts=[ToolReturnPart("other_tool", "NOT OURS", "call-1")]),
ModelResponse(parts=[TextPart("an answer")]),
ModelRequest(parts=[UserPromptPart("second")]),
]
compacted = compact_history(
messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE"
)
assert returns_of(compacted) == ["NOT OURS"]
def test_a_cite_acknowledgement_survives():
"""A receipt of the model's own action, not evidence."""
messages = [
ModelRequest(parts=[UserPromptPart("first")]),
ModelResponse(
parts=[ToolCallPart("rag_cite", {"chunk_ids": ["c1"]}, "call-1")]
),
ModelRequest(
parts=[ToolReturnPart("rag_cite", "Registered 1 citation.", "c1")]
),
ModelResponse(parts=[TextPart("an answer")]),
ModelRequest(parts=[UserPromptPart("second")]),
]
compacted = compact_history(
messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE"
)
assert returns_of(compacted) == ["Registered 1 citation."]
def test_an_uncited_earlier_image_is_dropped_with_its_label():
messages = [
*answered_question("first", evidence="OLD", images=True),
ModelRequest(parts=[UserPromptPart("second")]),
]
compacted = compact_history(
messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE"
)
assert images_of(compacted) == []
assert texts_of(compacted) == []
def test_cited_pictures_are_attached_beside_the_capsule():
messages = [
*answered_question("first", evidence="OLD", images=True),
ModelRequest(parts=[UserPromptPart("second")]),
]
fresh = BinaryContent(data=b"cited-bytes", media_type="image/png")
compacted = compact_history(
messages,
boundary=4,
owned_tools=OWNED,
capsule_text="CAPSULE",
capsule_images=[picture_label("cited-chunk", "#/pictures/3"), fresh],
)
assert images_of(compacted) == [fresh]
assert texts_of(compacted) == [picture_label("cited-chunk", "#/pictures/3")]
carrier = [
index
for index, message in enumerate(compacted)
if isinstance(message, ModelRequest)
and any(
isinstance(part, ToolReturnPart) and part.content == "CAPSULE"
for part in message.parts
)
]
attached = [
index
for index, message in enumerate(compacted)
if isinstance(message, ModelRequest)
and any(
isinstance(part, UserPromptPart) and not isinstance(part.content, str)
for part in message.parts
)
]
assert carrier == attached
def test_a_user_attached_image_is_never_dropped():
"""The user's own picture is not ours to remove, even in an earlier question."""
mine = UserPromptPart(content=["look at this", PNG])
messages = [
ModelRequest(parts=[mine]),
ModelResponse(parts=[ToolCallPart("rag_search", {"query": "q"}, "call-1")]),
ModelRequest(parts=[ToolReturnPart("rag_search", "OLD", "call-1")]),
ModelResponse(parts=[TextPart("an answer")]),
ModelRequest(parts=[UserPromptPart("second")]),
]
compacted = compact_history(
messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE"
)
assert images_of(compacted) == [PNG]
assert texts_of(compacted) == ["look at this"]
def test_the_stored_messages_are_never_mutated():
messages = [
*answered_question("first", evidence="OLD EVIDENCE", images=True),
ModelRequest(parts=[UserPromptPart("second")]),
]
before = [list(message.parts) for message in messages]
compact_history(messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE")
assert [list(message.parts) for message in messages] == before
assert "OLD EVIDENCE" in returns_of(messages)
def test_nothing_cited_leaves_only_receipts():
messages = [
*answered_question("first", evidence="OLD"),
ModelRequest(parts=[UserPromptPart("second")]),
]
compacted = compact_history(
messages, boundary=4, owned_tools=OWNED, capsule_text=""
)
assert returns_of(compacted) == [RECEIPT]
@dataclass
class Deps:
state: dict[str, Any] = field(default_factory=dict)
def rag_and_compactor(temp_db_path):
return (
create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False),
create_compaction(),
)
def in_flight_history() -> list[Any]:
"""A question already asked and searched, still awaiting its answer."""
return [
ModelRequest(parts=[UserPromptPart("what does the supervisor do?")]),
ModelResponse(parts=[ToolCallPart("rag_search", {"query": "s"}, "call-1")]),
ModelRequest(
parts=[ToolReturnPart("rag_search", "EVIDENCE FOR THE LIVE TURN", "call-1")]
),
]
def resuming_deps(question: int = 0) -> Deps:
"""State as a resumption always finds it: the question already identified."""
return Deps(
state={
"rag": RAGState(
evidence=CapabilityEvidenceRecord(question=question)
).model_dump(mode="json")
}
)
@pytest.mark.asyncio
async def test_without_the_compactor_the_history_is_untouched(temp_db_path):
"""Omission is the switch: there is no flag to test, only absence."""
rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False)
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])
settled = [*in_flight_history(), ModelResponse(parts=[TextPart("first answer")])]
await agent.run("a different question", deps=Deps(), message_history=settled)
assert returns_of(wire[-1]) == ["EVIDENCE FOR THE LIVE TURN"]
@pytest.mark.asyncio
async def test_with_the_compactor_a_new_question_compacts_the_previous_one(
temp_db_path,
):
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])
settled = [*in_flight_history(), ModelResponse(parts=[TextPart("first answer")])]
await agent.run("a different question", deps=Deps(), message_history=settled)
assert returns_of(wire[-1]) == [RECEIPT]
@pytest.mark.asyncio
async def test_a_resumed_question_keeps_the_evidence_it_is_answering_from(
temp_db_path,
):
"""The boundary is the stored identity of the question in progress.
An earlier question below it is compacted; the evidence the model is still
answering from sits above it and survives. Deriving the boundary from message
shape instead would put the live evidence below it and answer with a receipt
where the search result should be.
"""
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 FOR THE OLD TURN"),
*in_flight_history(),
]
await agent.run(deps=resuming_deps(question=4), message_history=history)
assert returns_of(wire[-1]) == [RECEIPT, "EVIDENCE FOR THE LIVE TURN"]
@pytest.mark.asyncio
async def test_compaction_never_reaches_the_stored_message_history(temp_db_path):
"""Rewriting is for the wire; hosts keep the evidence they gathered."""
rag, compactor = rag_and_compactor(temp_db_path)
turns = iter(
[
ModelResponse(parts=[ToolCallPart("rag_search", {"query": "q"}, "call-1")]),
ModelResponse(parts=[TextPart("first answer")]),
ModelResponse(parts=[TextPart("second answer")]),
]
)
async def model(_messages, _info):
return next(turns)
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag, compactor])
deps = Deps()
with patch.object(
RAGCapability, "_search", AsyncMock(return_value="REAL EVIDENCE")
):
first = await agent.run("old question", deps=deps)
second = await agent.run(
"current question", deps=deps, message_history=first.all_messages()
)
assert "REAL EVIDENCE" in returns_of(second.all_messages())
assert RECEIPT not in returns_of(second.all_messages())
REAL_PNG = base64.b64decode(
"iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8DwHwAFBQIAX8jx0gAAAABJRU5ErkJggg=="
)
async def _search_with_a_picture(self, query: str, _limit: int | None) -> str:
"""Record a result carrying a page image, the way a real search does."""
cast(Any, self.state).searches[query] = [
SearchResult(
content="evidence",
score=1.0,
chunk_id="chunk-1",
document_id="doc-1",
doc_item_refs=["#/pictures/0"],
)
]
self._note_evidence()
return "EVIDENCE"
async def _cite_a_picture_chunk(temp_db_path, fetched: bytes | None):
"""Two questions: cite a picture chunk, then ask something else."""
rag, compactor = rag_and_compactor(temp_db_path)
calls = iter(
[
[ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")],
[ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-2")],
[TextPart("first answer")],
[TextPart("second answer")],
]
)
wire: list[list[Any]] = []
async def model(messages, _info):
wire.append(list(messages))
return ModelResponse(parts=next(calls))
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag, compactor])
deps = Deps()
with (
patch.object(RAGCapability, "_search", _search_with_a_picture),
patch.object(
RAGCapability, "get_picture_bytes", AsyncMock(return_value=fetched)
),
):
first = await agent.run("what does the supervisor do?", deps=deps)
await agent.run(
"and what else?", deps=deps, message_history=first.all_messages()
)
return wire
@pytest.mark.asyncio
async def test_a_cited_picture_is_fetched_and_attached_with_its_label(temp_db_path):
wire = await _cite_a_picture_chunk(temp_db_path, REAL_PNG)
assert [picture.data for picture in images_of(wire[-1])] == [REAL_PNG]
assert texts_of(wire[-1]) == [picture_label("chunk-1", "#/pictures/0")]
@pytest.mark.asyncio
async def test_a_picture_that_cannot_be_fetched_emits_neither_image_nor_label(
temp_db_path,
):
"""A label without its picture tells the model a figure is there when it is not."""
wire = await _cite_a_picture_chunk(temp_db_path, None)
assert images_of(wire[-1]) == []
assert texts_of(wire[-1]) == []
@pytest.mark.asyncio
async def test_a_picture_that_will_not_decode_emits_neither_image_nor_label(
temp_db_path,
):
"""One vision placeholder is rendered per attachment, so a corrupt one miscounts."""
wire = await _cite_a_picture_chunk(temp_db_path, b"not-an-image")
assert images_of(wire[-1]) == []
assert texts_of(wire[-1]) == []
@pytest.mark.asyncio
async def test_the_capsule_is_built_once_per_request_and_again_for_the_next(
temp_db_path,
):
"""Two hook passes for one request must not rebuild; the next request must."""
rag, compactor = rag_and_compactor(temp_db_path)
deps = Deps()
ctx = RunContext(
deps=deps, model=TestModel(), usage=RunUsage(), run_id="run-1", run_step=1
)
run_rag = await rag.for_run(ctx)
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})
builds = 0
def counting_build(evidence):
nonlocal builds
builds += 1
return Capsule(text="CAPSULE")
async def handler(_request_context):
return ModelResponse(parts=[TextPart("answer")])
request = ModelRequestContext(
messages=[*answered_question("first", evidence="OLD")],
model=TestModel(),
model_request_parameters=ModelRequestParameters(),
model_settings=None,
)
with patch("haiku.rag.capabilities.compaction.build_capsule", counting_build):
await run_compactor.wrap_model_request(
ctx, request_context=request, handler=handler
)
await run_compactor.wrap_model_request(
ctx, request_context=request, handler=handler
)
assert builds == 1
await run_compactor.wrap_model_request(
replace(ctx, run_step=2), request_context=request, handler=handler
)
assert builds == 2
def test_a_user_quoting_our_wording_keeps_their_picture_and_their_text():
"""Prose is not proof of ownership: a user can write any phrase.
Recognising our own pictures by a natural-language substring removed a user's
image, its text, and with it the whole message part.
"""
quoted = UserPromptPart(
content=[
"Here is a page image retrieved from the knowledge base for my report",
PNG,
]
)
messages = [
ModelRequest(parts=[quoted]),
ModelResponse(parts=[ToolCallPart("rag_search", {"query": "q"}, "call-1")]),
ModelRequest(parts=[ToolReturnPart("rag_search", "OLD", "call-1")]),
ModelResponse(parts=[TextPart("an answer")]),
ModelRequest(parts=[UserPromptPart("second")]),
]
compacted = compact_history(
messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE"
)
assert images_of(compacted) == [PNG]
assert texts_of(compacted) == [
"Here is a page image retrieved from the knowledge base for my report"
]
assert all(message.parts for message in compacted)
def test_a_label_of_ours_with_no_picture_after_it_is_kept():
"""Only a genuine pair is ours to remove; a lone label is someone else's text."""
lonely = UserPromptPart(
content=[picture_label("chunk-1", "#/pictures/0"), "and more"]
)
messages = [
ModelRequest(parts=[lonely]),
ModelResponse(parts=[ToolCallPart("rag_search", {"query": "q"}, "call-1")]),
ModelRequest(parts=[ToolReturnPart("rag_search", "OLD", "call-1")]),
ModelResponse(parts=[TextPart("an answer")]),
ModelRequest(parts=[UserPromptPart("second")]),
]
compacted = compact_history(
messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE"
)
assert texts_of(compacted) == [
picture_label("chunk-1", "#/pictures/0"),
"and more",
]
@pytest.mark.asyncio
async def test_a_picture_whose_fetch_raises_costs_the_picture_not_the_answer(
temp_db_path,
):
rag, compactor = rag_and_compactor(temp_db_path)
calls = iter(
[
[ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")],
[ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-2")],
[TextPart("first answer")],
[TextPart("second answer")],
]
)
wire: list[list[Any]] = []
async def model(messages, _info):
wire.append(list(messages))
return ModelResponse(parts=next(calls))
agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag, compactor])
deps = Deps()
with (
patch.object(RAGCapability, "_search", _search_with_a_picture),
patch.object(
RAGCapability,
"get_picture_bytes",
AsyncMock(side_effect=OSError("the read failed")),
),
):
first = await agent.run("what does the supervisor do?", deps=deps)
second = await agent.run(
"and what else?", deps=deps, message_history=first.all_messages()
)
assert second.output == "second answer"
assert images_of(wire[-1]) == []
assert texts_of(wire[-1]) == []
def test_a_request_is_never_left_with_no_parts():
"""Emptying a message would leave something that is not a message.
Our own pictures always travel with the tool return in their request, so this
shape does not come from us — but a rewritten history can hold it, and a
partless request is invalid whatever produced it.
"""
ours_alone = ModelRequest(
parts=[UserPromptPart(content=[picture_label("chunk-1", "#/pictures/0"), PNG])]
)
messages = [
ours_alone,
ModelResponse(parts=[ToolCallPart("rag_search", {"query": "q"}, "call-1")]),
ModelRequest(parts=[ToolReturnPart("rag_search", "OLD", "call-1")]),
ModelResponse(parts=[TextPart("an answer")]),
ModelRequest(parts=[UserPromptPart("second")]),
]
compacted = compact_history(
messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE"
)
assert all(message.parts for message in compacted)
assert compacted[0] is ours_alone