diff --git a/CHANGELOG.md b/CHANGELOG.md index b9f86243..f3175b97 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -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 diff --git a/app/backend/main.py b/app/backend/main.py index ad8a451d..252e938c 100644 --- a/app/backend/main.py +++ b/app/backend/main.py @@ -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, ) diff --git a/docs/capabilities/index.md b/docs/capabilities/index.md index 17c4ed6b..bff95f8e 100644 --- a/docs/capabilities/index.md +++ b/docs/capabilities/index.md @@ -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. diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index 93f5ac0c..65c43ff3 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py @@ -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) diff --git a/haiku_rag_slim/haiku/rag/capabilities/compaction.py b/haiku_rag_slim/haiku/rag/capabilities/compaction.py index 7179ba27..e9b9e805 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/compaction.py +++ b/haiku_rag_slim/haiku/rag/capabilities/compaction.py @@ -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", diff --git a/haiku_rag_slim/haiku/rag/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py index 0350e5d4..6aa42799 100644 --- a/haiku_rag_slim/haiku/rag/chat/app.py +++ b/haiku_rag_slim/haiku/rag/chat/app.py @@ -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: diff --git a/haiku_rag_slim/haiku/rag/tools/search.py b/haiku_rag_slim/haiku/rag/tools/search.py index 787210b6..4f192a74 100644 --- a/haiku_rag_slim/haiku/rag/tools/search.py +++ b/haiku_rag_slim/haiku/rag/tools/search.py @@ -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 diff --git a/tests/capabilities/test_capabilities.py b/tests/capabilities/test_capabilities.py index 2e140f6f..f4f7943f 100644 --- a/tests/capabilities/test_capabilities.py +++ b/tests/capabilities/test_capabilities.py @@ -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" + ) diff --git a/tests/capabilities/test_evidence_wire.py b/tests/capabilities/test_evidence_wire.py new file mode 100644 index 00000000..bcce85f5 --- /dev/null +++ b/tests/capabilities/test_evidence_wire.py @@ -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