diff --git a/CHANGELOG.md b/CHANGELOG.md index 0c5c1819..f3175b97 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,12 +1,30 @@ # Changelog ## [Unreleased] +### 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. - `haiku-rag` and `haiku-ingester` CLI startup no longer imports `lancedb`, `pyarrow` and `pydantic_ai`. - `haiku.rag.store` no longer re-exports `Store`; import it from `haiku.rag.store.engine`. - `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. ## [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/benchmarks.md b/docs/benchmarks.md index d4d50c92..6bb1975a 100644 --- a/docs/benchmarks.md +++ b/docs/benchmarks.md @@ -133,21 +133,21 @@ Two approaches are benchmarked separately: | Embedding Model | Reranker | Cases | MAP | |------------------------------------------|------------------------------------------------------|------:|-------:| | `Qwen/Qwen3-VL-Embedding-8B` | none | 3045 | 0.9774 | -| `nvidia/llama-nemotron-embed-vl-1b-v2` | none | 3045 | 0.9709 | +| `nvidia/llama-nemotron-embed-vl-1b-v2` | none | 3045 | 0.9798 | | `nvidia/llama-nemotron-embed-vl-1b-v2` | `nvidia/llama-nemotron-rerank-vl-1b-v2` (multimodal) | 3045 | 0.9913 | -*The reranked row uses `reranking.multimodal: true`: picture chunks reach the vision reranker as images alongside their description text. Measured on haiku.rag main post-v0.67.3 (multimodal reranking ships in the next release).* +*The nemotron row without a reranker is measured on this release. The reranked row uses `reranking.multimodal: true`: picture chunks reach the vision reranker as images alongside their description text, measured on haiku.rag main post-v0.67.3.* ##### QA accuracy + citation retrieval | Embedding Model | Target | Capability model | Cases | QA accuracy | Mean `cited_map` | |------------------------------------------|-----------------|-----------------------------------|------:|-------------|------------------| | `Qwen/Qwen3-VL-Embedding-8B` | `rag-capability` | `vllm:Gemma-4-26B-A4B-NVFP4` | 1409 | 0.89 | — | -| `nvidia/llama-nemotron-embed-vl-1b-v2` | `rag-capability` | `vllm:Gemma-4-26B-A4B-NVFP4` | 3045 | 0.92 | 0.93 | -| `nvidia/llama-nemotron-embed-vl-1b-v2` | `analysis-capability`| `vllm:Gemma-4-26B-A4B-NVFP4` | 3045 | 0.94 | 0.78 | +| `nvidia/llama-nemotron-embed-vl-1b-v2` | `rag-capability` | `vllm:Gemma-4-26B-A4B-NVFP4` | 3039 | 0.9263 | 0.9761 | +| `nvidia/llama-nemotron-embed-vl-1b-v2` | `analysis-capability`| `vllm:Gemma-4-26B-A4B-NVFP4` | 3040 | 0.9362 | 0.9343 | | `nvidia/llama-nemotron-embed-vl-1b-v2` | `analysis-capability`| `vllm:Qwen3.6-35B-A3B-NVFP4` | 3045 | 0.95 | 0.93 | -*Measured on haiku.rag v0.52.0, no reranker, judged by `vllm:Qwen3.6-35B-A3B-NVFP4`. Qwen3-VL covered 1409 / 3045 cases.* +*Both nemotron `Gemma-4` rows are measured on this release, no reranker, judged by `vllm:Qwen3.6-35B-A3B-NVFP4` with thinking on, and exclude the cases that errored (6 of 3045 for `rag-capability`, 5 for `analysis-capability`). The `rag-capability` row cites at 99.64% with a mean of 1.08 citations per case, at a median 4.7s per case against 5.0s for `analysis-capability`. Citation coverage is what moved on this release: 4.9% of analysis cases register no citation, against 26.3% before, at unchanged searches and code executions per case. The remaining rows are from haiku.rag v0.52.0, where Qwen3-VL covered 1409 / 3045 cases.* #### Text embedder + VLM picture descriptions diff --git a/docs/capabilities/analysis.md b/docs/capabilities/analysis.md index 19fbe4d6..5effc127 100644 --- a/docs/capabilities/analysis.md +++ b/docs/capabilities/analysis.md @@ -46,6 +46,8 @@ async with HaikuRAG("my.lancedb") as client: ## State -When dependencies expose a state dictionary, `AnalysisState` is stored under `"analysis"`. It contains the document filter, code execution log, searches, and citations. Per-run searches and executions reset automatically; the filter and citation index persist. +When dependencies expose a state dictionary, `AnalysisState` is stored under `"analysis"`. It contains the document filter, code execution log, searches, citations, and the `evidence` record of what was retrieved and cited per question. Searches and executions are cleared when a new question starts, and a resumed question keeps them; the filter, citation index and evidence record persist. + +This capability does not alter the message history either. Register the [compaction capability](index.md#multi-turn-conversations) to compact earlier questions. The capability lazily opens both LanceDB and the sandbox only after it is loaded and a tool requires them. Resources close at the end of the agent run. diff --git a/docs/capabilities/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/docs/capabilities/rag.md b/docs/capabilities/rag.md index 00a912bb..35c3fbf2 100644 --- a/docs/capabilities/rag.md +++ b/docs/capabilities/rag.md @@ -37,16 +37,19 @@ class RAGState(BaseModel): citation_index: dict[str, Citation] citations: list[str] document_filter: str | None + evidence: CapabilityEvidenceRecord searches: dict[str, list[SearchResult]] ``` -`document_filter` persists between runs. Current citations and searches reset for each run, while the citation index remains available to the host application. +`document_filter`, `citation_index` and `evidence` persist across runs. Citations and searches are cleared when a new question starts; a run that resumes a question keeps the evidence it is still answering from. + +`evidence` records which chunks this capability retrieved and cited, and in which question. `haiku.rag.capabilities.ledger.citation_status(records, question=...)` derives `missing`, `grounded` or `ungrounded` from it, across capabilities. State is ordinary application state; the capability does not depend on AG-UI. An AG-UI application can expose it using Pydantic AI's standard adapter. ## Context management -Large RAG tool results from earlier user turns are replaced with a short marker before model requests. Tool-call pairing and current-turn evidence are retained. This prevents long conversations from repeatedly sending old retrieved content. +This capability does not alter the message history. To stop long conversations resending old retrieved content, register the [compaction capability](index.md#multi-turn-conversations) alongside it. ## Domain context and vision diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index ff7a2b0a..65c43ff3 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py @@ -15,8 +15,6 @@ from pydantic_ai.messages import ( ModelResponse, ToolCallPart, ToolReturn, - ToolReturnPart, - UserPromptPart, ) from pydantic_ai.models import ModelRequestContext from pydantic_ai.run import AgentRunResult @@ -24,11 +22,12 @@ from pydantic_ai.tools import ToolDefinition from pydantic_ai.toolsets import AgentToolset from haiku.rag.capabilities._tools import CodeExecutionEntry, search_corpus +from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord, EvidenceRef from haiku.rag.client import HaikuRAG from haiku.rag.config.models import AppConfig from haiku.rag.store.models.chunk import SearchResult from haiku.rag.store.models.citation import Citation, resolve_citations -from haiku.rag.tools.search import build_binary_parts_from_results +from haiku.rag.tools.search import build_image_content_from_results CITATION_GRACE_REQUESTS = 2 """Requests calling this capability's tools that its cite tool outlives the rest by. @@ -70,47 +69,42 @@ def resolve_db_path(db_path: Path | None, config: AppConfig) -> Path: def _clear_invocation_state(state: BaseModel) -> None: + """Drop the working evidence of the previous question. + + Only ever called when a new question starts. A resumption keeps it: the + results belong to the question still being answered, and dropping them leaves + a later citation unable to resolve against the expanded result the model saw, + recording no provenance for it. + """ for field_name in ("citations", "searches", "executions"): value = getattr(state, field_name, None) if hasattr(value, "clear"): value.clear() -def _compact_old_tool_returns( - messages: list[ModelMessage], tool_names: frozenset[str] -) -> list[ModelMessage]: - """Remove bulky prior-turn evidence while retaining the current turn. +def _is_resumption(prompt: Any, messages: list[ModelMessage]) -> bool: + """Whether this run continues a question rather than asking a new one. - Tool call and return parts remain paired; only the old return payload is - replaced. This keeps provider histories valid and preserves all evidence - gathered since the most recent user prompt. + Two signals, either of which is enough, because getting this wrong hands the + model a notice where its own evidence should be: + + - no prompt: how pydantic-ai resumes for interruptions and suspensions. + - an unfinished tail: the history ends with a request the model has not + answered, or with a response whose tool calls have no returns yet. Deferred + tool results may arrive *with* a prompt, so the prompt alone is not enough. + + A settled history ends with the previous answer, so a genuinely new question + is not mistaken for a continuation. The framework's own first-new-message + index would be better than either signal, but it is not public here. """ - latest_user_message = -1 - for index, message in enumerate(messages): - if isinstance(message, ModelRequest) and any( - isinstance(part, UserPromptPart) for part in message.parts - ): - latest_user_message = index - - if latest_user_message < 0: - return messages - - compacted = list(messages) - for index, message in enumerate(messages[:latest_user_message]): - if not isinstance(message, ModelRequest): - continue - parts = [ - replace( - part, - content="[Prior-turn RAG tool output removed; citations remain in state.]", - ) - 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 + if prompt is None: + return True + if not messages: + return False + last = messages[-1] + if isinstance(last, ModelRequest): + return True + return any(isinstance(part, ToolCallPart) for part in last.parts) def _called_own_tool(messages: list[ModelMessage], tool_names: frozenset[str]) -> bool: @@ -142,13 +136,38 @@ 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) + epoch: int = field(default=0, repr=False) async def for_run(self, ctx: RunContext[Any]) -> "RAGCapabilityBase[StateT]": + """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 + the identity already recorded: the question is the one in progress, and + adopting the current count would relabel it as a new one and judge its + declarations against the wrong question. A resumption with no recorded + identity is a state this design does not produce, so it is reported rather + than guessed at. With no history at all there is nothing in progress: an + absent prompt is then an instructions-only first question, which takes an + identity like any other. + """ outer = getattr(ctx.deps, "state", None) outer_state = outer if isinstance(outer, dict) else None raw_state = outer_state.get(self.state_namespace) if outer_state else None + resuming = _is_resumption(ctx.prompt, ctx.messages) + continuing = resuming and bool(ctx.messages) state = self.state_type.model_validate(raw_state or {}) - _clear_invocation_state(state) + record = cast(CapabilityEvidenceRecord, cast(Any, state).evidence) + if continuing and record.question is None: + raise RuntimeError( + f"The {self.state_namespace} capability is resuming a question with " + "no stored question identity. Capabilities cannot be added, removed " + "or migrated while a question is unfinished, and the run's state " + "must be carried between its runs." + ) + if not continuing: + _clear_invocation_state(state) + record.begin_question(len(ctx.messages)) run_capability = replace( self, state=state, @@ -159,6 +178,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): search_count=0, request_count=0, grace_requests_used=0, + epoch=0, ) run_capability._sync_state() return run_capability @@ -171,9 +191,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): async def before_model_request( self, ctx: RunContext[Any], request_context: ModelRequestContext ) -> ModelRequestContext: - request_context.messages = _compact_old_tool_returns( - request_context.messages, self.tool_names - ) + self.epoch = len(ctx.messages) if instruction := self._budget_notice(): current_request = request_context.messages[-1] if isinstance(current_request, ModelRequest): @@ -217,7 +235,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): ) if spent := self._spent_tool_names(): names = ", ".join(sorted(spent)) - if remaining := sorted(self._evidence_tool_names() - spent): + if remaining := sorted(self.evidence_tool_names() - spent): return ( f"The {self.state_namespace} capability has spent its budget " f"for {names}; further calls to them fail. Gather any further " @@ -255,8 +273,13 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): if tool.capability_id != self.id or tool.name == self._cite_tool_name ] - def _evidence_tool_names(self) -> set[str]: - """Tools that can bring new evidence into the run.""" + def evidence_tool_names(self) -> set[str]: + """Tools that can bring new evidence into the run. + + Public because compaction needs to know whose output on the wire is + evidence: a cite acknowledgement is a receipt of the model's own action and + must survive, while a code execution that reached the corpus is evidence. + """ return {f"{self.state_namespace}_search"} def _spent_tool_names(self) -> set[str]: @@ -306,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) @@ -326,6 +361,41 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): finally: self._sync_state() + def _evidence_record(self) -> CapabilityEvidenceRecord: + assert self.state is not None + return cast(CapabilityEvidenceRecord, cast(Any, self.state).evidence) + + def _note_evidence(self) -> None: + """Record an outcome the model can ground an answer on. + + Includes an empty search result and a failed execution that still printed + output: negative evidence grounds a refusal. Excludes a spent budget, which + yields nothing to ground anything on. + """ + self._evidence_record().note_evidence(self.epoch) + + def _declare(self, citations: list[Citation]) -> None: + """Record what the model cited, once the ids have resolved. + + Declaring earlier would let a call naming only unresolvable ids read as a + grounded answer. + """ + state = cast(Any, self.state) + retrieved = { + result.chunk_id + for results in state.searches.values() + for result in results + if result.chunk_id + } + self._evidence_record().declare( + [ + EvidenceRef(capability=self.state_namespace, chunk_id=c.chunk_id) + for c in citations + ], + epoch=self.epoch, + retrieved_now=retrieved, + ) + async def _search(self, query: str, limit: int | None) -> str | ToolReturn: assert self.state is not None self.search_count += 1 @@ -343,7 +413,8 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): ) state = cast(Any, self.state) state.searches[query] = results - if self.vision and (parts := build_binary_parts_from_results(results)): + self._note_evidence() + if self.vision and (parts := build_image_content_from_results(results)): return ToolReturn(return_value=formatted, content=parts) return formatted @@ -390,6 +461,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): "Copy chunk_ids verbatim from search results." ) self._register_citations(citations) + self._declare(citations) resolved = {citation.chunk_id for citation in citations} unresolved = [cid for cid in missing if cid not in resolved] if unresolved: diff --git a/haiku_rag_slim/haiku/rag/capabilities/analysis.py b/haiku_rag_slim/haiku/rag/capabilities/analysis.py index 95fa155e..e2de8868 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/analysis.py +++ b/haiku_rag_slim/haiku/rag/capabilities/analysis.py @@ -13,6 +13,7 @@ from haiku.rag.capabilities._base import ( RAGCapabilityBase, resolve_db_path, ) +from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord from haiku.rag.config.models import AppConfig from haiku.rag.sandbox import AnalysisContext, Sandbox from haiku.rag.store.models.chunk import SearchResult @@ -29,6 +30,7 @@ class AnalysisState(BaseModel): executions: list[CodeExecutionEntry] = Field(default_factory=list) citation_index: dict[str, Citation] = Field(default_factory=dict) citations: list[str] = Field(default_factory=list) + evidence: CapabilityEvidenceRecord = Field(default_factory=CapabilityEvidenceRecord) searches: dict[str, list[SearchResult]] = Field(default_factory=dict) @@ -85,11 +87,11 @@ class AnalysisCapability(RAGCapabilityBase[AnalysisState]): self.sandbox = None await super()._close() - def _evidence_tool_names(self) -> set[str]: + def evidence_tool_names(self) -> set[str]: # Searching from inside the sandbox does not count against # `qa.max_searches`, so code execution outlives a spent search budget # as a way to reach new evidence. - return super()._evidence_tool_names() | {"analysis_execute_code"} + return super().evidence_tool_names() | {"analysis_execute_code"} def _spent_tool_names(self) -> set[str]: spent = super()._spent_tool_names() @@ -107,6 +109,8 @@ class AnalysisCapability(RAGCapabilityBase[AnalysisState]): ) sandbox = await self._ensure_sandbox() result = await sandbox.execute(code) + if result.success or result.stdout: + self._note_evidence() if sandbox._search_results: existing = self.state.searches.get("_sandbox", []) seen = {item.chunk_id for item in existing} diff --git a/haiku_rag_slim/haiku/rag/capabilities/compaction.py b/haiku_rag_slim/haiku/rag/capabilities/compaction.py new file mode 100644 index 00000000..e2c051f0 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/capabilities/compaction.py @@ -0,0 +1,445 @@ +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field, replace +from typing import Any, cast + +from pydantic_ai import RunContext +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" + +CAPSULE_HEADER = ( + "[Evidence cited earlier in this conversation, kept so later questions can " + "rely on it. Cite these chunk_ids directly when you use them.]" +) + +RECEIPT = ( + "[Evidence retrieved for an earlier question, no longer shown. It does not " + "count as cited for the current question.]" +) + +ENTRY_SEPARATOR = "\n\n" + + +def group_label(position: int) -> str: + """Name a group by its position among the groups, not by question number. + + A question identity is a message count, so a header built from it would present + an index as a turn number, and an ordinal over the groups is not the + conversation's ordinal either whenever a question in between cited nothing. The + label claims only what it is: a grouping, newest first. + """ + return f"[Cited evidence group {position}]" + + +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. {RETRIEVED_IMAGE_TAG}" + ) + + +@dataclass(frozen=True) +class DiscoveredEvidence: + """One evidence capability's records, as the compactor found them. + + Read-only and rebuilt per request: the compactor merges these into a view and + persists nothing about evidence itself. + """ + + capability: str + record: CapabilityEvidenceRecord + citations: Mapping[str, Citation] + tool_names: frozenset[str] + + +@dataclass(frozen=True) +class RetainedPicture: + """A picture to re-attach, with the label that must accompany it. + + Addressed by owner, document and reference, because a reference such as + ``#/pictures/0`` repeats across documents and capabilities. The label travels + with it so it can never be emitted without its image. + """ + + capability: str + chunk_id: str + document_id: str + self_ref: str + label: str + + +@dataclass(frozen=True) +class Capsule: + """Everything the compactor would insert, and nothing about where it goes.""" + + text: str = "" + pictures: tuple[RetainedPicture, ...] = () + + +@dataclass(frozen=True) +class _Entry: + capability: str + chunk_id: str + question: int + citation: Citation + + def render(self) -> str: + title = self.citation.document_title + uri = self.citation.document_uri + source = f'"{title}"' if title else uri + if title and uri and uri != title: + source = f"{source} ({uri})" + return f"[{self.chunk_id}] Source: {source}\n{self.citation.content}" + + +def _eligible_entries(evidence: Sequence[DiscoveredEvidence]) -> list[_Entry]: + """Cited evidence with content, newest citing question first. + + Evidence cited in several questions belongs to the most recent one, so it is + rendered once and grouped where the model last used it. + + An occurrence and its canonical ``Citation`` are written by the same call, so a + cited chunk without one is not a state this design produces. Rendering the rest + regardless would quietly drop evidence an answer rested on, so it is reported. + """ + entries = [] + for discovered in evidence: + for chunk_id, occurrence in discovered.record.occurrences.items(): + if not occurrence.cited_in_questions: + continue + citation = discovered.citations.get(chunk_id) + if citation is None: + raise ValueError( + f"{discovered.capability} cited {chunk_id} in question(s) " + f"{occurrence.cited_in_questions} but has no citation record " + "for it, so its content cannot be retained." + ) + entries.append( + _Entry( + capability=discovered.capability, + chunk_id=chunk_id, + question=max(occurrence.cited_in_questions), + citation=citation, + ) + ) + entries.sort(key=lambda entry: (-entry.question, entry.capability, entry.chunk_id)) + return entries + + +def build_capsule(evidence: Sequence[DiscoveredEvidence]) -> Capsule: + """Render every cited piece of evidence, grouped by the question that cited it. + + Everything cited is kept whole and everything else is dropped. There is no + character budget: what a model can hold is the model's business, and a knob for + it would only half-rescue models that fail on long conversations regardless. + + A host that needs earlier evidence pruned can compact further on top, on the wire + only. Removing or reordering the stored history breaks the message counts that + question identities and epochs are derived from, and the next record written is + refused. + + Pure: no I/O and no message history, so what goes on the wire stays separable + from what it should contain. Picture bytes are fetched by the caller, which is + why a picture travels with its label rather than beside it. + """ + entries = _eligible_entries(evidence) + if not entries: + return Capsule() + + lines = [CAPSULE_HEADER] + pictures: list[RetainedPicture] = [] + seen: set[tuple[str, str, str]] = set() + position = 0 + current_question: int | None = None + for entry in entries: + if entry.question != current_question: + position += 1 + current_question = entry.question + lines.append(group_label(position)) + lines.append(entry.render()) + for self_ref in entry.citation.picture_refs: + # Overlapping chunks cite one figure, and a provider counts it twice. + # Identity is owner plus document plus reference, so the same reference + # in another document stays a different picture. + identity = (entry.capability, entry.citation.document_id, self_ref) + if identity in seen: + continue + seen.add(identity) + pictures.append( + RetainedPicture( + capability=entry.capability, + chunk_id=entry.chunk_id, + document_id=entry.citation.document_id, + self_ref=self_ref, + label=picture_label(entry.chunk_id, self_ref), + ) + ) + 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 position, part in enumerate(message.parts): + if isinstance(part, ToolReturnPart) and part.tool_name in owned_tools: + carries = (index, position) == carrier + body = capsule_text or RECEIPT if carries 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 carrier is not None and index == carrier[0] 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] +) -> tuple[int, int] | None: + """Where the last of our evidence returns is, as message and part. + + The part matters: a model can call search twice in one response, so one request + can hold several of our returns, and giving the capsule to each duplicates the + whole of it. + """ + for index in range(min(boundary, len(messages)) - 1, -1, -1): + message = messages[index] + if not isinstance(message, ModelRequest): + continue + for position in range(len(message.parts) - 1, -1, -1): + part = message.parts[position] + if isinstance(part, ToolReturnPart) and part.tool_name in owned_tools: + return index, position + return None + + +@dataclass +class EvidenceCompactionCapability(AbstractCapability[Any]): + """Rewrites the history from what the evidence capabilities recorded. + + Registering it is what turns compaction on: a host that leaves it out gets an + untouched transcript, which is why it has no enable flag. It reads the evidence + capabilities through the run's registry and holds no reference to any of them, + so a host running one capability, both, or neither needs no wiring change. + + Registering two is rejected by pydantic-ai before the run starts, since they + 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. + + The registry holds the per-run instances, which are the ones carrying + state; the registered objects never do. That includes a deferred capability + the model has not loaded, whose record is simply empty. + """ + discovered = [] + for capability in ctx.capabilities.values(): + if not isinstance(capability, RAGCapabilityBase): + continue + state = capability.state + discovered.append( + DiscoveredEvidence( + capability=capability.state_namespace, + record=cast(CapabilityEvidenceRecord, cast(Any, state).evidence), + citations=cast(Any, state).citation_index, + tool_names=frozenset(capability.evidence_tool_names()), + ) + ) + return sorted(discovered, key=lambda evidence: evidence.capability) + + +def create_capability() -> EvidenceCompactionCapability: + """Create the capability that compacts history from recorded evidence.""" + return EvidenceCompactionCapability( + id=CAPABILITY_ID, + description=( + "Replaces earlier questions' evidence on the wire with a capsule of " + "what was cited." + ), + ) + + +__all__ = [ + "CAPABILITY_ID", + "CAPSULE_HEADER", + "RECEIPT", + "Capsule", + "DiscoveredEvidence", + "EvidenceCompactionCapability", + "RetainedPicture", + "build_capsule", + "compact_history", + "create_capability", + "group_label", + "picture_label", +] diff --git a/haiku_rag_slim/haiku/rag/capabilities/ledger.py b/haiku_rag_slim/haiku/rag/capabilities/ledger.py new file mode 100644 index 00000000..2750bd6c --- /dev/null +++ b/haiku_rag_slim/haiku/rag/capabilities/ledger.py @@ -0,0 +1,213 @@ +from collections.abc import Iterable +from typing import Literal + +from pydantic import BaseModel, Field + +CitationStatus = Literal["missing", "grounded", "ungrounded"] + + +class EvidenceRef(BaseModel): + """One piece of evidence, identified by its owner as well as its chunk. + + A chunk id alone is not an identity: the same id can be reported by more than + one capability, and ownership is what tells compaction whose output it may + touch. + """ + + capability: str + chunk_id: str + + +class EvidenceOccurrence(BaseModel): + """Which questions retrieved a piece of evidence, and which cited it.""" + + capability: str + chunk_id: str + retrieved_in_questions: list[int] = Field(default_factory=list) + cited_in_questions: list[int] = Field(default_factory=list) + + +class CitationDeclaration(BaseModel): + """What a question declared as its grounding, and when. + + Bound to a question *and* an epoch: the epoch outlives a question, so matching + it alone would let a question that gathered no evidence inherit the previous + declaration and read as compliant having declared nothing. + """ + + question: int + epoch: int + refs: list[EvidenceRef] = Field(default_factory=list) + + +class CapabilityEvidenceRecord(BaseModel): + """What one evidence capability wrote, in its own state namespace. + + Holds no content: ``Citation`` in ``citation_index`` is the canonical record + and already persists content, document id and picture refs. This is the index + over it that compaction needs and the transcript cannot provide. + + Single-writer by construction. A record shared between capabilities would be + overwritten by whichever of them synced its state last; merging happens in the + transient views built by ``citation_status`` and the optional capabilities. + + ``question`` is the number of messages that existed when the question arrived, + and ``epoch`` the number when an outcome occurred. Both are derived from the + conversation rather than counted locally, so every participant computes the + same values without sharing a counter. ``question`` is unset until a run + establishes it, so a record a host merely created is distinguishable from one + that has been through a question. + """ + + occurrences: dict[str, EvidenceOccurrence] = Field(default_factory=dict) + question: int | None = None + latest_evidence_epoch: int = 0 + declaration: CitationDeclaration | None = None + + def _reject_regression(self, count: int, what: str) -> None: + """Refuse a message count below one already recorded in this question. + + Identities and epochs are both message counts, and every comparison + between them assumes the conversation only grows while a question is + being answered. One capability truncating or reordering the history + breaks that, and each way of recording it has to refuse the same way: an + unchecked evidence outcome freezes every later declaration as stale, while + an unchecked declaration replaces a newer one with an older one and + revives the answer it grounded. + """ + recorded = max( + self.question or 0, + self.latest_evidence_epoch, + self.declaration.epoch if self.declaration else 0, + ) + if count < recorded: + raise ValueError( + f"{what} at message count {count} is behind {recorded}, which is " + "already recorded: message history must be append-only for " + "question identities and epochs to hold." + ) + + def begin_question(self, identity: int) -> None: + """Take the identity of a question that has just arrived. + + The evidence epoch and the declaration describe the question that has + just ended, and outlive their meaning the moment the next one begins: + kept, they let the last question's citations ground this one, and hold a + horizon this question's own declarations cannot pass. Clearing them + confines those comparisons to one question, so a host whose stored + history shifted between two of them is answered rather than refused. + + Identities themselves must still separate one question from the next. + Occurrences outlive the question that recorded them and carry the + identities that cited them, which compaction groups and orders the capsule + by: a reused identity merges two questions into one group, and a lower one + renders later evidence as though it were cited earlier. + """ + if self.question is not None and identity <= self.question: + raise ValueError( + f"A question at message count {identity} is not past question " + f"{self.question}, which is already answered: identities separate " + "one question from the next and are compared as recency." + ) + self.question = identity + self.latest_evidence_epoch = 0 + self.declaration = None + + def note_evidence(self, epoch: int) -> None: + """Record that the model has seen an evidence outcome. + + Called for anything an answer could rest on, including a search that + returned nothing and a failed execution that still printed output — a + fruitless search grounds a refusal. Not called for a failure that yields + no evidence at all, such as an exhausted budget. + """ + self._reject_regression(epoch, "Evidence") + self.latest_evidence_epoch = epoch + + def declare( + self, + refs: list[EvidenceRef], + *, + epoch: int, + retrieved_now: set[str] | None = None, + ) -> None: + """Record validated citations for the current question. + + Repeated calls at the same epoch merge, so citing again cannot narrow what + was already declared: an empty call after a grounded one leaves it + grounded. A call at a later epoch declares afresh, because evidence the + model saw in between may be what it is now citing. + """ + if self.question is None: + raise ValueError( + "Citations cannot be declared before a run establishes the " + "question identity." + ) + self._reject_regression(epoch, "A declaration") + current = self.declaration + if current is not None and (current.question, current.epoch) == ( + self.question, + epoch, + ): + known = {(ref.capability, ref.chunk_id) for ref in current.refs} + current.refs.extend( + ref for ref in refs if (ref.capability, ref.chunk_id) not in known + ) + else: + self.declaration = CitationDeclaration( + question=self.question, epoch=epoch, refs=list(refs) + ) + + for ref in refs: + occurrence = self.occurrences.setdefault( + ref.chunk_id, + EvidenceOccurrence(capability=ref.capability, chunk_id=ref.chunk_id), + ) + if self.question not in occurrence.cited_in_questions: + occurrence.cited_in_questions.append(self.question) + if ( + retrieved_now + and ref.chunk_id in retrieved_now + and self.question not in occurrence.retrieved_in_questions + ): + occurrence.retrieved_in_questions.append(self.question) + + +def citation_status( + records: Iterable[CapabilityEvidenceRecord], *, question: int +) -> CitationStatus: + """Derived, never stored, so refs and status cannot contradict. + + A declaration is current only for the question it was made in, and only if it + followed the newest evidence outcome of *every* capability: a question where + one capability cited and another then searched without citing is not grounded. + Strictly later, since a citation made in the same request as an outcome cannot + have read it. + + A grounding *violation* is not one of these: that is an enforcement outcome + recorded by the policy capability, not something the model declared. + """ + records = list(records) + horizon = max((record.latest_evidence_epoch for record in records), default=0) + current = [ + record.declaration + for record in records + if record.declaration is not None + and record.declaration.question == question + and record.declaration.epoch > horizon + ] + if not current: + return "missing" + return ( + "grounded" if any(declaration.refs for declaration in current) else "ungrounded" + ) + + +__all__ = [ + "CapabilityEvidenceRecord", + "CitationDeclaration", + "CitationStatus", + "EvidenceOccurrence", + "EvidenceRef", + "citation_status", +] diff --git a/haiku_rag_slim/haiku/rag/capabilities/rag.py b/haiku_rag_slim/haiku/rag/capabilities/rag.py index 08931954..79216a87 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/rag.py +++ b/haiku_rag_slim/haiku/rag/capabilities/rag.py @@ -12,6 +12,7 @@ from haiku.rag.capabilities._base import ( RAGCapabilityBase, resolve_db_path, ) +from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord from haiku.rag.config.models import AppConfig from haiku.rag.store.models.chunk import SearchResult from haiku.rag.store.models.citation import Citation @@ -33,6 +34,7 @@ _instructions_path = Path(__file__).parent / "instructions" / "rag.md" class RAGState(BaseModel): citation_index: dict[str, Citation] = Field(default_factory=dict) citations: list[str] = Field(default_factory=list) + evidence: CapabilityEvidenceRecord = Field(default_factory=CapabilityEvidenceRecord) document_filter: str | None = None searches: dict[str, list[SearchResult]] = Field(default_factory=dict) diff --git a/haiku_rag_slim/haiku/rag/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py index 0350e5d4..beccebcd 100644 --- a/haiku_rag_slim/haiku/rag/chat/app.py +++ b/haiku_rag_slim/haiku/rag/chat/app.py @@ -1,6 +1,7 @@ import asyncio import uuid from collections.abc import Iterable, Sequence +from copy import deepcopy from dataclasses import dataclass, field from pathlib import Path from typing import TYPE_CHECKING, Any @@ -25,6 +26,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 +155,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: @@ -205,7 +209,12 @@ class ChatApp(App): await chat_history.show_thinking() message = None - deps = ChatDeps(state=self._state) + # The run gets a copy: state and message history have to advance together. + # A cancelled or failed run discards its messages, and state that advanced + # anyway would leave the next question deriving its identity from a shorter + # history than the evidence already recorded — refused as non-append-only, + # with the conversation stuck until it is cleared. + deps = ChatDeps(state=deepcopy(self._state)) try: async with self._agent.run_stream_events( diff --git a/haiku_rag_slim/haiku/rag/store/models/chunk.py b/haiku_rag_slim/haiku/rag/store/models/chunk.py index b2936e15..95666a39 100644 --- a/haiku_rag_slim/haiku/rag/store/models/chunk.py +++ b/haiku_rag_slim/haiku/rag/store/models/chunk.py @@ -211,7 +211,7 @@ class SearchResult(BaseModel): parts.append(f"Type: {primary_label}") # Surface picture captions when present. Order matches the binary - # attachments emitted by build_binary_parts_from_results, so the model + # attachments emitted by build_image_content_from_results, so the model # can correlate caption ↔ attached image by position (BinaryContent # identifiers don't survive serialization to the OpenAI vision API). if self.picture_captions: diff --git a/haiku_rag_slim/haiku/rag/tools/search.py b/haiku_rag_slim/haiku/rag/tools/search.py index 4e8bf8e5..4f192a74 100644 --- a/haiku_rag_slim/haiku/rag/tools/search.py +++ b/haiku_rag_slim/haiku/rag/tools/search.py @@ -10,11 +10,35 @@ 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. -def build_binary_parts_from_results( +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], -) -> list[BinaryContent]: - """Decode and validate picture bytes attached to search results. +) -> list[str | BinaryContent]: + """Decode and validate picture bytes attached to search results, labelled. Dedup keyed on ``(document_id, self_ref)`` so the same picture in different chunks is sent once. Pictures that fail @@ -22,8 +46,18 @@ def build_binary_parts_from_results( vision placeholder per ``BinaryContent``, so emitting one for an image the server can't decode leaves the processor with an off-by-one count. + + Every picture is preceded by a line naming the result it belongs to. + ``ToolReturn.content`` reaches the model as a user-role message, so + retrieved pictures are otherwise indistinguishable from ones the user + attached, and models narrate them as part of the question: unlabelled, + gemma4-26b answered about a figure from an unrelated document, and with a + single note ahead of the batch it still called them "images in the prompt". + The label also names the chunk to cite for a figure, which + ``BinaryContent.identifier`` cannot do — it does not survive serialization + to the vision API. """ - parts: list[BinaryContent] = [] + collected: list[tuple[str | None, str, BinaryContent]] = [] seen: set[tuple[str | None, str]] = set() for result in results: if not result.image_data: @@ -32,21 +66,22 @@ def build_binary_parts_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 - parts.append( - BinaryContent( - data=data, - media_type="image/png", - identifier=self_ref, - ) - ) + collected.append((result.chunk_id, self_ref, picture)) seen.add(key) - return parts + + content: list[str | BinaryContent] = [] + total = len(collected) + for position, (chunk_id, self_ref, picture) in enumerate(collected, 1): + content.append( + 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 def create_search_toolset( @@ -134,9 +169,9 @@ def create_search_toolset( if not config.qa.model.vision: return text - binary_parts = build_binary_parts_from_results(results_list) - if binary_parts: - return ToolReturn(return_value=text, content=binary_parts) + image_content = build_image_content_from_results(results_list) + if image_content: + return ToolReturn(return_value=text, content=image_content) return text toolset: FunctionToolset[RAGDeps] = FunctionToolset() diff --git a/tests/capabilities/test_capabilities.py b/tests/capabilities/test_capabilities.py index 6bbf6141..e8af1d14 100644 --- a/tests/capabilities/test_capabilities.py +++ b/tests/capabilities/test_capabilities.py @@ -4,7 +4,7 @@ from typing import Any, cast from unittest.mock import AsyncMock, patch import pytest -from pydantic_ai import Agent, ModelRetry, RunContext, ToolFailed +from pydantic_ai import Agent, DeferredToolResults, ModelRetry, RunContext, ToolFailed from pydantic_ai.messages import ( ModelRequest, ModelResponse, @@ -20,10 +20,13 @@ from pydantic_ai.usage import RunUsage from haiku.rag.capabilities._base import ( CITATION_GRACE_REQUESTS, _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 +from haiku.rag.capabilities.ledger import ( + CapabilityEvidenceRecord, + citation_status, +) from haiku.rag.capabilities.rag import AGENT_PREAMBLE, RAGCapability, RAGState from haiku.rag.capabilities.rag import create_capability as create_rag from haiku.rag.config.models import AppConfig, PromptsConfig @@ -308,7 +311,7 @@ async def test_search_and_empty_citation_limits(temp_db_path): @pytest.mark.asyncio async def test_cite_resolves_direct_chunk_ids_and_reuses_document_lookup(temp_db_path): capability = create_rag(db_path=temp_db_path, config=AppConfig()) - capability.state = RAGState() + capability.state = RAGState(evidence=CapabilityEvidenceRecord(question=0)) client = AsyncMock() client.get_chunk_by_id.side_effect = [ Chunk(id="chunk-1", document_id="doc-1", content="first"), @@ -334,7 +337,7 @@ async def test_cite_resolves_direct_chunk_ids_and_reuses_document_lookup(temp_db @pytest.mark.asyncio async def test_cite_reports_unresolved_ids_on_partial_success(temp_db_path): capability = create_rag(db_path=temp_db_path, config=AppConfig()) - capability.state = RAGState() + capability.state = RAGState(evidence=CapabilityEvidenceRecord(question=0)) client = AsyncMock() client.get_chunk_by_id.side_effect = [ Chunk(id="chunk-1", document_id="doc-1", content="first"), @@ -364,6 +367,7 @@ async def test_cite_repairs_chunk_ids_damaged_in_transcription(temp_db_path): unrelated = "9c2cd07e-5a3f-45a6-968d-cbd6f06ab57b" capability = create_rag(db_path=temp_db_path, config=AppConfig()) capability.state = RAGState( + evidence=CapabilityEvidenceRecord(question=0), searches={ "q": [ SearchResult( @@ -374,7 +378,7 @@ async def test_cite_repairs_chunk_ids_damaged_in_transcription(temp_db_path): document_uri="test://document", ) ] - } + }, ) client = AsyncMock() client.get_chunk_by_id.return_value = None @@ -461,6 +465,58 @@ async def test_analysis_execution_limit_fails_the_tool(temp_db_path): await capability._execute_code("print('done')") +@pytest.mark.asyncio +async def test_a_spent_execution_budget_is_not_evidence(temp_db_path): + """Nothing was produced to ground an answer on, so nothing is recorded.""" + config = AppConfig() + config.analysis.max_executions = 0 + capability = create_analysis(db_path=temp_db_path, config=config) + capability.state = AnalysisState() + capability.epoch = 5 + + with pytest.raises(ToolFailed): + await capability._execute_code("print('done')") + + assert capability.state.evidence.latest_evidence_epoch == 0 + + +@pytest.mark.parametrize( + ("success", "stdout", "expected_epoch"), + [ + pytest.param(True, "42", 5, id="succeeded"), + pytest.param(True, "", 5, id="succeeded without output"), + pytest.param(False, "42", 5, id="failed after printing"), + pytest.param(False, "", 0, id="failed without printing"), + ], +) +@pytest.mark.asyncio +async def test_only_a_code_execution_the_model_can_read_is_evidence( + temp_db_path, success, stdout, expected_epoch +): + """A raised error with nothing printed grounds nothing, so it is not evidence. + + A failure that printed first does ground an answer, and so does a successful + run whose outcome is that it printed nothing. + """ + capability = create_analysis(db_path=temp_db_path, config=AppConfig()) + capability.state = AnalysisState(evidence=CapabilityEvidenceRecord(question=0)) + capability.epoch = 5 + sandbox = AsyncMock(spec=Sandbox) + sandbox._search_results = [] + sandbox.execute.return_value = SandboxResult( + stdout=stdout, stderr="" if success else "boom", success=success + ) + capability.sandbox = sandbox + + if success: + await capability._execute_code("print(42)") + else: + with pytest.raises(ToolFailed): + await capability._execute_code("print(42)") + + assert capability.state.evidence.latest_evidence_epoch == expected_epoch + + @pytest.mark.asyncio async def test_spent_search_budget_is_announced_but_keeps_the_tool(rag_db): """A spent budget is announced; the tool stays declared to avoid a dead run. @@ -555,7 +611,7 @@ async def test_spent_search_notice_points_at_code_while_it_has_budget(temp_db_pa notice = capability._budget_notice() assert notice is not None assert "analysis_execute_code" in notice - assert capability._evidence_tool_names() <= capability._spent_tool_names() + assert capability.evidence_tool_names() <= capability._spent_tool_names() @pytest.mark.asyncio @@ -570,7 +626,7 @@ async def test_spent_search_notice_tells_rag_to_answer(temp_db_path): assert notice is not None assert "rag_search" in notice - assert capability._evidence_tool_names() == {"rag_search"} + assert capability.evidence_tool_names() == {"rag_search"} @pytest.mark.asyncio @@ -746,7 +802,9 @@ async def test_native_agent_composition_initializes_host_state(temp_db_path): result = await agent.run("Hello", deps=deps) assert result.output == "success (no tool calls)" - assert deps.state["rag"] == RAGState().model_dump(mode="json") + assert deps.state["rag"] == RAGState( + evidence=CapabilityEvidenceRecord(question=0) + ).model_dump(mode="json") @pytest.mark.asyncio @@ -787,41 +845,400 @@ 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")]), +def _resuming_deps() -> Deps: + """State as a resumption always finds it: the question already identified. + + A run that resumes has been through ``for_run`` before, so the identity of the + question in progress is stored. Fabricating the history without it is a state + the design does not produce, and is rejected rather than guessed at. + """ + return Deps( + state={ + "rag": RAGState(evidence=CapabilityEvidenceRecord(question=0)).model_dump( + mode="json" + ) + } + ) + + +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", "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")] + parts=[ToolReturnPart("rag_search", "EVIDENCE FOR THE LIVE TURN", "call-1")] ), ] - compacted = _compact_old_tool_returns(messages, frozenset({"rag_search"})) - old_return = compacted[2].parts[0] - current_return = compacted[5].parts[0] - assert isinstance(old_return, ToolReturnPart) - assert "removed" in str(old_return.content) - assert isinstance(current_return, ToolReturnPart) - assert current_return.content == "current evidence" +def _record(deps: Deps, namespace: str) -> CapabilityEvidenceRecord: + return CapabilityEvidenceRecord.model_validate(deps.state[namespace]["evidence"]) -def test_tool_results_are_unchanged_when_history_has_no_user_prompt(): - messages = [ - ModelResponse(parts=[ToolCallPart("rag_search", {}, "current-call")]), - ModelRequest( - parts=[ToolReturnPart("rag_search", "current evidence", "current-call")] - ), +async def _stub_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") + ] + self._note_evidence() + return "EVIDENCE" + + +@pytest.mark.asyncio +async def test_a_question_takes_its_own_identity_and_both_capabilities_agree( + temp_db_path, +): + """Identity is derived from the conversation, so no counter is shared.""" + rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False) + analysis = create_analysis( + db_path=temp_db_path, config=AppConfig(), defer_loading=False + ) + + async def model(_messages, _info): + return ModelResponse(parts=[TextPart("answer")]) + + agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag, analysis]) + deps = Deps() + + first = await agent.run("first question", deps=deps) + first_identity = _record(deps, "rag").question + await agent.run("second question", deps=deps, message_history=first.all_messages()) + second_identity = _record(deps, "rag").question + + assert first_identity == 0 + assert second_identity is not None and second_identity > 0 + assert _record(deps, "analysis").question == second_identity + + +@pytest.mark.asyncio +async def test_a_resumption_keeps_the_identity_of_the_question_in_progress( + temp_db_path, +): + rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False) + + async def model(_messages, _info): + return ModelResponse(parts=[TextPart("answer")]) + + agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag]) + deps = Deps( + state={ + "rag": RAGState(evidence=CapabilityEvidenceRecord(question=7)).model_dump( + mode="json" + ) + } + ) + history = [ + *_in_flight_history(), + ModelResponse(parts=[ToolCallPart("external_tool", {}, "call-2")]), ] - compacted = _compact_old_tool_returns(messages, frozenset({"rag_search"})) + await agent.run( + "carry on", + deferred_tool_results=DeferredToolResults(calls={"call-2": "external result"}), + message_history=history, + deps=deps, + ) - assert compacted is messages - current_return = compacted[1].parts[0] - assert isinstance(current_return, ToolReturnPart) - assert current_return.content == "current evidence" + assert _record(deps, "rag").question == 7 + + +@pytest.mark.asyncio +async def test_resuming_without_a_stored_identity_fails_instead_of_guessing( + temp_db_path, +): + """Adopting the message count would relabel a question already in progress. + + Every declaration and epoch comparison in it would then be judged against the + wrong question, silently. This state is not one the design produces, so it is + reported rather than repaired. + """ + rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False) + + async def model(_messages, _info): # pragma: no cover - never reached + return ModelResponse(parts=[TextPart("answer")]) + + agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag]) + + with pytest.raises(RuntimeError, match="no stored question identity"): + await agent.run( + "carry on", + deferred_tool_results=DeferredToolResults( + calls={"call-2": "external result"} + ), + message_history=[ + *_in_flight_history(), + ModelResponse(parts=[ToolCallPart("external_tool", {}, "call-2")]), + ], + deps=Deps(), + ) + + +@pytest.mark.asyncio +async def test_citing_after_searching_grounds_the_question(temp_db_path): + """The whole rule, end to end, with no compactor and no policy capability.""" + rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False) + calls = iter( + [ + [ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")], + [ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-2")], + [TextPart("answer")], + ] + ) + + async def model(_messages, _info): + return ModelResponse(parts=next(calls)) + + agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag]) + deps = Deps() + + with patch.object(RAGCapability, "_search", _stub_search): + await agent.run("what does the supervisor do?", deps=deps) + + record = _record(deps, "rag") + question = record.question + assert question is not None + assert record.declaration is not None + assert [ref.chunk_id for ref in record.declaration.refs] == ["chunk-1"] + assert record.occurrences["chunk-1"].retrieved_in_questions == [question] + assert citation_status([record], question=question) == "grounded" + + +@pytest.mark.asyncio +async def test_searching_after_citing_leaves_the_question_uncited(temp_db_path): + rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False) + calls = iter( + [ + [ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")], + [ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-2")], + [ToolCallPart("rag_search", {"query": "again"}, "call-3")], + [TextPart("answer")], + ] + ) + + async def model(_messages, _info): + return ModelResponse(parts=next(calls)) + + agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag]) + deps = Deps() + + with patch.object(RAGCapability, "_search", _stub_search): + await agent.run("what does the supervisor do?", deps=deps) + + record = _record(deps, "rag") + question = record.question + assert question is not None + assert record.declaration is not None + assert citation_status([record], question=question) == "missing" + + +@pytest.mark.asyncio +async def test_a_citation_in_the_same_request_as_its_search_is_not_current( + temp_db_path, +): + """Two calls in one response share an epoch, and citing must follow seeing.""" + rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False) + calls = iter( + [ + [ + ToolCallPart("rag_search", {"query": "supervisor"}, "call-1"), + ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-2"), + ], + [TextPart("answer")], + ] + ) + + async def model(_messages, _info): + return ModelResponse(parts=next(calls)) + + agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag]) + deps = Deps() + + with patch.object(RAGCapability, "_search", _stub_search): + await agent.run("what does the supervisor do?", deps=deps) + + record = _record(deps, "rag") + question = record.question + assert question is not None + assert record.declaration is not None + assert record.declaration.epoch == record.latest_evidence_epoch + assert citation_status([record], question=question) == "missing" + + +@pytest.mark.asyncio +async def test_evidence_cited_in_two_questions_keeps_both_in_the_record(temp_db_path): + """Occurrences outlive the question that wrote them, through the state dict.""" + rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False) + calls = iter( + [ + [ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")], + [ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-2")], + [TextPart("first answer")], + [ToolCallPart("rag_search", {"query": "supervisor again"}, "call-3")], + [ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-4")], + [TextPart("second answer")], + ] + ) + + async def model(_messages, _info): + return ModelResponse(parts=next(calls)) + + agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag]) + deps = Deps() + + with patch.object(RAGCapability, "_search", _stub_search): + first = await agent.run("who supervises?", deps=deps) + first_question = _record(deps, "rag").question + await agent.run( + "and who supervises them?", deps=deps, message_history=first.all_messages() + ) + + record = _record(deps, "rag") + assert record.occurrences["chunk-1"].cited_in_questions == [ + first_question, + record.question, + ] + assert record.question != first_question + + +@pytest.mark.asyncio +async def test_a_run_with_no_prompt_and_no_history_starts_a_question(temp_db_path): + """An instructions-only run is a first question, not a resumption. + + There is no question in progress to keep an identity for, so nothing is + missing and the run proceeds with a fresh one. + """ + rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False) + + async def model(_messages, _info): + return ModelResponse(parts=[TextPart("answer")]) + + agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag]) + deps = Deps() + + await agent.run(deps=deps) + + assert _record(deps, "rag").question == 0 + + +@pytest.mark.asyncio +async def test_citing_without_searching_grounds_the_question(temp_db_path): + """A direct chunk-id citation stands on its own, with no evidence outcome. + + Epochs count messages and so start above zero, which is what lets a + declaration made in the first request still beat an empty evidence horizon. + """ + rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False) + calls = iter( + [ + [ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-1")], + [TextPart("answer")], + ] + ) + + async def model(_messages, _info): + return ModelResponse(parts=next(calls)) + + client = AsyncMock() + client.get_chunk_by_id.return_value = Chunk( + id="chunk-1", document_id="doc-1", content="evidence" + ) + client.get_document_by_id.return_value = None + agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag]) + deps = Deps() + + with patch.object(RAGCapability, "_ensure_rag", AsyncMock(return_value=client)): + await agent.run("cite chunk-1", deps=deps) + + record = _record(deps, "rag") + question = record.question + assert question is not None + assert record.latest_evidence_epoch == 0 + assert record.declaration is not None + assert record.declaration.epoch > 0 + assert record.occurrences["chunk-1"].retrieved_in_questions == [] + assert citation_status([record], question=question) == "grounded" + + +@pytest.mark.asyncio +async def test_a_host_seeded_record_does_not_pass_for_a_resumption(temp_db_path): + """A default record is truthy, so its presence cannot stand in for identity. + + Seeding one is what a host does when it has no state to send, and taking it + at face value would silently answer as question zero. + """ + rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False) + + async def model(_messages, _info): # pragma: no cover - never reached + return ModelResponse(parts=[TextPart("answer")]) + + agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag]) + + with pytest.raises(RuntimeError, match="no stored question identity"): + await agent.run( + "carry on", + message_history=_in_flight_history(), + deps=Deps(state={"rag": RAGState().model_dump(mode="json")}), + ) + + +@pytest.mark.asyncio +async def test_a_resumption_keeps_the_evidence_the_question_already_gathered( + temp_db_path, +): + """Clearing it would lose the results the model is still answering from. + + A citation after the resumption then records no provenance, and cannot resolve + against the expanded search result the model actually saw. + """ + rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False) + calls = iter( + [ + [ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")], + [TextPart("partial answer")], + [ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-3")], + [TextPart("answer")], + ] + ) + + async def model(_messages, _info): + return ModelResponse(parts=next(calls)) + + agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag]) + deps = Deps() + + with patch.object(RAGCapability, "_search", _stub_search): + interrupted = await agent.run("what does the supervisor do?", deps=deps) + identity = _record(deps, "rag").question + assert identity is not None + await agent.run( + deferred_tool_results=DeferredToolResults( + calls={"call-2": "external result"} + ), + message_history=[ + *interrupted.all_messages(), + ModelResponse(parts=[ToolCallPart("external_tool", {}, "call-2")]), + ], + deps=deps, + ) + + record = _record(deps, "rag") + 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_capsule.py b/tests/capabilities/test_evidence_capsule.py new file mode 100644 index 00000000..6f1a51b5 --- /dev/null +++ b/tests/capabilities/test_evidence_capsule.py @@ -0,0 +1,386 @@ +from dataclasses import dataclass, field, replace +from typing import Any +from unittest.mock import patch + +import pytest +from pydantic_ai import Agent +from pydantic_ai.exceptions import UserError +from pydantic_ai.messages import ModelResponse, TextPart +from pydantic_ai.models.function import FunctionModel + +from haiku.rag.capabilities.analysis import create_capability as create_analysis +from haiku.rag.capabilities.compaction import ( + CAPSULE_HEADER, + DiscoveredEvidence, + EvidenceCompactionCapability, + build_capsule, + group_label, +) +from haiku.rag.capabilities.compaction import create_capability as create_compaction +from haiku.rag.capabilities.ledger import ( + CapabilityEvidenceRecord, + EvidenceOccurrence, +) +from haiku.rag.capabilities.rag import create_capability as create_rag +from haiku.rag.config.models import AppConfig +from haiku.rag.store.models.citation import Citation + + +@dataclass +class Deps: + state: dict[str, Any] = field(default_factory=dict) + + +def citation(chunk_id: str, content: str = "evidence body", pictures=()) -> Citation: + return Citation( + document_id=f"doc-of-{chunk_id}", + chunk_id=chunk_id, + document_uri=f"test://{chunk_id}", + document_title=f"Title {chunk_id}", + content=content, + picture_refs=list(pictures), + ) + + +def replace_citation(cited: Citation, **changes: Any) -> Citation: + return cited.model_copy(update=changes) + + +def discovered( + capability: str = "rag", + *, + cited: dict[str, list[int]] | None = None, + contents: dict[str, str] | None = None, + pictures: dict[str, list[str]] | None = None, +) -> DiscoveredEvidence: + """One capability's records, as the compactor would find them.""" + cited = cited or {} + contents = contents or {} + pictures = pictures or {} + record = CapabilityEvidenceRecord(question=max(max(cited.values(), default=[0]))) + for chunk_id, questions in cited.items(): + record.occurrences[chunk_id] = EvidenceOccurrence( + capability=capability, + chunk_id=chunk_id, + retrieved_in_questions=list(questions), + cited_in_questions=list(questions), + ) + return DiscoveredEvidence( + capability=capability, + record=record, + citations={ + chunk_id: citation( + chunk_id, + contents.get(chunk_id, "evidence body"), + pictures.get(chunk_id, ()), + ) + for chunk_id in cited + }, + tool_names=frozenset({f"{capability}_search"}), + ) + + +def test_nothing_cited_produces_no_capsule(): + capsule = build_capsule([discovered()]) + + assert capsule.text == "" + assert capsule.pictures == () + + +def test_cited_evidence_is_grouped_newest_question_first(): + capsule = build_capsule([discovered(cited={"old": [2], "new": [8]})]) + + assert capsule.text.index(group_label(1)) < capsule.text.index(group_label(2)) + assert capsule.text.index("[new]") < capsule.text.index("[old]") + assert CAPSULE_HEADER in capsule.text + + +def test_an_entry_is_rendered_once_in_its_most_recent_citing_group(): + capsule = build_capsule([discovered(cited={"reused": [2, 8], "only-old": [2]})]) + + assert capsule.text.count("[reused]") == 1 + assert capsule.text.index("[reused]") < capsule.text.index("[only-old]") + + +def test_evidence_cited_in_one_question_forms_one_group(): + capsule = build_capsule([discovered(cited={"a": [4], "b": [4]})]) + + assert group_label(1) in capsule.text + assert group_label(2) not in capsule.text + + +def test_every_cited_entry_is_kept_whole(): + """No budget: a long citation is retained in full rather than truncated.""" + body = "L" * 20_000 + capsule = build_capsule([discovered(cited={"long": [4]}, contents={"long": body})]) + + assert body in capsule.text + + +def test_both_capabilities_share_one_capsule(): + capsule = build_capsule( + [ + discovered("rag", cited={"from-rag": [4]}), + discovered("analysis", cited={"from-analysis": [6]}), + ] + ) + + assert capsule.text.count(CAPSULE_HEADER) == 1 + assert "[from-rag]" in capsule.text + assert "[from-analysis]" in capsule.text + + +def test_the_same_chunk_id_under_two_capabilities_is_kept_apart(): + capsule = build_capsule( + [ + discovered("rag", cited={"shared": [4]}, contents={"shared": "rag body"}), + discovered( + "analysis", cited={"shared": [4]}, contents={"shared": "analysis body"} + ), + ] + ) + + assert "rag body" in capsule.text + assert "analysis body" in capsule.text + + +def test_cited_evidence_with_no_canonical_citation_is_an_error(): + """Both are written by the same call, so divergence is not a valid state. + + Rendering the rest would quietly drop evidence an answer rested on, against + the one guarantee this capsule makes. + """ + evidence = discovered(cited={"present": [4]}) + evidence.record.occurrences["absent"] = EvidenceOccurrence( + capability="rag", chunk_id="absent", cited_in_questions=[4] + ) + + with pytest.raises(ValueError, match="absent"): + build_capsule([evidence]) + + +def test_retrieved_but_uncited_evidence_is_not_kept(): + evidence = discovered(cited={"cited": [4]}) + evidence.record.occurrences["seen-only"] = EvidenceOccurrence( + capability="rag", chunk_id="seen-only", retrieved_in_questions=[4] + ) + + capsule = build_capsule([evidence]) + + assert "[cited]" in capsule.text + assert "seen-only" not in capsule.text + + +def test_a_source_is_named_once_when_the_title_is_the_uri(): + """Real corpora set both to the document id, which reads as a stutter.""" + evidence = discovered(cited={"a": [4]}) + evidence.citations["a"].document_title = "2410.11843v5" + evidence.citations["a"].document_uri = "2410.11843v5" + + capsule = build_capsule([evidence]) + + assert 'Source: "2410.11843v5"' in capsule.text + assert "2410.11843v5)" not in capsule.text + + +def test_pictures_of_cited_evidence_are_all_retained_newest_first(): + capsule = build_capsule( + [ + discovered( + cited={"a": [2], "c": [6]}, + pictures={"a": ["#/pictures/0"], "c": ["#/pictures/1", "#/pictures/2"]}, + ) + ] + ) + + assert [picture.self_ref for picture in capsule.pictures] == [ + "#/pictures/1", + "#/pictures/2", + "#/pictures/0", + ] + assert capsule.pictures[0].document_id == "doc-of-c" + assert capsule.pictures[0].capability == "rag" + + +def test_a_picture_of_uncited_evidence_is_not_retained(): + found = discovered(cited={"cited": [4]}, pictures={"cited": ["#/pictures/0"]}) + found.record.occurrences["seen-only"] = EvidenceOccurrence( + capability="rag", chunk_id="seen-only", retrieved_in_questions=[4] + ) + evidence = replace( + found, + citations={ + **found.citations, + "seen-only": citation("seen-only", pictures=["#/pictures/9"]), + }, + ) + + capsule = build_capsule([evidence]) + + assert [picture.self_ref for picture in capsule.pictures] == ["#/pictures/0"] + + +def test_one_picture_cited_through_two_chunks_is_attached_once(): + """Overlapping chunks of one document share a figure, counted twice by a provider.""" + found = discovered( + cited={"first": [4], "second": [4]}, + pictures={"first": ["#/pictures/1"], "second": ["#/pictures/1"]}, + ) + shared = { + chunk_id: replace_citation(cited, document_id="doc-shared") + for chunk_id, cited in found.citations.items() + } + + capsule = build_capsule([replace(found, citations=shared)]) + + assert len(capsule.pictures) == 1 + assert capsule.pictures[0].chunk_id == "first" + + +def test_the_same_reference_in_two_documents_is_kept_twice(): + """``#/pictures/1`` means a different figure in a different document. + + One capability throughout, so only the document differs: dropping the document + from the identity would have to fail this. + """ + capsule = build_capsule( + [ + discovered( + "rag", + cited={"a": [4], "b": [4]}, + pictures={"a": ["#/pictures/1"], "b": ["#/pictures/1"]}, + ) + ] + ) + + assert len(capsule.pictures) == 2 + assert {picture.capability for picture in capsule.pictures} == {"rag"} + assert {picture.document_id for picture in capsule.pictures} == { + "doc-of-a", + "doc-of-b", + } + + +def test_a_picture_label_names_the_chunk_it_belongs_to(): + capsule = build_capsule( + [discovered(cited={"a": [4]}, pictures={"a": ["#/pictures/0"]})] + ) + + label = capsule.pictures[0].label + assert "[a]" in label + assert "#/pictures/0" in label + assert "knowledge base" in label + assert "Not provided by the user" in label + + +def _spy_discovery(found: list[list[DiscoveredEvidence]]): + """Discover from ``before_run``, the earliest point the registry is reliable.""" + original = EvidenceCompactionCapability.before_run + + async def spy(self, ctx): + await original(self, ctx) + found.append(self.discover(ctx)) + + return patch.object(EvidenceCompactionCapability, "before_run", spy) + + +async def _answer(_messages, _info): + return ModelResponse(parts=[TextPart("answer")]) + + +@pytest.mark.asyncio +async def test_the_compactor_discovers_both_evidence_capabilities(temp_db_path): + """Discovery runs one way through the registry, so nothing needs wiring.""" + compactor = create_compaction() + rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False) + analysis = create_analysis( + db_path=temp_db_path, config=AppConfig(), defer_loading=False + ) + found: list[list[DiscoveredEvidence]] = [] + + with _spy_discovery(found): + agent = Agent( + FunctionModel(_answer), + deps_type=Deps, + capabilities=[rag, analysis, compactor], + ) + await agent.run("a question", deps=Deps()) + + assert {evidence.capability: set(evidence.tool_names) for evidence in found[0]} == { + "rag": {"rag_search"}, + "analysis": {"analysis_search", "analysis_execute_code"}, + } + + +@pytest.mark.asyncio +async def test_discovery_sees_the_run_instances_not_the_registered_ones(temp_db_path): + """A registered capability holds no state; only its per-run copy does.""" + compactor = create_compaction() + rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False) + found: list[list[DiscoveredEvidence]] = [] + + with _spy_discovery(found): + agent = Agent( + FunctionModel(_answer), deps_type=Deps, capabilities=[rag, compactor] + ) + await agent.run("a question", deps=Deps()) + + assert rag.state is None + assert found[0][0].record.question == 0 + + +def test_two_compactors_fail_fast(temp_db_path): + """Each would rewrite the same history and each would build its own capsule. + + They share this capability's id, so pydantic-ai refuses at construction and + nothing here has to police it. + """ + rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False) + + with pytest.raises(UserError, match="unique within a run"): + Agent( + FunctionModel(_answer), + deps_type=Deps, + capabilities=[rag, create_compaction(), create_compaction()], + ) + + +@pytest.mark.asyncio +async def test_a_compactor_alone_discovers_nothing_and_still_runs(): + found: list[list[DiscoveredEvidence]] = [] + + with _spy_discovery(found): + agent = Agent( + FunctionModel(_answer), deps_type=Deps, capabilities=[create_compaction()] + ) + result = await agent.run("a question", deps=Deps()) + + assert found == [[]] + assert result.output == "answer" + + +@pytest.mark.asyncio +async def test_a_deferred_capability_the_model_never_loaded_has_an_empty_record( + temp_db_path, +): + """It is still discovered, because every registered capability gets a run copy. + + Nothing was retrieved under it, so its record contributes no entries and the + compactor needs no special case for it. + """ + deferred = create_rag(db_path=temp_db_path, config=AppConfig()) + found: list[list[DiscoveredEvidence]] = [] + + with _spy_discovery(found): + agent = Agent( + FunctionModel(_answer), + deps_type=Deps, + capabilities=[deferred, create_compaction()], + ) + await agent.run("a question", deps=Deps()) + + assert deferred.defer_loading is True + assert deferred.state is None + assert [evidence.capability for evidence in found[0]] == ["rag"] + assert found[0][0].record.occurrences == {} + assert build_capsule(found[0]).text == "" diff --git a/tests/capabilities/test_evidence_ledger.py b/tests/capabilities/test_evidence_ledger.py new file mode 100644 index 00000000..c9f94353 --- /dev/null +++ b/tests/capabilities/test_evidence_ledger.py @@ -0,0 +1,242 @@ +import pytest + +from haiku.rag.capabilities.ledger import ( + CapabilityEvidenceRecord, + CitationDeclaration, + EvidenceRef, + citation_status, +) + + +def rag_ref(chunk_id: str = "c1") -> EvidenceRef: + return EvidenceRef(capability="rag", chunk_id=chunk_id) + + +def test_a_record_survives_the_state_round_trip(): + """Capability state is persisted as JSON, so the schema must survive it. + + A dict keyed by ``(capability, chunk_id)`` does not: the key serialises to + ``"rag,c1"`` and fails revalidation as a tuple. + """ + record = CapabilityEvidenceRecord(question=4) + record.note_evidence(5) + record.declare([rag_ref()], epoch=7, retrieved_now={"c1"}) + + restored = CapabilityEvidenceRecord.model_validate(record.model_dump(mode="json")) + + assert restored == record + assert citation_status([restored], question=4) == "grounded" + assert restored.occurrences["c1"].cited_in_questions == [4] + assert restored.occurrences["c1"].retrieved_in_questions == [4] + + +def test_no_declaration_reads_as_missing(): + assert citation_status([CapabilityEvidenceRecord()], question=0) == "missing" + assert citation_status([], question=0) == "missing" + + +def test_refs_make_it_grounded_and_no_refs_make_it_ungrounded(): + grounded = CapabilityEvidenceRecord(question=0) + grounded.declare([rag_ref()], epoch=1) + + ungrounded = CapabilityEvidenceRecord(question=0) + ungrounded.declare([], epoch=1) + + assert citation_status([grounded], question=0) == "grounded" + assert citation_status([ungrounded], question=0) == "ungrounded" + + +def test_an_earlier_questions_declaration_is_never_current(): + """A later question inherits nothing: it has declared nothing yet.""" + record = CapabilityEvidenceRecord(question=2) + record.declare([rag_ref()], epoch=3) + assert citation_status([record], question=2) == "grounded" + + record.begin_question(8) + + assert citation_status([record], question=8) == "missing" + + +def test_a_citation_in_the_same_request_as_the_evidence_is_not_current(): + """Citing must follow seeing: equal epochs mean one request.""" + record = CapabilityEvidenceRecord(question=0) + record.note_evidence(5) + record.declare([rag_ref()], epoch=5) + + assert citation_status([record], question=0) == "missing" + + record.declare([rag_ref()], epoch=7) + + assert citation_status([record], question=0) == "grounded" + + +def test_evidence_from_another_capability_after_citing_makes_it_uncited(): + """Currency spans capabilities, which only works because epochs are global.""" + cited = CapabilityEvidenceRecord(question=0) + cited.note_evidence(3) + cited.declare([rag_ref()], epoch=5) + searched_after = CapabilityEvidenceRecord(question=0) + searched_after.note_evidence(7) + + assert citation_status([cited], question=0) == "grounded" + assert citation_status([cited, searched_after], question=0) == "missing" + + +def test_declarations_at_the_same_epoch_merge(): + record = CapabilityEvidenceRecord(question=0) + record.declare([rag_ref("c1")], epoch=3) + record.declare([rag_ref("c2")], epoch=3) + + assert record.declaration is not None + assert [ref.chunk_id for ref in record.declaration.refs] == ["c1", "c2"] + + +def test_repeating_a_ref_at_the_same_epoch_does_not_duplicate_it(): + record = CapabilityEvidenceRecord(question=0) + record.declare([rag_ref()], epoch=3) + record.declare([rag_ref()], epoch=3) + + assert record.declaration is not None + assert len(record.declaration.refs) == 1 + + +def test_neither_cite_order_downgrades_a_grounded_declaration(): + grounded_then_empty = CapabilityEvidenceRecord(question=0) + grounded_then_empty.declare([rag_ref()], epoch=3) + grounded_then_empty.declare([], epoch=3) + + empty_then_grounded = CapabilityEvidenceRecord(question=0) + empty_then_grounded.declare([], epoch=3) + empty_then_grounded.declare([rag_ref()], epoch=3) + + assert citation_status([grounded_then_empty], question=0) == "grounded" + assert citation_status([empty_then_grounded], question=0) == "grounded" + + +def test_the_same_chunk_id_under_two_capabilities_stays_separate(): + rag = CapabilityEvidenceRecord(question=0) + rag.declare([EvidenceRef(capability="rag", chunk_id="shared")], epoch=3) + analysis = CapabilityEvidenceRecord(question=0) + analysis.declare([EvidenceRef(capability="analysis", chunk_id="shared")], epoch=3) + + assert rag.occurrences["shared"].capability == "rag" + assert analysis.occurrences["shared"].capability == "analysis" + + +def test_citing_the_same_chunk_in_two_questions_records_both(): + record = CapabilityEvidenceRecord(question=2) + record.declare([rag_ref()], epoch=3, retrieved_now={"c1"}) + record.begin_question(8) + record.declare([rag_ref()], epoch=9) + + occurrence = record.occurrences["c1"] + assert occurrence.cited_in_questions == [2, 8] + assert occurrence.retrieved_in_questions == [2] + + +def test_a_declaration_records_the_question_and_epoch_it_was_made_at(): + record = CapabilityEvidenceRecord(question=6) + record.declare([rag_ref()], epoch=11) + + assert record.declaration == CitationDeclaration( + question=6, epoch=11, refs=[rag_ref()] + ) + + +def test_a_fresh_record_has_no_question_identity(): + """The identity is established by the run, and its absence must be detectable. + + A default record is truthy, so its mere presence cannot stand in for having + been through ``for_run``: a host that seeds one would otherwise pass the + resumption check with a fabricated identity of zero. + """ + assert CapabilityEvidenceRecord().question is None + + +def test_evidence_cannot_move_backwards_in_the_conversation(): + """Epochs are message counts, and currency depends on them only growing. + + Silently keeping the newer value would leave every later declaration stale + for the rest of the conversation, permanently and invisibly. + """ + record = CapabilityEvidenceRecord(question=0) + record.note_evidence(9) + + with pytest.raises(ValueError, match="append-only"): + record.note_evidence(4) + + +def test_citing_before_a_run_establishes_the_question_is_refused(): + with pytest.raises(ValueError, match="question identity"): + CapabilityEvidenceRecord().declare([rag_ref()], epoch=3) + + +def test_a_declaration_cannot_move_backwards(): + """Otherwise a stale citation replaces a newer one and revives the answer.""" + record = CapabilityEvidenceRecord(question=0) + record.declare([rag_ref("newer")], epoch=5) + + with pytest.raises(ValueError, match="append-only"): + record.declare([rag_ref("older")], epoch=3) + + assert record.declaration is not None + assert [ref.chunk_id for ref in record.declaration.refs] == ["newer"] + + +def test_evidence_cannot_predate_a_recorded_declaration(): + """The declaration's epoch is a recorded message count like any other.""" + record = CapabilityEvidenceRecord(question=0) + record.declare([rag_ref()], epoch=5) + + with pytest.raises(ValueError, match="append-only"): + record.note_evidence(3) + + +def test_a_question_starts_behind_the_epochs_of_the_one_before_it(): + """A question's identity is the history it arrives on, not a continuation. + + Epochs are compared only within the question that recorded them, so a host + whose stored history shifted between two questions is answered rather than + refused. + """ + record = CapabilityEvidenceRecord(question=0) + record.note_evidence(9) + + record.begin_question(4) + + assert record.question == 4 + + +def test_a_question_cannot_reuse_the_identity_of_the_one_before_it(): + """Occurrences outlive their question and are ordered by identity. + + Two questions sharing one identity merge into a single capsule group, and a + lower one is rendered as though its evidence were cited earlier. + """ + record = CapabilityEvidenceRecord(question=4) + + with pytest.raises(ValueError, match="already answered"): + record.begin_question(4) + + with pytest.raises(ValueError, match="already answered"): + record.begin_question(3) + + assert record.question == 4 + + +def test_a_question_starts_clear_of_the_one_before_it(): + """Evidence and declarations describe a single question and end with it. + + Carrying either into the next question makes it answerable by the last + question's citations, and freezes its own declarations behind an epoch no + message in it can reach. + """ + record = CapabilityEvidenceRecord(question=4) + record.note_evidence(6) + record.declare([rag_ref()], epoch=7) + + record.begin_question(9) + + assert record.latest_evidence_epoch == 0 + assert record.declaration is None + assert citation_status([record], question=9) == "missing" diff --git a/tests/capabilities/test_evidence_wire.py b/tests/capabilities/test_evidence_wire.py new file mode 100644 index 00000000..3915fe17 --- /dev/null +++ b/tests/capabilities/test_evidence_wire.py @@ -0,0 +1,729 @@ +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, DeferredToolResults, 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.parametrize( + "resume_kwargs", + [ + pytest.param({}, id="no prompt"), + pytest.param( + {"deferred_tool_results": DeferredToolResults()}, id="deferred results" + ), + ], +) +@pytest.mark.asyncio +async def test_a_resumed_question_keeps_the_evidence_it_is_answering_from( + temp_db_path, resume_kwargs +): + """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, **resume_kwargs + ) + + 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 + + +def test_two_owned_returns_in_one_request_yield_one_capsule(): + """A model can search twice in one response, so a request can hold two returns. + + Identifying the carrier by message alone gave every return in it the capsule, + which duplicates the whole thing — and it is unbounded. + """ + messages = [ + ModelRequest(parts=[UserPromptPart("first")]), + ModelResponse( + parts=[ + ToolCallPart("rag_search", {"query": "a"}, "call-1"), + ToolCallPart("rag_search", {"query": "b"}, "call-2"), + ] + ), + ModelRequest( + parts=[ + ToolReturnPart("rag_search", "FIRST EVIDENCE", "call-1"), + ToolReturnPart("rag_search", "SECOND EVIDENCE", "call-2"), + ] + ), + 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) == [RECEIPT, "CAPSULE"] + + +def test_the_capsule_is_attached_beside_the_newest_return_of_that_request(): + messages = [ + ModelRequest(parts=[UserPromptPart("first")]), + ModelResponse( + parts=[ + ToolCallPart("rag_search", {"query": "a"}, "call-1"), + ToolCallPart("rag_search", {"query": "b"}, "call-2"), + ] + ), + ModelRequest( + parts=[ + ToolReturnPart("rag_search", "FIRST", "call-1"), + ToolReturnPart("rag_search", "SECOND", "call-2"), + ] + ), + ModelResponse(parts=[TextPart("an answer")]), + 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", "#/pictures/1"), fresh], + ) + + assert images_of(compacted) == [fresh] + assert returns_of(compacted) == [RECEIPT, "CAPSULE"] diff --git a/tests/chat/test_chat_app.py b/tests/chat/test_chat_app.py index 40057501..b1fdd345 100644 --- a/tests/chat/test_chat_app.py +++ b/tests/chat/test_chat_app.py @@ -1,3 +1,4 @@ +from copy import deepcopy from pathlib import Path from unittest.mock import AsyncMock, patch @@ -412,3 +413,55 @@ async def test_chat_app_open_failure_surfaces_real_error(tmp_path: Path): with pytest.raises(FileNotFoundError): async with app.run_test(): pass + + +@pytest.mark.asyncio +async def test_a_cancelled_run_does_not_advance_persisted_state(temp_db_path: Path): + """State and message history have to move together, or the thread bricks. + + A cancelled run keeps whatever the tools wrote but discards the run's messages. + If the state advanced, the next question derives its identity from the shorter + history, lands behind the recorded evidence epoch, and is refused as + non-append-only — leaving the conversation unusable until cleared. + """ + import asyncio + + app, mock_client = _make_app(temp_db_path) + + class CancellingRun: + """A run that writes evidence through the tools, then is cancelled.""" + + def __init__(self, deps): + self._deps = deps + + async def __aenter__(self): + self._deps.state["rag"] = { + "evidence": {"question": 0, "latest_evidence_epoch": 7} + } + return self + + async def __aexit__(self, *_): + return False + + def __aiter__(self): + return self + + async def __anext__(self): + raise asyncio.CancelledError + + with patch("haiku.rag.chat.app.HaikuRAG", return_value=mock_client): + async with app.run_test(): + app._state = { + "rag": {"evidence": {"question": 0, "latest_evidence_epoch": 0}} + } + before = deepcopy(app._state) + + class Agent: + def run_stream_events(self, *_, deps, **__): + return CancellingRun(deps) + + app._agent = Agent() # type: ignore[assignment] + await app._run_agent("a question") + + assert app._state == before + assert app._messages == [] diff --git a/tests/test_picture_in_context.py b/tests/test_picture_in_context.py index 55f2f368..a9b32a47 100644 --- a/tests/test_picture_in_context.py +++ b/tests/test_picture_in_context.py @@ -402,8 +402,9 @@ async def test_search_tool_returns_multimodal_when_picture_present(): assert isinstance(result.return_value, str) assert "Type: picture" in result.return_value or "rank 1" in result.return_value assert result.content is not None - assert len(result.content) == 1 - part = result.content[0] + images = [c for c in result.content if isinstance(c, BinaryContent)] + assert len(images) == 1 + part = images[0] assert isinstance(part, BinaryContent) assert part.media_type == "image/png" assert part.identifier == "#/pictures/0" @@ -457,11 +458,12 @@ async def test_search_tool_attaches_same_self_ref_from_different_documents(): assert isinstance(result, ToolReturn) assert result.content is not None - assert len(result.content) == 2, ( + images = [c for c in result.content if isinstance(c, BinaryContent)] + assert len(images) == 2, ( "Both documents' figures must reach the model — dedup keyed on " "self_ref alone would drop doc-B's bytes." ) - payloads = {part.data for part in result.content} # type: ignore[attr-defined] + payloads = {part.data for part in images} assert PICTURE_BYTES in payloads assert other_bytes in payloads @@ -905,7 +907,7 @@ async def test_search_tool_drops_invalid_image_bytes(): assert isinstance(result, ToolReturn) assert result.content is not None - identifiers = {p.identifier for p in result.content} # type: ignore[attr-defined] + identifiers = {p.identifier for p in result.content if isinstance(p, BinaryContent)} assert identifiers == {"#/pictures/0"}, ( "Only the decodable PNG should reach the model — the corrupt " "ref must be dropped so we don't emit a placeholder for an " diff --git a/tests/tools/test_search.py b/tests/tools/test_search.py index 60bf3404..8f09ad04 100644 --- a/tests/tools/test_search.py +++ b/tests/tools/test_search.py @@ -210,30 +210,35 @@ def search_config(): return Config -class TestBuildBinaryPartsFromResults: - """Picture bytes are attached once per (document, self_ref) pair.""" +def _png_b64(): + import base64 + from io import BytesIO + + from PIL import Image as PILImage + + buf = BytesIO() + PILImage.new("RGB", (4, 4), "red").save(buf, format="PNG") + return base64.b64encode(buf.getvalue()).decode() + + +class TestBuildImageContentFromResults: + """Picture bytes are attached once per (document, self_ref) pair, and labelled.""" def test_results_without_image_data_contribute_nothing(self): - from haiku.rag.tools.search import build_binary_parts_from_results + from haiku.rag.tools.search import build_image_content_from_results results = [ SearchResult(content="text only", score=0.5, chunk_id="c1", image_data=None) ] - assert build_binary_parts_from_results(results) == [] + assert build_image_content_from_results(results) == [] def test_duplicate_document_and_ref_is_attached_once(self): - import base64 - from io import BytesIO + from pydantic_ai.messages import BinaryContent - from PIL import Image as PILImage + from haiku.rag.tools.search import build_image_content_from_results - from haiku.rag.tools.search import build_binary_parts_from_results - - buf = BytesIO() - PILImage.new("RGB", (4, 4), "red").save(buf, format="PNG") - png = base64.b64encode(buf.getvalue()).decode() - shared = {"#/pictures/0": png} + shared = {"#/pictures/0": _png_b64()} results = [ SearchResult( content="a", @@ -251,6 +256,53 @@ class TestBuildBinaryPartsFromResults: ), ] - parts = build_binary_parts_from_results(results) + content = build_image_content_from_results(results) - assert len(parts) == 1 + images = [item for item in content if isinstance(item, BinaryContent)] + assert len(images) == 1 + + def test_each_image_is_labelled_with_the_result_it_belongs_to(self): + """Label every picture, not just the batch. + + ``ToolReturn.content`` reaches the model as a user-role message, and one + leading note does not override that: with a single note on the wire, + gemma4-26b still reasoned "the user also provided images in the prompt". + A label adjacent to each picture also names the chunk to cite for it, + which ``BinaryContent.identifier`` cannot do — it does not survive + serialization to the vision API. + """ + from pydantic_ai.messages import BinaryContent + + from haiku.rag.tools.search import build_image_content_from_results + + results = [ + SearchResult( + content="a", + score=0.9, + chunk_id="c1", + document_id="doc-1", + image_data={"#/pictures/0": _png_b64()}, + ), + SearchResult( + content="b", + score=0.8, + chunk_id="c2", + document_id="doc-2", + image_data={"#/pictures/3": _png_b64()}, + ), + ] + + content = build_image_content_from_results(results) + + # label, image, label, image — each picture preceded by its own line. + assert [type(item) is str for item in content] == [True, False, True, False] + assert isinstance(content[1], BinaryContent) + assert isinstance(content[3], BinaryContent) + + first, second = content[0], content[2] + assert isinstance(first, str) and isinstance(second, str) + assert "c1" in first and "#/pictures/0" in first + assert "c2" in second and "#/pictures/3" in second + assert "1 of 2" in first and "2 of 2" in second + for label in (first, second): + assert "not provided by the user" in label.lower()