From 5d76461b9f0b9519df0fa933364be54144bc3657 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 7 Aug 2026 12:40:05 +0300 Subject: [PATCH 01/13] Trim earlier-question evidence off the wire, keeping the images _compact_old_tool_returns ran in before_model_request, whose result core assigns back onto ctx.state.message_history, so the trim reached all_messages() and every host that persists a thread. It now runs in wrap_model_request, which operates on a detached list: the model sees the trimmed history, the host keeps what it retrieved. Content attached to a ToolReturn arrives as its own UserPromptPart in the same ModelRequest as the ToolReturnPart, so the turn-boundary scan read an image-bearing search result as a new user turn and discarded evidence retrieved earlier in the same turn. _is_user_turn() now requires a request with no tool returns. Page images on a replaced return stay. Dropping them with their text bounds context growth, but a follow-up about a figure already shown ("what colour is that box?") carries no terms that could retrieve it again: measured on gemma4-26b against two ORB figures, removing the image turned both answers into "I cannot find enough information", and keeping it answers correctly. Bounding that growth needs to preserve cited figures, which is a separate change. The replacement notice no longer claims citations remain in state; _clear_invocation_state has cleared them by then. --- CHANGELOG.md | 2 + .../haiku/rag/capabilities/_base.py | 63 +++++++--- tests/capabilities/test_capabilities.py | 117 +++++++++++++++++- 3 files changed, 167 insertions(+), 15 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0c5c1819..4a7b7930 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,8 @@ - `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. +- Prior-question tool output is trimmed from the model request only; `all_messages()` retains what the run gathered. +- A search result carrying page images no longer counts as a turn boundary; evidence retrieved earlier in the same turn survives. ## [0.73.0] - 2026-08-06 diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index ff7a2b0a..ae933464 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py @@ -7,7 +7,7 @@ from typing import Any, cast from pydantic import BaseModel from pydantic_ai import ModelRetry, RunContext, ToolFailed -from pydantic_ai.capabilities import AbstractCapability +from pydantic_ai.capabilities import AbstractCapability, WrapModelRequestHandler from pydantic_ai.messages import ( InstructionPart, ModelMessage, @@ -76,20 +76,43 @@ def _clear_invocation_state(state: BaseModel) -> None: value.clear() +PRIOR_TURN_NOTICE = ( + "[Evidence retrieved for an earlier question, no longer shown. It does not " + "count as cited for the current question.]" +) + + +def _is_user_turn(message: ModelMessage) -> bool: + """Whether this message is a user turn rather than tool output. + + Content attached to a ``ToolReturn`` (page images) arrives as its own + ``UserPromptPart`` in the same ``ModelRequest`` as the ``ToolReturnPart``, + so a bare ``UserPromptPart`` check reads tool output as a new turn. + """ + if not isinstance(message, ModelRequest): + return False + return any(isinstance(part, UserPromptPart) for part in message.parts) and not any( + isinstance(part, ToolReturnPart) for part in message.parts + ) + + def _compact_old_tool_returns( messages: list[ModelMessage], tool_names: frozenset[str] ) -> list[ModelMessage]: - """Remove bulky prior-turn evidence while retaining the current turn. + """Remove bulky earlier-question evidence while retaining the current 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. + replaced. + + Page images attached to a replaced return are deliberately left in place. + Dropping them alongside their text is tempting — they are the bulk, and + they accumulate — but a follow-up about a figure already shown ("what + colour is that box?") carries no terms that could retrieve it again, so + removing the image turns an answerable question into a refusal. """ latest_user_message = -1 for index, message in enumerate(messages): - if isinstance(message, ModelRequest) and any( - isinstance(part, UserPromptPart) for part in message.parts - ): + if _is_user_turn(message): latest_user_message = index if latest_user_message < 0: @@ -100,10 +123,7 @@ def _compact_old_tool_returns( if not isinstance(message, ModelRequest): continue parts = [ - replace( - part, - content="[Prior-turn RAG tool output removed; citations remain in state.]", - ) + replace(part, content=PRIOR_TURN_NOTICE) if isinstance(part, ToolReturnPart) and part.tool_name in tool_names else part for part in message.parts @@ -168,12 +188,27 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): return f"{self.config.prompts.domain_preamble}\n\n{self.instruction_text}" return self.instruction_text - async def before_model_request( - self, ctx: RunContext[Any], request_context: ModelRequestContext - ) -> ModelRequestContext: + async def wrap_model_request( + self, + ctx: RunContext[Any], + *, + request_context: ModelRequestContext, + handler: WrapModelRequestHandler, + ) -> ModelResponse: + """Trim earlier-question evidence off the wire only. + + Deliberately not ``before_model_request``: that hook's result is + assigned back onto the run's message history, so trimming there would + destroy the host's record of what was retrieved. + """ request_context.messages = _compact_old_tool_returns( request_context.messages, self.tool_names ) + return await handler(request_context) + + async def before_model_request( + self, ctx: RunContext[Any], request_context: ModelRequestContext + ) -> ModelRequestContext: if instruction := self._budget_notice(): current_request = request_context.messages[-1] if isinstance(current_request, ModelRequest): diff --git a/tests/capabilities/test_capabilities.py b/tests/capabilities/test_capabilities.py index 6bbf6141..f107df51 100644 --- a/tests/capabilities/test_capabilities.py +++ b/tests/capabilities/test_capabilities.py @@ -6,6 +6,7 @@ from unittest.mock import AsyncMock, patch import pytest from pydantic_ai import Agent, ModelRetry, RunContext, ToolFailed from pydantic_ai.messages import ( + BinaryContent, ModelRequest, ModelResponse, TextPart, @@ -19,6 +20,7 @@ from pydantic_ai.usage import RunUsage from haiku.rag.capabilities._base import ( CITATION_GRACE_REQUESTS, + PRIOR_TURN_NOTICE, _called_own_tool, _compact_old_tool_returns, ) @@ -806,7 +808,7 @@ def test_prior_turn_tool_results_are_compacted_but_current_evidence_is_kept(): 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 old_return.content == PRIOR_TURN_NOTICE assert isinstance(current_return, ToolReturnPart) assert current_return.content == "current evidence" @@ -825,3 +827,116 @@ def test_tool_results_are_unchanged_when_history_has_no_user_prompt(): current_return = compacted[1].parts[0] assert isinstance(current_return, ToolReturnPart) assert current_return.content == "current evidence" + + +PAGE_IMAGE = BinaryContent(data=b"\x89PNG" + b"\x00" * 64, media_type="image/png") + + +def _search_exchange(call_id: str, evidence: str, *, images: bool): + """One search round-trip in the shape pydantic-ai produces. + + Images on a ``ToolReturn`` arrive as a separate ``UserPromptPart`` appended + to the same ``ModelRequest`` as the ``ToolReturnPart``. + """ + parts: list[Any] = [ToolReturnPart("rag_search", evidence, call_id)] + if images: + parts.append(UserPromptPart(content=[PAGE_IMAGE])) + return [ + ModelResponse(parts=[ToolCallPart("rag_search", {}, call_id)]), + ModelRequest(parts=parts), + ] + + +def test_image_bearing_tool_return_does_not_end_the_current_turn(): + messages = [ + ModelRequest(parts=[UserPromptPart("current question")]), + *_search_exchange("first-call", "first evidence", images=False), + *_search_exchange("second-call", "second evidence", images=True), + ] + + compacted = _compact_old_tool_returns(messages, frozenset({"rag_search"})) + + first_return = compacted[2].parts[0] + assert isinstance(first_return, ToolReturnPart) + assert first_return.content == "first evidence" + + +def test_prior_turn_images_outlive_their_tool_return(): + """A follow-up about a figure cannot retrieve it again, so keep the image. + + "What colour is that box?" has no terms the search can use, so dropping the + image with its text turns an answerable question into a refusal. + """ + messages = [ + ModelRequest(parts=[UserPromptPart("old question")]), + *_search_exchange("old-call", "old evidence", images=True), + ModelResponse(parts=[TextPart("an answer")]), + ModelRequest(parts=[UserPromptPart("current question")]), + ] + + compacted = _compact_old_tool_returns(messages, frozenset({"rag_search"})) + + old_return = compacted[2].parts[0] + assert isinstance(old_return, ToolReturnPart) + assert old_return.content == PRIOR_TURN_NOTICE + assert any( + isinstance(part, UserPromptPart) and not isinstance(part.content, str) + for message in compacted + for part in message.parts + ) + + +def test_user_attached_image_starts_a_turn_and_is_never_dropped(): + messages = [ + ModelRequest(parts=[UserPromptPart("old question")]), + *_search_exchange("old-call", "old evidence", images=False), + ModelResponse(parts=[TextPart("an answer")]), + ModelRequest(parts=[UserPromptPart(content=[PAGE_IMAGE, "what is this?"])]), + ] + + compacted = _compact_old_tool_returns(messages, frozenset({"rag_search"})) + + old_return = compacted[2].parts[0] + assert isinstance(old_return, ToolReturnPart) + assert old_return.content == PRIOR_TURN_NOTICE + attached = compacted[-1].parts[0] + assert isinstance(attached, UserPromptPart) + assert attached.content == [PAGE_IMAGE, "what is this?"] + + +async def test_compaction_never_reaches_the_stored_message_history(temp_db_path): + """Trimming is for the wire; hosts keep the evidence they gathered.""" + capability = create_rag( + db_path=temp_db_path, config=AppConfig(), defer_loading=False + ) + turns = iter( + [ + ModelResponse(parts=[ToolCallPart("rag_search", {"query": "q"}, "call-1")]), + ModelResponse(parts=[TextPart("first answer")]), + ModelResponse(parts=[TextPart("second answer")]), + ] + ) + + async def model(_messages, _info): + return next(turns) + + agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[capability]) + deps = Deps(state={"rag": RAGState().model_dump(mode="json")}) + + with patch.object( + RAGCapability, "_search", AsyncMock(return_value="REAL EVIDENCE") + ): + first = await agent.run("old question", deps=deps) + second = await agent.run( + "current question", deps=deps, message_history=first.all_messages() + ) + + returns = [ + str(part.content) + for message in second.all_messages() + if isinstance(message, ModelRequest) + for part in message.parts + if isinstance(part, ToolReturnPart) + ] + assert "REAL EVIDENCE" in returns + assert PRIOR_TURN_NOTICE not in returns From d9bd3a701fde1320096b50e4321341f9d2d7670e Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Sat, 8 Aug 2026 14:43:29 +0300 Subject: [PATCH 02/13] Take the compaction turn boundary from the run, not the message shape --- CHANGELOG.md | 3 +- .../haiku/rag/capabilities/_base.py | 42 ++++---- tests/capabilities/test_capabilities.py | 96 +++++++++++++++++-- 3 files changed, 106 insertions(+), 35 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 4a7b7930..f75350e7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,7 +8,8 @@ - `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. - Prior-question tool output is trimmed from the model request only; `all_messages()` retains what the run gathered. -- A search result carrying page images no longer counts as a turn boundary; evidence retrieved earlier in the same turn survives. +- Evidence retrieved for the current question is no longer trimmed mid-question, whether or not it carries page images. +- `rag_cite` / `analysis_cite` returns are no longer replaced by the prior-question notice. ## [0.73.0] - 2026-08-06 diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index ae933464..41994ed8 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py @@ -16,7 +16,6 @@ from pydantic_ai.messages import ( ToolCallPart, ToolReturn, ToolReturnPart, - UserPromptPart, ) from pydantic_ai.models import ModelRequestContext from pydantic_ai.run import AgentRunResult @@ -82,22 +81,11 @@ PRIOR_TURN_NOTICE = ( ) -def _is_user_turn(message: ModelMessage) -> bool: - """Whether this message is a user turn rather than tool output. - - Content attached to a ``ToolReturn`` (page images) arrives as its own - ``UserPromptPart`` in the same ``ModelRequest`` as the ``ToolReturnPart``, - so a bare ``UserPromptPart`` check reads tool output as a new turn. - """ - if not isinstance(message, ModelRequest): - return False - return any(isinstance(part, UserPromptPart) for part in message.parts) and not any( - isinstance(part, ToolReturnPart) for part in message.parts - ) - - def _compact_old_tool_returns( - messages: list[ModelMessage], tool_names: frozenset[str] + messages: list[ModelMessage], + tool_names: frozenset[str], + *, + turn_start: int, ) -> list[ModelMessage]: """Remove bulky earlier-question evidence while retaining the current one. @@ -109,17 +97,19 @@ def _compact_old_tool_returns( they accumulate — but a follow-up about a figure already shown ("what colour is that box?") carries no terms that could retrieve it again, so removing the image turns an answerable question into a refusal. - """ - latest_user_message = -1 - for index, message in enumerate(messages): - if _is_user_turn(message): - latest_user_message = index - if latest_user_message < 0: + ``turn_start`` is how many messages existed when the current question + arrived, so everything below it belongs to an earlier one. The run reports + it rather than this function deriving it from message shape: a + ``UserPromptPart`` mid-question is as likely to be page images on a tool + return, or a notice a capability injected, and reading either as the next + question strips evidence the model is still answering from. + """ + if turn_start <= 0: return messages compacted = list(messages) - for index, message in enumerate(messages[:latest_user_message]): + for index, message in enumerate(messages[:turn_start]): if not isinstance(message, ModelRequest): continue parts = [ @@ -162,6 +152,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): search_count: int = field(default=0, repr=False) request_count: int = field(default=0, repr=False) grace_requests_used: int = field(default=0, repr=False) + turn_start: int = field(default=0, repr=False) async def for_run(self, ctx: RunContext[Any]) -> "RAGCapabilityBase[StateT]": outer = getattr(ctx.deps, "state", None) @@ -179,6 +170,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): search_count=0, request_count=0, grace_requests_used=0, + turn_start=len(ctx.messages), ) run_capability._sync_state() return run_capability @@ -202,7 +194,9 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): destroy the host's record of what was retrieved. """ request_context.messages = _compact_old_tool_returns( - request_context.messages, self.tool_names + request_context.messages, + self.tool_names - {self._cite_tool_name}, + turn_start=self.turn_start, ) return await handler(request_context) diff --git a/tests/capabilities/test_capabilities.py b/tests/capabilities/test_capabilities.py index f107df51..aa360a73 100644 --- a/tests/capabilities/test_capabilities.py +++ b/tests/capabilities/test_capabilities.py @@ -803,7 +803,9 @@ def test_prior_turn_tool_results_are_compacted_but_current_evidence_is_kept(): ), ] - compacted = _compact_old_tool_returns(messages, frozenset({"rag_search"})) + compacted = _compact_old_tool_returns( + messages, frozenset({"rag_search"}), turn_start=3 + ) old_return = compacted[2].parts[0] current_return = compacted[5].parts[0] @@ -813,7 +815,7 @@ def test_prior_turn_tool_results_are_compacted_but_current_evidence_is_kept(): assert current_return.content == "current evidence" -def test_tool_results_are_unchanged_when_history_has_no_user_prompt(): +def test_nothing_is_compacted_on_the_first_question(): messages = [ ModelResponse(parts=[ToolCallPart("rag_search", {}, "current-call")]), ModelRequest( @@ -821,7 +823,9 @@ def test_tool_results_are_unchanged_when_history_has_no_user_prompt(): ), ] - compacted = _compact_old_tool_returns(messages, frozenset({"rag_search"})) + compacted = _compact_old_tool_returns( + messages, frozenset({"rag_search"}), turn_start=0 + ) assert compacted is messages current_return = compacted[1].parts[0] @@ -847,18 +851,33 @@ def _search_exchange(call_id: str, evidence: str, *, images: bool): ] -def test_image_bearing_tool_return_does_not_end_the_current_turn(): +@pytest.mark.parametrize( + ("label", "trailing"), + [ + ("page images on a tool return", [UserPromptPart(content=[PAGE_IMAGE])]), + ("a notice injected mid-run", [UserPromptPart("You answered without citing")]), + ], +) +def test_current_turn_evidence_survives_later_user_prompt_parts(label, trailing): + """Nothing that arrives mid-question may be read as the next question. + + Page images on a tool return and a notice this capability injects both + appear as a ``UserPromptPart`` after the question, so deriving the turn from + message shape stripped evidence the model was still answering from. + """ messages = [ ModelRequest(parts=[UserPromptPart("current question")]), *_search_exchange("first-call", "first evidence", images=False), - *_search_exchange("second-call", "second evidence", images=True), + ModelRequest(parts=trailing), ] - compacted = _compact_old_tool_returns(messages, frozenset({"rag_search"})) + compacted = _compact_old_tool_returns( + messages, frozenset({"rag_search"}), turn_start=0 + ) first_return = compacted[2].parts[0] assert isinstance(first_return, ToolReturnPart) - assert first_return.content == "first evidence" + assert first_return.content == "first evidence", label def test_prior_turn_images_outlive_their_tool_return(): @@ -874,7 +893,9 @@ def test_prior_turn_images_outlive_their_tool_return(): ModelRequest(parts=[UserPromptPart("current question")]), ] - compacted = _compact_old_tool_returns(messages, frozenset({"rag_search"})) + compacted = _compact_old_tool_returns( + messages, frozenset({"rag_search"}), turn_start=4 + ) old_return = compacted[2].parts[0] assert isinstance(old_return, ToolReturnPart) @@ -886,7 +907,7 @@ def test_prior_turn_images_outlive_their_tool_return(): ) -def test_user_attached_image_starts_a_turn_and_is_never_dropped(): +def test_user_attached_image_is_never_dropped(): messages = [ ModelRequest(parts=[UserPromptPart("old question")]), *_search_exchange("old-call", "old evidence", images=False), @@ -894,7 +915,9 @@ def test_user_attached_image_starts_a_turn_and_is_never_dropped(): ModelRequest(parts=[UserPromptPart(content=[PAGE_IMAGE, "what is this?"])]), ] - compacted = _compact_old_tool_returns(messages, frozenset({"rag_search"})) + compacted = _compact_old_tool_returns( + messages, frozenset({"rag_search"}), turn_start=4 + ) old_return = compacted[2].parts[0] assert isinstance(old_return, ToolReturnPart) @@ -940,3 +963,56 @@ async def test_compaction_never_reaches_the_stored_message_history(temp_db_path) ] assert "REAL EVIDENCE" in returns assert PRIOR_TURN_NOTICE not in returns + + +@pytest.mark.asyncio +async def test_prior_turn_search_is_compacted_but_its_cite_receipt_is_not(temp_db_path): + """Only evidence is compacted, and the boundary comes from the run. + + A cite acknowledgement is a receipt, not evidence: replacing it lengthened + the request and erased the record that citations had been registered. + """ + capability = create_rag( + db_path=temp_db_path, config=AppConfig(), defer_loading=False + ) + turns = iter( + [ + ModelResponse(parts=[ToolCallPart("rag_search", {"query": "q"}, "call-1")]), + ModelResponse( + parts=[ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-2")] + ), + ModelResponse(parts=[TextPart("first answer")]), + ModelResponse(parts=[TextPart("second answer")]), + ] + ) + wire: list[list[Any]] = [] + + async def model(messages, _info): + wire.append(list(messages)) + return next(turns) + + agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[capability]) + deps = Deps(state={"rag": RAGState().model_dump(mode="json")}) + + async def search(self, query: str, _limit: int | None) -> str: + """Record a result the way the real search does, so citing resolves.""" + cast(Any, self.state).searches[query] = [ + SearchResult(content="evidence", score=1.0, chunk_id="chunk-1") + ] + return "REAL EVIDENCE" + + with patch.object(RAGCapability, "_search", search): + first = await agent.run("old question", deps=deps) + await agent.run( + "current question", deps=deps, message_history=first.all_messages() + ) + + prior = { + part.tool_name: str(part.content) + for message in wire[-1] + if isinstance(message, ModelRequest) + for part in message.parts + if isinstance(part, ToolReturnPart) + } + assert prior["rag_search"] == PRIOR_TURN_NOTICE + assert prior["rag_cite"] != PRIOR_TURN_NOTICE From 46f7ab8d97c78233d8c76b65270e33b6bb8d5b3e Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Sat, 8 Aug 2026 17:51:06 +0300 Subject: [PATCH 03/13] Name the search result each page image belongs to ToolReturn.content reaches the model as a user-role message and the pictures arrive bare, so nothing connects a figure to the chunk it came from: BinaryContent.identifier does not survive serialization to the vision API, and the captions in the result text correlate only by position. Precede each picture with its position, source chunk id and self_ref. build_binary_parts_from_results becomes build_image_content_from_results and returns the labels interleaved with the pictures, so both attachment sites emit them the same way. This does not stop a model narrating retrieved pictures as user-supplied. Measured on gemma4-26b with a single note ahead of the batch, and again with per-image labels: it quotes the label and still says the user provided them. The message role wins over its text. --- CHANGELOG.md | 1 + .../haiku/rag/capabilities/_base.py | 4 +- .../haiku/rag/store/models/chunk.py | 2 +- haiku_rag_slim/haiku/rag/tools/search.py | 46 +++++++---- tests/test_picture_in_context.py | 12 +-- tests/tools/test_search.py | 82 +++++++++++++++---- 6 files changed, 109 insertions(+), 38 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f75350e7..5bda89f7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,7 @@ - `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. +- Each page image attached to a search result is preceded by a line giving its position and the chunk id it came from. `build_binary_parts_from_results` is now `build_image_content_from_results` and returns those labels interleaved with the pictures. - Prior-question tool output is trimmed from the model request only; `all_messages()` retains what the run gathered. - Evidence retrieved for the current question is no longer trimmed mid-question, whether or not it carries page images. - `rag_cite` / `analysis_cite` returns are no longer replaced by the prior-question notice. diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index 41994ed8..5be8e73e 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py @@ -27,7 +27,7 @@ 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. @@ -372,7 +372,7 @@ 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)): + if self.vision and (parts := build_image_content_from_results(results)): return ToolReturn(return_value=formatted, content=parts) return formatted 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..787210b6 100644 --- a/haiku_rag_slim/haiku/rag/tools/search.py +++ b/haiku_rag_slim/haiku/rag/tools/search.py @@ -11,10 +11,10 @@ from haiku.rag.store.models import SearchResult from haiku.rag.tools.context import RAGDeps -def build_binary_parts_from_results( +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 +22,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, bytes]] = [] seen: set[tuple[str | None, str]] = set() for result in results: if not result.image_data: @@ -38,15 +48,21 @@ def build_binary_parts_from_results( img.verify() except Exception: continue - parts.append( - BinaryContent( - data=data, - media_type="image/png", - identifier=self_ref, - ) - ) + collected.append((result.chunk_id, self_ref, data)) seen.add(key) - return parts + + content: list[str | BinaryContent] = [] + total = len(collected) + for position, (chunk_id, self_ref, data) in enumerate(collected, 1): + content.append( + f"Page image {position} of {total}, retrieved from the knowledge " + f"base for search result [{chunk_id}] ({self_ref}). " + "Not provided by the user." + ) + content.append( + BinaryContent(data=data, media_type="image/png", identifier=self_ref) + ) + return content def create_search_toolset( @@ -134,9 +150,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/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() From 8eef0b37344f8123fc5de8faa2ec425233a2ef3a Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 10 Aug 2026 12:15:47 +0300 Subject: [PATCH 04/13] Leave a resumed run's evidence alone MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A run carrying no prompt is continuing a question rather than asking one: pydantic-ai resumes that way for deferred tool results, interruptions and suspended responses. len(ctx.messages) then counts the live question's own messages, so its search result was replaced by the earlier-question notice and the model was asked to answer with the evidence removed. Reproduced: resuming with an in-flight history left a notice where the only evidence was. Switch compaction off for the whole run when ctx.prompt is None. The absence of a prompt is the signal rather than the message layout — the resume shapes differ from each other, and deriving the boundary from layout is what broke this to begin with. --- CHANGELOG.md | 1 + .../haiku/rag/capabilities/_base.py | 13 ++- tests/capabilities/test_capabilities.py | 81 ++++++++++++++++++- 3 files changed, 93 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5bda89f7..3944f227 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,7 @@ - Each page image attached to a search result is preceded by a line giving its position and the chunk id it came from. `build_binary_parts_from_results` is now `build_image_content_from_results` and returns those labels interleaved with the pictures. - Prior-question tool output is trimmed from the model request only; `all_messages()` retains what the run gathered. - Evidence retrieved for the current question is no longer trimmed mid-question, whether or not it carries page images. +- A run resumed without a prompt (deferred tool results, interruption, suspension) no longer has the active question's evidence trimmed as if it belonged to an earlier one. - `rag_cite` / `analysis_cite` returns are no longer replaced by the prior-question notice. ## [0.73.0] - 2026-08-06 diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index 5be8e73e..789204d4 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py @@ -155,6 +155,17 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): turn_start: int = field(default=0, repr=False) async def for_run(self, ctx: RunContext[Any]) -> "RAGCapabilityBase[StateT]": + """Start a run's own copy, and decide what counts as an earlier question. + + A run carrying no prompt is continuing a question rather than asking one: + pydantic-ai resumes that way for deferred tool results, interruptions and + suspended responses. ``len(ctx.messages)`` would then count the live + question's own messages and hand the model a notice where its search + result should be, so compaction is switched off for the whole run + (``turn_start=0``). The absence of a prompt is the signal — the three + resume shapes differ in message layout, and reading the layout is what + broke this in the first place. + """ 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 @@ -170,7 +181,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): search_count=0, request_count=0, grace_requests_used=0, - turn_start=len(ctx.messages), + turn_start=0 if ctx.prompt is None else len(ctx.messages), ) run_capability._sync_state() return run_capability diff --git a/tests/capabilities/test_capabilities.py b/tests/capabilities/test_capabilities.py index aa360a73..7a533848 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 ( BinaryContent, ModelRequest, @@ -965,6 +965,85 @@ async def test_compaction_never_reaches_the_stored_message_history(temp_db_path) assert PRIOR_TURN_NOTICE not in returns +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 _wire_returns(sent: list[Any]) -> list[str]: + return [ + str(part.content) + for message in sent + if isinstance(message, ModelRequest) + for part in message.parts + if isinstance(part, ToolReturnPart) + ] + + +@pytest.mark.asyncio +async def test_a_new_question_compacts_the_previous_one(temp_db_path): + """The baseline the resume cases are contrasted against.""" + capability = create_rag( + db_path=temp_db_path, config=AppConfig(), defer_loading=False + ) + wire: list[list[Any]] = [] + + async def model(messages, _info): + wire.append(list(messages)) + return ModelResponse(parts=[TextPart("answer")]) + + agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[capability]) + + await agent.run( + "a different question", deps=Deps(), message_history=_in_flight_history() + ) + + assert _wire_returns(wire[-1]) == [PRIOR_TURN_NOTICE] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "resume_kwargs", + [ + pytest.param({}, id="no prompt"), + pytest.param( + {"deferred_tool_results": DeferredToolResults()}, id="deferred results" + ), + ], +) +async def test_a_resumed_run_keeps_the_active_questions_evidence( + temp_db_path, resume_kwargs +): + """A run without a prompt continues a question; nothing in it is prior. + + ``len(ctx.messages)`` cannot tell the two apart — on a resumption it counts + the live question's own messages and marks its evidence as earlier-question + evidence, leaving the model to answer with a notice where its search result + used to be. Deferred, interrupted and suspended resumes all differ in shape, + so the absence of a prompt is the signal rather than the message layout. + """ + capability = create_rag( + db_path=temp_db_path, config=AppConfig(), defer_loading=False + ) + wire: list[list[Any]] = [] + + async def model(messages, _info): + wire.append(list(messages)) + return ModelResponse(parts=[TextPart("answer")]) + + agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[capability]) + + await agent.run(deps=Deps(), message_history=_in_flight_history(), **resume_kwargs) + + assert _wire_returns(wire[-1]) == ["EVIDENCE FOR THE LIVE TURN"] + + @pytest.mark.asyncio async def test_prior_turn_search_is_compacted_but_its_cite_receipt_is_not(temp_db_path): """Only evidence is compacted, and the boundary comes from the run. From fa0c4a4f500cd07989d3fde5d918e2de3117a5ba Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 10 Aug 2026 14:13:19 +0300 Subject: [PATCH 05/13] Treat an unfinished history tail as a continuation too MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit deferred_tool_results may arrive with a non-empty prompt, so the absence of a prompt cannot be the only test for a resumption. Reproduced: resuming a live question that way replaced its search result with the earlier-question notice while the deferred result arrived alongside it. _is_resumption accepts either signal — no prompt, or a history ending with a request the model has not answered or a response whose tool calls have no returns. A settled history ends with the previous answer, so a genuinely new question is unaffected. A new prompt on top of an unanswered tail is ambiguous and now counts as a continuation: compacting costs the answer if it is one, while not compacting only costs a larger request. --- .../haiku/rag/capabilities/_base.py | 42 ++++++++--- tests/capabilities/test_capabilities.py | 70 ++++++++++++++++++- 2 files changed, 101 insertions(+), 11 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index 789204d4..5803865c 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py @@ -123,6 +123,31 @@ def _compact_old_tool_returns( return compacted +def _is_resumption(prompt: Any, messages: list[ModelMessage]) -> bool: + """Whether this run continues a question rather than asking a new one. + + 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. + """ + 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: """Whether the model's most recent response called one of these tools.""" for message in reversed(messages): @@ -157,14 +182,11 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): async def for_run(self, ctx: RunContext[Any]) -> "RAGCapabilityBase[StateT]": """Start a run's own copy, and decide what counts as an earlier question. - A run carrying no prompt is continuing a question rather than asking one: - pydantic-ai resumes that way for deferred tool results, interruptions and - suspended responses. ``len(ctx.messages)`` would then count the live - question's own messages and hand the model a notice where its search - result should be, so compaction is switched off for the whole run - (``turn_start=0``). The absence of a prompt is the signal — the three - resume shapes differ in message layout, and reading the layout is what - broke this in the first place. + On a resumption (see ``_is_resumption``) compaction is switched off for the + whole run: ``len(ctx.messages)`` would count the live question's own + messages and replace its evidence with the earlier-question notice, leaving + the model to answer with the evidence taken away. Failing this way costs a + larger request; failing the other way costs the answer. """ outer = getattr(ctx.deps, "state", None) outer_state = outer if isinstance(outer, dict) else None @@ -181,7 +203,9 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): search_count=0, request_count=0, grace_requests_used=0, - turn_start=0 if ctx.prompt is None else len(ctx.messages), + turn_start=( + 0 if _is_resumption(ctx.prompt, ctx.messages) else len(ctx.messages) + ), ) run_capability._sync_state() return run_capability diff --git a/tests/capabilities/test_capabilities.py b/tests/capabilities/test_capabilities.py index 7a533848..7b37b11f 100644 --- a/tests/capabilities/test_capabilities.py +++ b/tests/capabilities/test_capabilities.py @@ -988,7 +988,39 @@ def _wire_returns(sent: list[Any]) -> list[str]: @pytest.mark.asyncio async def test_a_new_question_compacts_the_previous_one(temp_db_path): - """The baseline the resume cases are contrasted against.""" + """The baseline the resume cases are contrasted against. + + A settled history — the previous question answered — is what a genuinely new + question follows. An unfinished tail is ambiguous instead, and treated as a + continuation. + """ + capability = create_rag( + db_path=temp_db_path, config=AppConfig(), defer_loading=False + ) + wire: list[list[Any]] = [] + + async def model(messages, _info): + wire.append(list(messages)) + return ModelResponse(parts=[TextPart("answer")]) + + agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[capability]) + settled = [*_in_flight_history(), ModelResponse(parts=[TextPart("first answer")])] + + await agent.run("a different question", deps=Deps(), message_history=settled) + + assert _wire_returns(wire[-1]) == [PRIOR_TURN_NOTICE] + + +@pytest.mark.asyncio +async def test_a_prompt_on_an_unanswered_tail_is_treated_as_a_continuation( + temp_db_path, +): + """Ambiguous shape, resolved the safe way. + + A history ending in a request the model never answered, plus a new prompt, + could be a fresh question or a continuation. Compacting would cost the answer + if it is a continuation; not compacting only costs a larger request. + """ capability = create_rag( db_path=temp_db_path, config=AppConfig(), defer_loading=False ) @@ -1004,7 +1036,41 @@ async def test_a_new_question_compacts_the_previous_one(temp_db_path): "a different question", deps=Deps(), message_history=_in_flight_history() ) - assert _wire_returns(wire[-1]) == [PRIOR_TURN_NOTICE] + assert _wire_returns(wire[-1]) == ["EVIDENCE FOR THE LIVE TURN"] + + +@pytest.mark.asyncio +async def test_a_resume_carrying_a_prompt_keeps_the_active_evidence(temp_db_path): + """Deferred results may arrive with a prompt, and that is still a continuation. + + ``ctx.prompt`` is non-null here, so the absence of a prompt cannot be the only + signal: the history tail is unfinished — a response whose tool call has no + return yet — and the question it belongs to is still being answered. + """ + capability = create_rag( + db_path=temp_db_path, config=AppConfig(), defer_loading=False + ) + wire: list[list[Any]] = [] + + async def model(messages, _info): + wire.append(list(messages)) + return ModelResponse(parts=[TextPart("answer")]) + + agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[capability]) + history = [ + *_in_flight_history(), + ModelResponse(parts=[ToolCallPart("external_tool", {}, "call-2")]), + ] + + await agent.run( + "carry on", + deferred_tool_results=DeferredToolResults(calls={"call-2": "external result"}), + message_history=history, + deps=Deps(), + ) + + assert "EVIDENCE FOR THE LIVE TURN" in _wire_returns(wire[-1]) + assert PRIOR_TURN_NOTICE not in _wire_returns(wire[-1]) @pytest.mark.asyncio From 85594a1fd198a75e2c1d1384463ef7e354ac01bf Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 10 Aug 2026 15:09:53 +0300 Subject: [PATCH 06/13] Record what each evidence capability retrieved and cited `CapabilityEvidenceRecord` holds the relationships a transcript cannot express: which chunks a capability retrieved, which it cited, in which questions, and at which point in the conversation. RAG and analysis each own one in their own state namespace. Nothing is co-written: the host's state is JSON storage, so a shared record would be overwritten by whichever capability synced last, and merging happens in transient per-request views instead. Both clocks are derived from the conversation rather than counted locally, so every participant computes the same values without sharing a counter. Question identity is the message count when the question arrived; epoch is the message count at an outcome. Epochs are therefore globally comparable, which is what lets `citation_status` require a declaration to follow the newest evidence of every capability, and what makes equal epochs mean one request. A declaration is written only after `resolve_citations` succeeds, so a call naming only unresolvable ids is not a citation. Status is derived, never stored, so refs and status cannot contradict. Resuming a question requires the host to carry the capability state from the run being resumed. Without it the identity of the question in progress is unknowable, and adopting the current message count would relabel that question as a new one and judge every declaration in it against the wrong identity. Nothing reads the records yet and no wire behaviour changes. --- CHANGELOG.md | 8 + .../haiku/rag/capabilities/_base.py | 64 +++- .../haiku/rag/capabilities/analysis.py | 3 + .../haiku/rag/capabilities/ledger.py | 155 ++++++++ haiku_rag_slim/haiku/rag/capabilities/rag.py | 2 + tests/capabilities/test_capabilities.py | 346 +++++++++++++++++- tests/capabilities/test_evidence_ledger.py | 150 ++++++++ 7 files changed, 722 insertions(+), 6 deletions(-) create mode 100644 haiku_rag_slim/haiku/rag/capabilities/ledger.py create mode 100644 tests/capabilities/test_evidence_ledger.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 3944f227..9bd1b8b2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,14 @@ # Changelog ## [Unreleased] +### Added + +- `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. + +### Changed + +- 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. + ### Fixed - `cross-encoder` reranking no longer ties the scores of strongly-relevant candidates, which left their order to the sort. Scores remain 0-1. diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index 5803865c..626a5652 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py @@ -23,6 +23,7 @@ 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 @@ -178,6 +179,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): request_count: int = field(default=0, repr=False) grace_requests_used: int = field(default=0, repr=False) turn_start: int = field(default=0, repr=False) + epoch: int = field(default=0, repr=False) async def for_run(self, ctx: RunContext[Any]) -> "RAGCapabilityBase[StateT]": """Start a run's own copy, and decide what counts as an earlier question. @@ -187,12 +189,31 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): messages and replace its evidence with the earlier-question notice, leaving the model to answer with the evidence taken away. Failing this way costs a larger request; failing the other way costs the answer. + + 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 — unless there is no history at all, where an absent prompt + means an instructions-only first question and nothing is in progress. """ 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) + if resuming and ctx.messages and not (raw_state or {}).get("evidence"): + 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." + ) state = self.state_type.model_validate(raw_state or {}) _clear_invocation_state(state) + if not resuming: + cast(Any, state).evidence.question = len(ctx.messages) run_capability = replace( self, state=state, @@ -203,9 +224,8 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): search_count=0, request_count=0, grace_requests_used=0, - turn_start=( - 0 if _is_resumption(ctx.prompt, ctx.messages) else len(ctx.messages) - ), + epoch=0, + turn_start=0 if resuming else len(ctx.messages), ) run_capability._sync_state() return run_capability @@ -238,6 +258,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): async def before_model_request( self, ctx: RunContext[Any], request_context: ModelRequestContext ) -> ModelRequestContext: + self.epoch = len(ctx.messages) if instruction := self._budget_notice(): current_request = request_context.messages[-1] if isinstance(current_request, ModelRequest): @@ -390,6 +411,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 @@ -407,6 +463,7 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): ) state = cast(Any, self.state) state.searches[query] = results + self._note_evidence() if self.vision and (parts := build_image_content_from_results(results)): return ToolReturn(return_value=formatted, content=parts) return formatted @@ -454,6 +511,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..a9495e7f 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) @@ -107,6 +109,7 @@ class AnalysisCapability(RAGCapabilityBase[AnalysisState]): ) sandbox = await self._ensure_sandbox() result = await sandbox.execute(code) + 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/ledger.py b/haiku_rag_slim/haiku/rag/capabilities/ledger.py new file mode 100644 index 00000000..83868aad --- /dev/null +++ b/haiku_rag_slim/haiku/rag/capabilities/ledger.py @@ -0,0 +1,155 @@ +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. + """ + + occurrences: dict[str, EvidenceOccurrence] = Field(default_factory=dict) + question: int = 0 + latest_evidence_epoch: int = 0 + declaration: CitationDeclaration | None = 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.latest_evidence_epoch = max(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. + """ + 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/tests/capabilities/test_capabilities.py b/tests/capabilities/test_capabilities.py index 7b37b11f..98e6cc43 100644 --- a/tests/capabilities/test_capabilities.py +++ b/tests/capabilities/test_capabilities.py @@ -26,6 +26,10 @@ from haiku.rag.capabilities._base import ( ) 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 @@ -463,6 +467,44 @@ 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", [True, False]) +@pytest.mark.asyncio +async def test_a_code_execution_is_evidence_even_when_it_fails(temp_db_path, success): + """Output the model can read grounds an answer, whether the code raised or not.""" + capability = create_analysis(db_path=temp_db_path, config=AppConfig()) + capability.state = AnalysisState() + capability.epoch = 5 + sandbox = AsyncMock(spec=Sandbox) + sandbox._search_results = [] + sandbox.execute.return_value = SandboxResult( + stdout="42", 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 == 5 + + @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. @@ -965,6 +1007,22 @@ async def test_compaction_never_reaches_the_stored_message_history(temp_db_path) assert PRIOR_TURN_NOTICE not in returns +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 [ @@ -1033,7 +1091,9 @@ async def test_a_prompt_on_an_unanswered_tail_is_treated_as_a_continuation( agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[capability]) await agent.run( - "a different question", deps=Deps(), message_history=_in_flight_history() + "a different question", + deps=_resuming_deps(), + message_history=_in_flight_history(), ) assert _wire_returns(wire[-1]) == ["EVIDENCE FOR THE LIVE TURN"] @@ -1066,7 +1126,7 @@ async def test_a_resume_carrying_a_prompt_keeps_the_active_evidence(temp_db_path "carry on", deferred_tool_results=DeferredToolResults(calls={"call-2": "external result"}), message_history=history, - deps=Deps(), + deps=_resuming_deps(), ) assert "EVIDENCE FOR THE LIVE TURN" in _wire_returns(wire[-1]) @@ -1105,7 +1165,9 @@ async def test_a_resumed_run_keeps_the_active_questions_evidence( agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[capability]) - await agent.run(deps=Deps(), message_history=_in_flight_history(), **resume_kwargs) + await agent.run( + deps=_resuming_deps(), message_history=_in_flight_history(), **resume_kwargs + ) assert _wire_returns(wire[-1]) == ["EVIDENCE FOR THE LIVE TURN"] @@ -1161,3 +1223,281 @@ async def test_prior_turn_search_is_compacted_but_its_cite_receipt_is_not(temp_d } assert prior["rag_search"] == PRIOR_TURN_NOTICE assert prior["rag_cite"] != PRIOR_TURN_NOTICE + + +def _record(deps: Deps, namespace: str) -> CapabilityEvidenceRecord: + return CapabilityEvidenceRecord.model_validate(deps.state[namespace]["evidence"]) + + +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()) + + assert first_identity == 0 + assert _record(deps, "rag").question > first_identity + assert _record(deps, "analysis").question == _record(deps, "rag").question + + +@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")]), + ] + + await agent.run( + "carry on", + deferred_tool_results=DeferredToolResults(calls={"call-2": "external result"}), + message_history=history, + deps=deps, + ) + + 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") + 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 == [record.question] + assert citation_status([record], question=record.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") + assert record.declaration is not None + assert citation_status([record], question=record.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") + assert record.declaration is not None + assert record.declaration.epoch == record.latest_evidence_epoch + assert citation_status([record], question=record.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, + ] + + +@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") + 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=record.question) == "grounded" diff --git a/tests/capabilities/test_evidence_ledger.py b/tests/capabilities/test_evidence_ledger.py new file mode 100644 index 00000000..d9c264c9 --- /dev/null +++ b/tests/capabilities/test_evidence_ledger.py @@ -0,0 +1,150 @@ +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() + grounded.declare([rag_ref()], epoch=1) + + ungrounded = CapabilityEvidenceRecord() + 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(): + """Epochs outlive a question, so the epoch alone would inherit it.""" + record = CapabilityEvidenceRecord(question=2) + record.declare([rag_ref()], epoch=3) + assert citation_status([record], question=2) == "grounded" + + record.question = 8 + + assert record.declaration is not None + 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() + 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() + cited.note_evidence(3) + cited.declare([rag_ref()], epoch=5) + searched_after = CapabilityEvidenceRecord() + 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() + 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() + 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() + grounded_then_empty.declare([rag_ref()], epoch=3) + grounded_then_empty.declare([], epoch=3) + + empty_then_grounded = CapabilityEvidenceRecord() + 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() + rag.declare([EvidenceRef(capability="rag", chunk_id="shared")], epoch=3) + analysis = CapabilityEvidenceRecord() + analysis.declare([EvidenceRef(capability="analysis", chunk_id="shared")], epoch=3) + + assert rag.occurrences["shared"].capability == "rag" + assert analysis.occurrences["shared"].capability == "analysis" + + +def test_an_evidence_epoch_never_moves_backwards(): + record = CapabilityEvidenceRecord() + record.note_evidence(9) + record.note_evidence(4) + + assert record.latest_evidence_epoch == 9 + + +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.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()] + ) From 42d923fe4bb16555c594ca329ebd8a6bfe297dad Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 11 Aug 2026 10:08:10 +0300 Subject: [PATCH 07/13] Make the ledger's clocks defensible against the host A question identity is unset until a run establishes it. Testing whether the record exists cannot stand in for that: a host with no state to send seeds a default record, and a default record is truthy, so a resumption missing its real state would have proceeded as question zero. Every way of recording a message count now refuses one behind what is already stored, through one shared check: a new identity, an evidence outcome and a declaration, each against the newest identity, evidence epoch and declaration epoch. Identities and epochs are only comparable while the conversation grows, and `before_model_request` results are assigned back onto history, so that is a constraint on the host rather than a guarantee of the framework. Left unchecked, an evidence outcome moving backwards freezes every later declaration as stale, and a declaration moving backwards replaces a newer one with an older one and revives the answer it grounded. The searches, citations and executions of a question in progress survive a resumption. Clearing them cost the results the model was still answering from: a citation afterwards recorded no provenance and could not resolve against the expanded result it had seen, falling through to a database lookup. A code execution counts as evidence when it succeeded or printed something. A raised error with an empty stdout grounds nothing. --- CHANGELOG.md | 1 + .../haiku/rag/capabilities/_base.py | 24 +++- .../haiku/rag/capabilities/analysis.py | 3 +- .../haiku/rag/capabilities/ledger.py | 42 +++++- tests/capabilities/test_capabilities.py | 129 +++++++++++++++--- tests/capabilities/test_evidence_ledger.py | 93 ++++++++++--- 6 files changed, 243 insertions(+), 49 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9bd1b8b2..b9f86243 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,6 +15,7 @@ - `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. - Prior-question tool output is trimmed from the model request only; `all_messages()` retains what the run gathered. - Evidence retrieved for the current question is no longer trimmed mid-question, whether or not it carries page images. diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index 626a5652..9752d2b0 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py @@ -70,6 +70,13 @@ 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"): @@ -196,24 +203,27 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): 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 — unless there is no history at all, where an absent prompt - means an instructions-only first question and nothing is in progress. + 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) - if resuming and ctx.messages and not (raw_state or {}).get("evidence"): + continuing = resuming and bool(ctx.messages) + state = self.state_type.model_validate(raw_state or {}) + 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." ) - state = self.state_type.model_validate(raw_state or {}) - _clear_invocation_state(state) - if not resuming: - cast(Any, state).evidence.question = len(ctx.messages) + if not continuing: + _clear_invocation_state(state) + record.begin_question(len(ctx.messages)) run_capability = replace( self, state=state, diff --git a/haiku_rag_slim/haiku/rag/capabilities/analysis.py b/haiku_rag_slim/haiku/rag/capabilities/analysis.py index a9495e7f..90bb8147 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/analysis.py +++ b/haiku_rag_slim/haiku/rag/capabilities/analysis.py @@ -109,7 +109,8 @@ class AnalysisCapability(RAGCapabilityBase[AnalysisState]): ) sandbox = await self._ensure_sandbox() result = await sandbox.execute(code) - self._note_evidence() + 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/ledger.py b/haiku_rag_slim/haiku/rag/capabilities/ledger.py index 83868aad..4d837f79 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/ledger.py +++ b/haiku_rag_slim/haiku/rag/capabilities/ledger.py @@ -54,14 +54,43 @@ class CapabilityEvidenceRecord(BaseModel): ``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. + 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 = 0 + 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. + + Identities and epochs are both message counts, and every comparison + between them assumes the conversation only grows. 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.""" + self._reject_regression(identity, "A question") + self.question = identity + def note_evidence(self, epoch: int) -> None: """Record that the model has seen an evidence outcome. @@ -70,7 +99,8 @@ class CapabilityEvidenceRecord(BaseModel): fruitless search grounds a refusal. Not called for a failure that yields no evidence at all, such as an exhausted budget. """ - self.latest_evidence_epoch = max(self.latest_evidence_epoch, epoch) + self._reject_regression(epoch, "Evidence") + self.latest_evidence_epoch = epoch def declare( self, @@ -86,6 +116,12 @@ class CapabilityEvidenceRecord(BaseModel): 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, diff --git a/tests/capabilities/test_capabilities.py b/tests/capabilities/test_capabilities.py index 98e6cc43..c0ab7f3f 100644 --- a/tests/capabilities/test_capabilities.py +++ b/tests/capabilities/test_capabilities.py @@ -314,7 +314,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"), @@ -340,7 +340,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"), @@ -370,6 +370,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( @@ -380,7 +381,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 @@ -482,17 +483,31 @@ async def test_a_spent_execution_budget_is_not_evidence(temp_db_path): assert capability.state.evidence.latest_evidence_epoch == 0 -@pytest.mark.parametrize("success", [True, False]) +@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_a_code_execution_is_evidence_even_when_it_fails(temp_db_path, success): - """Output the model can read grounds an answer, whether the code raised or not.""" +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() + capability.state = AnalysisState(evidence=CapabilityEvidenceRecord(question=0)) capability.epoch = 5 sandbox = AsyncMock(spec=Sandbox) sandbox._search_results = [] sandbox.execute.return_value = SandboxResult( - stdout="42", stderr="" if success else "boom", success=success + stdout=stdout, stderr="" if success else "boom", success=success ) capability.sandbox = sandbox @@ -502,7 +517,7 @@ async def test_a_code_execution_is_evidence_even_when_it_fails(temp_db_path, suc with pytest.raises(ToolFailed): await capability._execute_code("print(42)") - assert capability.state.evidence.latest_evidence_epoch == 5 + assert capability.state.evidence.latest_evidence_epoch == expected_epoch @pytest.mark.asyncio @@ -790,7 +805,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 @@ -1257,10 +1274,11 @@ async def test_a_question_takes_its_own_identity_and_both_capabilities_agree( 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 _record(deps, "rag").question > first_identity - assert _record(deps, "analysis").question == _record(deps, "rag").question + assert second_identity is not None and second_identity > 0 + assert _record(deps, "analysis").question == second_identity @pytest.mark.asyncio @@ -1348,10 +1366,12 @@ async def test_citing_after_searching_grounds_the_question(temp_db_path): 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 == [record.question] - assert citation_status([record], question=record.question) == "grounded" + assert record.occurrences["chunk-1"].retrieved_in_questions == [question] + assert citation_status([record], question=question) == "grounded" @pytest.mark.asyncio @@ -1376,8 +1396,10 @@ async def test_searching_after_citing_leaves_the_question_uncited(temp_db_path): 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=record.question) == "missing" + assert citation_status([record], question=question) == "missing" @pytest.mark.asyncio @@ -1406,9 +1428,11 @@ async def test_a_citation_in_the_same_request_as_its_search_is_not_current( 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=record.question) == "missing" + assert citation_status([record], question=question) == "missing" @pytest.mark.asyncio @@ -1444,6 +1468,7 @@ async def test_evidence_cited_in_two_questions_keeps_both_in_the_record(temp_db_ first_question, record.question, ] + assert record.question != first_question @pytest.mark.asyncio @@ -1496,8 +1521,78 @@ async def test_citing_without_searching_grounds_the_question(temp_db_path): 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=record.question) == "grounded" + 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" diff --git a/tests/capabilities/test_evidence_ledger.py b/tests/capabilities/test_evidence_ledger.py index d9c264c9..fa849c64 100644 --- a/tests/capabilities/test_evidence_ledger.py +++ b/tests/capabilities/test_evidence_ledger.py @@ -1,3 +1,5 @@ +import pytest + from haiku.rag.capabilities.ledger import ( CapabilityEvidenceRecord, CitationDeclaration, @@ -34,10 +36,10 @@ def test_no_declaration_reads_as_missing(): def test_refs_make_it_grounded_and_no_refs_make_it_ungrounded(): - grounded = CapabilityEvidenceRecord() + grounded = CapabilityEvidenceRecord(question=0) grounded.declare([rag_ref()], epoch=1) - ungrounded = CapabilityEvidenceRecord() + ungrounded = CapabilityEvidenceRecord(question=0) ungrounded.declare([], epoch=1) assert citation_status([grounded], question=0) == "grounded" @@ -50,7 +52,7 @@ def test_an_earlier_questions_declaration_is_never_current(): record.declare([rag_ref()], epoch=3) assert citation_status([record], question=2) == "grounded" - record.question = 8 + record.begin_question(8) assert record.declaration is not None assert citation_status([record], question=8) == "missing" @@ -58,7 +60,7 @@ def test_an_earlier_questions_declaration_is_never_current(): 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() + record = CapabilityEvidenceRecord(question=0) record.note_evidence(5) record.declare([rag_ref()], epoch=5) @@ -71,10 +73,10 @@ def test_a_citation_in_the_same_request_as_the_evidence_is_not_current(): def test_evidence_from_another_capability_after_citing_makes_it_uncited(): """Currency spans capabilities, which only works because epochs are global.""" - cited = CapabilityEvidenceRecord() + cited = CapabilityEvidenceRecord(question=0) cited.note_evidence(3) cited.declare([rag_ref()], epoch=5) - searched_after = CapabilityEvidenceRecord() + searched_after = CapabilityEvidenceRecord(question=0) searched_after.note_evidence(7) assert citation_status([cited], question=0) == "grounded" @@ -82,7 +84,7 @@ def test_evidence_from_another_capability_after_citing_makes_it_uncited(): def test_declarations_at_the_same_epoch_merge(): - record = CapabilityEvidenceRecord() + record = CapabilityEvidenceRecord(question=0) record.declare([rag_ref("c1")], epoch=3) record.declare([rag_ref("c2")], epoch=3) @@ -91,7 +93,7 @@ def test_declarations_at_the_same_epoch_merge(): def test_repeating_a_ref_at_the_same_epoch_does_not_duplicate_it(): - record = CapabilityEvidenceRecord() + record = CapabilityEvidenceRecord(question=0) record.declare([rag_ref()], epoch=3) record.declare([rag_ref()], epoch=3) @@ -100,11 +102,11 @@ def test_repeating_a_ref_at_the_same_epoch_does_not_duplicate_it(): def test_neither_cite_order_downgrades_a_grounded_declaration(): - grounded_then_empty = CapabilityEvidenceRecord() + grounded_then_empty = CapabilityEvidenceRecord(question=0) grounded_then_empty.declare([rag_ref()], epoch=3) grounded_then_empty.declare([], epoch=3) - empty_then_grounded = CapabilityEvidenceRecord() + empty_then_grounded = CapabilityEvidenceRecord(question=0) empty_then_grounded.declare([], epoch=3) empty_then_grounded.declare([rag_ref()], epoch=3) @@ -113,27 +115,19 @@ def test_neither_cite_order_downgrades_a_grounded_declaration(): def test_the_same_chunk_id_under_two_capabilities_stays_separate(): - rag = CapabilityEvidenceRecord() + rag = CapabilityEvidenceRecord(question=0) rag.declare([EvidenceRef(capability="rag", chunk_id="shared")], epoch=3) - analysis = CapabilityEvidenceRecord() + 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_an_evidence_epoch_never_moves_backwards(): - record = CapabilityEvidenceRecord() - record.note_evidence(9) - record.note_evidence(4) - - assert record.latest_evidence_epoch == 9 - - 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.question = 8 + record.begin_question(8) record.declare([rag_ref()], epoch=9) occurrence = record.occurrences["c1"] @@ -148,3 +142,60 @@ def test_a_declaration_records_the_question_and_epoch_it_was_made_at(): 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_cannot_start_before_what_is_already_recorded(): + record = CapabilityEvidenceRecord(question=0) + record.note_evidence(9) + + with pytest.raises(ValueError, match="append-only"): + record.begin_question(4) From 0aa6d79f88c175b057fed2476fd904e1c0e7909a Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 11 Aug 2026 12:59:02 +0300 Subject: [PATCH 08/13] Build one capsule of cited evidence from the records `EvidenceCompactionCapability` reads what the evidence capabilities recorded out of the run registry, and `build_capsule` renders it: every cited item, grouped by the question that last cited it, newest group first, each rendered once, with the pictures of cited evidence and the labels that must accompany them. Discovery runs one way and reads only, so no capability holds a reference to another and a host running one, both or neither needs no wiring change. Everything cited is kept whole and everything else is dropped. There is no character budget, no picture cap and nothing to configure: a cap would only half-rescue models that fail on long conversations regardless, and a host that needs earlier evidence pruned can compact its own requests further. A capability reports which of its tools produce evidence, so a cite acknowledgement is never mistaken for one. Pictures are identified by owner, document and reference, so one figure cited through overlapping chunks is attached once while the same reference in another document stays a different picture. The builder does no I/O and never sees the message history, so a picture travels with its label and the caller fetches the bytes. Nothing reaches the wire yet. --- .../haiku/rag/capabilities/_base.py | 11 +- .../haiku/rag/capabilities/analysis.py | 4 +- .../haiku/rag/capabilities/compaction.py | 243 +++++++++++ tests/capabilities/test_capabilities.py | 4 +- tests/capabilities/test_evidence_capsule.py | 386 ++++++++++++++++++ 5 files changed, 641 insertions(+), 7 deletions(-) create mode 100644 haiku_rag_slim/haiku/rag/capabilities/compaction.py create mode 100644 tests/capabilities/test_evidence_capsule.py diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index 9752d2b0..93f5ac0c 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py @@ -312,7 +312,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 " @@ -350,8 +350,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]: diff --git a/haiku_rag_slim/haiku/rag/capabilities/analysis.py b/haiku_rag_slim/haiku/rag/capabilities/analysis.py index 90bb8147..e2de8868 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/analysis.py +++ b/haiku_rag_slim/haiku/rag/capabilities/analysis.py @@ -87,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() 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..7179ba27 --- /dev/null +++ b/haiku_rag_slim/haiku/rag/capabilities/compaction.py @@ -0,0 +1,243 @@ +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from typing import Any, cast + +from pydantic_ai import RunContext +from pydantic_ai.capabilities import AbstractCapability + +from haiku.rag.capabilities._base import RAGCapabilityBase +from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord +from haiku.rag.store.models.citation import Citation + +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." + ) + + +@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)) + + +@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. + """ + + 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", + "create_capability", + "group_label", + "picture_label", +] diff --git a/tests/capabilities/test_capabilities.py b/tests/capabilities/test_capabilities.py index c0ab7f3f..2e140f6f 100644 --- a/tests/capabilities/test_capabilities.py +++ b/tests/capabilities/test_capabilities.py @@ -614,7 +614,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 @@ -629,7 +629,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 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 == "" From 2cd568847e7cee3a1494e4007fc5eb0e996e8311 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 11 Aug 2026 13:47:30 +0300 Subject: [PATCH 09/13] Move the wire rewrite into the compaction capability `_compact_old_tool_returns`, `PRIOR_TURN_NOTICE` and `turn_start` leave `RAGCapabilityBase`, along with its `wrap_model_request` hook. The evidence capabilities now retrieve and validate, and nothing else. Registering the compaction capability is what rewrites a request; leaving it out sends the transcript untouched, which was never a choice a host could make before. The boundary is the recorded question identity rather than message shape, so a resumption compacts what lies below the question in progress instead of switching compaction off for the whole run. The newest earlier evidence return carries the capsule and every other becomes a receipt, so one capsule exists by construction and every return stays paired with its call. Pictures of cited evidence are fetched through the capability that retrieved them and re-attached beside the capsule with fresh labels. Ownership of a picture on the wire requires the machine tag we write and an image directly after it, since neither position nor prose is proof: several tools' results can arrive in one request, and a user quoting our wording above their own picture had it removed. A picture that cannot be fetched or decoded is emitted with neither its image nor its label. The chat TUI and the example backend register the compactor, being multi-turn. `client.ask`, `client.analyze` and the MCP tools do not: a single-shot question has nothing earlier to compact. --- CHANGELOG.md | 12 +- app/backend/main.py | 7 +- docs/capabilities/index.md | 37 +- .../haiku/rag/capabilities/_base.py | 93 +-- .../haiku/rag/capabilities/compaction.py | 201 +++++- haiku_rag_slim/haiku/rag/chat/app.py | 5 +- haiku_rag_slim/haiku/rag/tools/search.py | 47 +- tests/capabilities/test_capabilities.py | 372 +--------- tests/capabilities/test_evidence_wire.py | 655 ++++++++++++++++++ 9 files changed, 969 insertions(+), 460 deletions(-) create mode 100644 tests/capabilities/test_evidence_wire.py diff --git a/CHANGELOG.md b/CHANGELOG.md index b9f86243..f3175b97 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -3,12 +3,20 @@ ### Added +- `EvidenceCompactionCapability` (`haiku.rag.capabilities.compaction.create_capability`): registering it replaces earlier questions' evidence on the model request with the evidence that was cited, grouped by the question that cited it, cited page images re-attached, other earlier evidence returns reduced to a receipt. Requests only; `all_messages()` is untouched. No configuration. - `RAGState.evidence` / `AnalysisState.evidence` (`CapabilityEvidenceRecord`): which evidence a capability retrieved and cited, per question, keyed by message-count question identities and epochs. `haiku.rag.capabilities.ledger.citation_status(records, question=...)` derives `missing` / `grounded` / `ungrounded` across capabilities. +- `RAGCapabilityBase.evidence_tool_names()` and `get_picture_bytes()`. +- `haiku.rag.tools.search.decode_picture()`. ### Changed +- `RAGCapability` and `AnalysisCapability` no longer rewrite the model request. Register `create_capability()` from `haiku.rag.capabilities.compaction` alongside them to keep earlier questions compacted. - Resuming a run (no prompt, deferred tool results, an unfinished history tail) raises `RuntimeError` unless the host carries the capability state from the run being resumed. +### Removed + +- `PRIOR_TURN_NOTICE` and `_compact_old_tool_returns` from `haiku.rag.capabilities._base`, and `RAGCapabilityBase.turn_start`. + ### Fixed - `cross-encoder` reranking no longer ties the scores of strongly-relevant candidates, which left their order to the sort. Scores remain 0-1. @@ -17,10 +25,6 @@ - `docling-local` reuses one docling `DocumentConverter` per set of conversion options instead of building one per document, so local layout, table and OCR models are no longer loaded per document. Conversions through a shared converter are serialized. - A resumed run keeps the searches, citations and executions of the question in progress instead of clearing them. - Each page image attached to a search result is preceded by a line giving its position and the chunk id it came from. `build_binary_parts_from_results` is now `build_image_content_from_results` and returns those labels interleaved with the pictures. -- Prior-question tool output is trimmed from the model request only; `all_messages()` retains what the run gathered. -- Evidence retrieved for the current question is no longer trimmed mid-question, whether or not it carries page images. -- A run resumed without a prompt (deferred tool results, interruption, suspension) no longer has the active question's evidence trimmed as if it belonged to an earlier one. -- `rag_cite` / `analysis_cite` returns are no longer replaced by the prior-question notice. ## [0.73.0] - 2026-08-06 diff --git a/app/backend/main.py b/app/backend/main.py index ad8a451d..252e938c 100644 --- a/app/backend/main.py +++ b/app/backend/main.py @@ -18,6 +18,9 @@ from starlette.requests import Request from starlette.responses import JSONResponse, Response, StreamingResponse from starlette.routing import Route +from haiku.rag.capabilities.compaction import ( + create_capability as create_compaction, +) from haiku.rag.capabilities.rag import AGENT_PREAMBLE, RAGState, create_capability from haiku.rag.client import HaikuRAG from haiku.rag.config import load_yaml_config @@ -81,7 +84,9 @@ capability = create_capability(db_path=db_path, config=Config, defer_loading=Fal agent = Agent( get_model(Config.qa.model, Config), instructions=AGENT_PREAMBLE, - capabilities=[capability], + # Conversations here are multi-turn, so earlier questions are reduced to the + # evidence they cited rather than carried whole. + capabilities=[capability, create_compaction()], deps_type=AppDeps, ) diff --git a/docs/capabilities/index.md b/docs/capabilities/index.md index 17c4ed6b..bff95f8e 100644 --- a/docs/capabilities/index.md +++ b/docs/capabilities/index.md @@ -1,13 +1,14 @@ # Capabilities -haiku.rag provides two native [Pydantic AI capabilities](https://ai.pydantic.dev/capabilities/): +haiku.rag provides native [Pydantic AI capabilities](https://ai.pydantic.dev/capabilities/): | Capability | Use it for | |---|---| | [`RAGCapability`](rag.md) | Grounded document search and citations. | | [`AnalysisCapability`](analysis.md) | Corpus computation and structural analysis with sandboxed Python. | +| `EvidenceCompactionCapability` | Optional. Shrinking a conversation's history to the evidence that was cited. | -Both capabilities are deferred by default. An agent initially sees only their descriptions and the standard `load_capability` tool. Instructions and tools enter the model context only when the model loads a capability. +The two evidence capabilities are deferred by default. An agent initially sees only their descriptions and the standard `load_capability` tool. Instructions and tools enter the model context only when the model loads a capability. ## Compose an agent @@ -34,6 +35,38 @@ agent = Agent( ) ``` +## Multi-turn conversations + +Every question adds its search results to the history, so requests grow turn after +turn, and can degrade answers or exceed a provider's limits as they do. Register the +compaction capability to replace earlier questions' evidence with the evidence that +was actually cited: + +```python +from haiku.rag.capabilities.compaction import create_capability as compaction +from haiku.rag.capabilities.rag import create_capability as rag + +agent = Agent( + "openai:gpt-5", + capabilities=[rag(db_path="my.lancedb"), compaction()], +) +``` + +Cited text and cited page images are kept in full, grouped by the question that +cited them, and stay citable by the same chunk ids. Everything else earlier becomes a +short receipt. Registering the capability is the only switch: leave it out and the +transcript reaches the model untouched. There is nothing to configure. + +Compaction rewrites the request, never the stored history, so `all_messages()` still +holds everything the run gathered. Retained evidence still grows with the +conversation — this reduces what a request carries, it does not bound it. A host that +needs more aggressive pruning can compact its own requests further, on the wire only. + +Resuming a question (deferred tool results, an interruption, a suspension) requires +the host to carry the capability state from the run being resumed, alongside the +message history. Without it the identity of the question in progress is unknowable +and the run fails rather than silently treating it as a new question. + ## State Capabilities use a plain `state: dict[str, Any]` attribute on agent dependencies when one is available. RAG state lives under `"rag"`; analysis state lives under `"analysis"`. This keeps state independent of any transport or UI protocol. diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index 93f5ac0c..65c43ff3 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py @@ -7,7 +7,7 @@ from typing import Any, cast from pydantic import BaseModel from pydantic_ai import ModelRetry, RunContext, ToolFailed -from pydantic_ai.capabilities import AbstractCapability, WrapModelRequestHandler +from pydantic_ai.capabilities import AbstractCapability from pydantic_ai.messages import ( InstructionPart, ModelMessage, @@ -15,7 +15,6 @@ from pydantic_ai.messages import ( ModelResponse, ToolCallPart, ToolReturn, - ToolReturnPart, ) from pydantic_ai.models import ModelRequestContext from pydantic_ai.run import AgentRunResult @@ -83,54 +82,6 @@ def _clear_invocation_state(state: BaseModel) -> None: value.clear() -PRIOR_TURN_NOTICE = ( - "[Evidence retrieved for an earlier question, no longer shown. It does not " - "count as cited for the current question.]" -) - - -def _compact_old_tool_returns( - messages: list[ModelMessage], - tool_names: frozenset[str], - *, - turn_start: int, -) -> list[ModelMessage]: - """Remove bulky earlier-question evidence while retaining the current one. - - Tool call and return parts remain paired; only the old return payload is - replaced. - - Page images attached to a replaced return are deliberately left in place. - Dropping them alongside their text is tempting — they are the bulk, and - they accumulate — but a follow-up about a figure already shown ("what - colour is that box?") carries no terms that could retrieve it again, so - removing the image turns an answerable question into a refusal. - - ``turn_start`` is how many messages existed when the current question - arrived, so everything below it belongs to an earlier one. The run reports - it rather than this function deriving it from message shape: a - ``UserPromptPart`` mid-question is as likely to be page images on a tool - return, or a notice a capability injected, and reading either as the next - question strips evidence the model is still answering from. - """ - if turn_start <= 0: - return messages - - compacted = list(messages) - for index, message in enumerate(messages[:turn_start]): - if not isinstance(message, ModelRequest): - continue - parts = [ - replace(part, content=PRIOR_TURN_NOTICE) - if isinstance(part, ToolReturnPart) and part.tool_name in tool_names - else part - for part in message.parts - ] - if parts != message.parts: - compacted[index] = replace(message, parts=parts) - return compacted - - def _is_resumption(prompt: Any, messages: list[ModelMessage]) -> bool: """Whether this run continues a question rather than asking a new one. @@ -185,17 +136,10 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): search_count: int = field(default=0, repr=False) request_count: int = field(default=0, repr=False) grace_requests_used: int = field(default=0, repr=False) - turn_start: int = field(default=0, repr=False) epoch: int = field(default=0, repr=False) async def for_run(self, ctx: RunContext[Any]) -> "RAGCapabilityBase[StateT]": - """Start a run's own copy, and decide what counts as an earlier question. - - On a resumption (see ``_is_resumption``) compaction is switched off for the - whole run: ``len(ctx.messages)`` would count the live question's own - messages and replace its evidence with the earlier-question notice, leaving - the model to answer with the evidence taken away. Failing this way costs a - larger request; failing the other way costs the answer. + """Start a run's own copy, and settle which question it is answering. A new question takes the message count as its identity, which every participant derives identically from the same history. A resumption keeps @@ -235,7 +179,6 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): request_count=0, grace_requests_used=0, epoch=0, - turn_start=0 if resuming else len(ctx.messages), ) run_capability._sync_state() return run_capability @@ -245,26 +188,6 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): return f"{self.config.prompts.domain_preamble}\n\n{self.instruction_text}" return self.instruction_text - async def wrap_model_request( - self, - ctx: RunContext[Any], - *, - request_context: ModelRequestContext, - handler: WrapModelRequestHandler, - ) -> ModelResponse: - """Trim earlier-question evidence off the wire only. - - Deliberately not ``before_model_request``: that hook's result is - assigned back onto the run's message history, so trimming there would - destroy the host's record of what was retrieved. - """ - request_context.messages = _compact_old_tool_returns( - request_context.messages, - self.tool_names - {self._cite_tool_name}, - turn_start=self.turn_start, - ) - return await handler(request_context) - async def before_model_request( self, ctx: RunContext[Any], request_context: ModelRequestContext ) -> ModelRequestContext: @@ -406,6 +329,18 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): self.rag = rag return self.rag + async def get_picture_bytes(self, document_id: str, self_ref: str) -> bytes | None: + """Fetch a picture of this capability's evidence, for whoever re-attaches it. + + Public because compaction rehydrates cited pictures and this capability + already holds the connection they came from; bytes are never kept in state. + """ + async with self.rag_lock: + rag = await self._ensure_rag() + return await rag.document_item_repository.get_picture_bytes( + document_id, self_ref + ) + async def _close(self) -> None: if self.rag is not None: await self.rag.__aexit__(None, None, None) diff --git a/haiku_rag_slim/haiku/rag/capabilities/compaction.py b/haiku_rag_slim/haiku/rag/capabilities/compaction.py index 7179ba27..e9b9e805 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/compaction.py +++ b/haiku_rag_slim/haiku/rag/capabilities/compaction.py @@ -1,13 +1,23 @@ from collections.abc import Mapping, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, field, replace from typing import Any, cast from pydantic_ai import RunContext -from pydantic_ai.capabilities import AbstractCapability +from pydantic_ai.capabilities import AbstractCapability, WrapModelRequestHandler +from pydantic_ai.messages import ( + BinaryContent, + ModelMessage, + ModelRequest, + ModelResponse, + ToolReturnPart, + UserPromptPart, +) +from pydantic_ai.models import ModelRequestContext from haiku.rag.capabilities._base import RAGCapabilityBase from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord from haiku.rag.store.models.citation import Citation +from haiku.rag.tools.search import RETRIEVED_IMAGE_TAG, decode_picture CAPABILITY_ID = "haiku-rag-evidence-compaction" @@ -38,7 +48,7 @@ def group_label(position: int) -> str: def picture_label(chunk_id: str, self_ref: str) -> str: return ( f"Page image retrieved from the knowledge base for cited evidence " - f"[{chunk_id}] ({self_ref}). Not provided by the user." + f"[{chunk_id}] ({self_ref}). Not provided by the user. {RETRIEVED_IMAGE_TAG}" ) @@ -181,6 +191,102 @@ def build_capsule(evidence: Sequence[DiscoveredEvidence]) -> Capsule: return Capsule(text=ENTRY_SEPARATOR.join(lines), pictures=tuple(pictures)) +def _strip_our_pictures(part: UserPromptPart) -> UserPromptPart | None: + """Drop the pictures we attached, together with the labels describing them. + + Ours is a label carrying the machine tag immediately followed by an image — + both halves required. Position alone is not ownership, since several tools' + results can arrive in one request; prose alone is not either, because a user can + write any phrase, and treating one as proof removed a user's own picture along + with their text. A label is only ever dropped with its picture: left behind it + would tell the model a figure is present when it is gone. + """ + if isinstance(part.content, str): + return part + items = list(part.content) + kept: list[Any] = [] + index = 0 + while index < len(items): + item = items[index] + following = items[index + 1] if index + 1 < len(items) else None + is_ours = ( + isinstance(item, str) + and RETRIEVED_IMAGE_TAG in item + and isinstance(following, BinaryContent) + ) + if is_ours: + index += 2 + continue + kept.append(item) + index += 1 + return replace(part, content=kept) if kept else None + + +def compact_history( + messages: list[ModelMessage], + *, + boundary: int, + owned_tools: frozenset[str], + capsule_text: str, + capsule_images: Sequence[str | BinaryContent] = (), +) -> list[ModelMessage]: + """Replace earlier questions' evidence with the capsule, on a copy. + + ``boundary`` is how many messages existed when the current question arrived, so + everything below it belongs to an earlier one. It comes from the recorded + question identity rather than from message shape: mid-question a user-role part + is as likely to be page images or an injected notice, and reading either as the + next question strips evidence the model is still answering from. + + The newest earlier return carries the capsule and every other becomes a receipt, + so exactly one capsule exists by construction. Returns are never removed, only + rewritten, which keeps each one paired with its call. Nothing outside this + capability's evidence tools is touched — not a cite acknowledgement, not another + capability's output, not a picture the user attached. + """ + if boundary <= 0: + return messages + + carrier = _newest_owned_return(messages, boundary, owned_tools) + compacted = list(messages) + for index, message in enumerate(messages[:boundary]): + if not isinstance(message, ModelRequest): + continue + parts: list[Any] = [] + for part in message.parts: + if isinstance(part, ToolReturnPart) and part.tool_name in owned_tools: + body = capsule_text or RECEIPT if index == carrier else RECEIPT + parts.append(replace(part, content=body)) + elif isinstance(part, UserPromptPart): + if (kept := _strip_our_pictures(part)) is not None: + parts.append(kept) + else: + parts.append(part) + if index == carrier and capsule_images: + parts.append(UserPromptPart(content=list(capsule_images))) + if not parts: + # A request with no parts is not a message; whatever emptied it was not + # ours to remove after all. + continue + if parts != message.parts: + compacted[index] = replace(message, parts=parts) + return compacted + + +def _newest_owned_return( + messages: list[ModelMessage], boundary: int, owned_tools: frozenset[str] +) -> int | None: + """Index of the last earlier request holding one of our evidence returns.""" + for index in range(min(boundary, len(messages)) - 1, -1, -1): + message = messages[index] + if isinstance(message, ModelRequest) and any( + isinstance(part, ToolReturnPart) and part.tool_name in owned_tools + for part in message.parts + ): + return index + return None + + @dataclass class EvidenceCompactionCapability(AbstractCapability[Any]): """Rewrites the history from what the evidence capabilities recorded. @@ -194,6 +300,94 @@ class EvidenceCompactionCapability(AbstractCapability[Any]): would share this capability's id. """ + built_for: tuple[str | None, int] | None = field(default=None, repr=False) + capsule: Capsule = field(default_factory=Capsule, repr=False) + images: tuple[str | BinaryContent, ...] = field(default=(), repr=False) + + async def for_run(self, ctx: RunContext[Any]) -> "EvidenceCompactionCapability": + """Give the run its own build cache, so concurrent runs cannot share one.""" + return replace(self, built_for=None, capsule=Capsule(), images=()) + + async def wrap_model_request( + self, + ctx: RunContext[Any], + *, + request_context: ModelRequestContext, + handler: WrapModelRequestHandler, + ) -> ModelResponse: + """Rewrite the request, never the stored history. + + Deliberately not ``before_model_request``: that hook's result is assigned + back onto the run's message history, which would destroy the host's record + of what was retrieved and break the message counts that question identities + and epochs are derived from. + """ + evidence = self.discover(ctx) + boundary = max((found.record.question or 0 for found in evidence), default=0) + if boundary > 0: + await self._build_once(ctx, evidence) + request_context.messages = compact_history( + request_context.messages, + boundary=boundary, + owned_tools=frozenset().union( + *(found.tool_names for found in evidence) + ), + capsule_text=self.capsule.text, + capsule_images=self.images, + ) + return await handler(request_context) + + async def _build_once( + self, ctx: RunContext[Any], evidence: Sequence[DiscoveredEvidence] + ) -> None: + """Build the capsule once per model request, however often the hook runs. + + Keyed on the run and its step rather than persisted: a stored key would + freeze one question's capsule across the next. + """ + key = (ctx.run_id, ctx.run_step) + if key == self.built_for: + return + self.capsule = build_capsule(evidence) + self.images = await self._rehydrate(ctx) + self.built_for = key + + async def _rehydrate(self, ctx: RunContext[Any]) -> tuple[str | BinaryContent, ...]: + """Fetch the cited pictures through the capability that retrieved them. + + Bytes are never stored in state, and the owner already holds an open + connection. A picture that cannot be fetched, for any reason, or that will + not decode, is emitted with neither its image nor its label: a label can + never outlive what it describes, and a figure the model has already been + given in text is not worth failing a question over. + """ + owners = { + capability.state_namespace: capability + for capability in ctx.capabilities.values() + if isinstance(capability, RAGCapabilityBase) + } + content: list[str | BinaryContent] = [] + for retained in self.capsule.pictures: + # Indexed, not looked up defensively: the capsule was built from these + # same capabilities in this same call, so a missing owner is a broken + # invariant rather than a picture to skip. + owner = owners[retained.capability] + try: + data = await owner.get_picture_bytes( + retained.document_id, retained.self_ref + ) + except Exception: + # A read that fails costs this picture, not the answer. + continue + if data is None: + continue + picture = decode_picture(data, retained.self_ref) + if picture is None: + continue + content.append(retained.label) + content.append(picture) + return tuple(content) + def discover(self, ctx: RunContext[Any]) -> list[DiscoveredEvidence]: """Read what each evidence capability recorded, without writing anything. @@ -237,6 +431,7 @@ __all__ = [ "EvidenceCompactionCapability", "RetainedPicture", "build_capsule", + "compact_history", "create_capability", "group_label", "picture_label", diff --git a/haiku_rag_slim/haiku/rag/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py index 0350e5d4..6aa42799 100644 --- a/haiku_rag_slim/haiku/rag/chat/app.py +++ b/haiku_rag_slim/haiku/rag/chat/app.py @@ -25,6 +25,7 @@ from textual.worker import Worker from haiku.rag.capabilities._base import RAGCapabilityBase from haiku.rag.capabilities.analysis import AnalysisState +from haiku.rag.capabilities.compaction import create_capability as create_compaction from haiku.rag.capabilities.rag import AGENT_PREAMBLE, RAGState from haiku.rag.chat.widgets.chat_history import ChatHistory, CitationWidget from haiku.rag.chat.widgets.image_select import ImageAdded @@ -153,7 +154,9 @@ class ChatApp(App): self._model, deps_type=ChatDeps, instructions=AGENT_PREAMBLE, - capabilities=self._capabilities, + # A chat is multi-turn by definition, so earlier questions are reduced + # to the evidence they cited rather than carried whole. + capabilities=[*self._capabilities, create_compaction()], ) self._state = {} for capability in self._capabilities: diff --git a/haiku_rag_slim/haiku/rag/tools/search.py b/haiku_rag_slim/haiku/rag/tools/search.py index 787210b6..4f192a74 100644 --- a/haiku_rag_slim/haiku/rag/tools/search.py +++ b/haiku_rag_slim/haiku/rag/tools/search.py @@ -10,6 +10,30 @@ from haiku.rag.config.models import AppConfig from haiku.rag.store.models import SearchResult from haiku.rag.tools.context import RAGDeps +RETRIEVED_IMAGE_TAG = "[haiku.rag/retrieved-image]" +"""Tag every label we attach to a retrieved picture ends with. + +Identifies our own pictures on the wire without inferring ownership from position, +which is wrong as soon as two tools' results arrive in one request. Deliberately not +a phrase: a user writing "retrieved from the knowledge base for my report" above +their own picture had it removed, along with their text. +""" + + +def decode_picture(data: bytes, self_ref: str) -> BinaryContent | None: + """Wrap picture bytes for the wire, or return nothing if they will not decode. + + The model adapter renders one vision placeholder per ``BinaryContent``, so + emitting one for an image the server cannot decode leaves the processor with an + off-by-one count. + """ + try: + with Image.open(BytesIO(data)) as image: + image.verify() + except Exception: + return None + return BinaryContent(data=data, media_type="image/png", identifier=self_ref) + def build_image_content_from_results( results: list[SearchResult], @@ -33,7 +57,7 @@ def build_image_content_from_results( ``BinaryContent.identifier`` cannot do — it does not survive serialization to the vision API. """ - collected: list[tuple[str | None, str, bytes]] = [] + collected: list[tuple[str | None, str, BinaryContent]] = [] seen: set[tuple[str | None, str]] = set() for result in results: if not result.image_data: @@ -42,26 +66,21 @@ def build_image_content_from_results( key = (result.document_id, self_ref) if key in seen: continue - data = base64.b64decode(b64) - try: - with Image.open(BytesIO(data)) as img: - img.verify() - except Exception: + picture = decode_picture(base64.b64decode(b64), self_ref) + if picture is None: continue - collected.append((result.chunk_id, self_ref, data)) + collected.append((result.chunk_id, self_ref, picture)) seen.add(key) content: list[str | BinaryContent] = [] total = len(collected) - for position, (chunk_id, self_ref, data) in enumerate(collected, 1): + for position, (chunk_id, self_ref, picture) in enumerate(collected, 1): content.append( - f"Page image {position} of {total}, retrieved from the knowledge " - f"base for search result [{chunk_id}] ({self_ref}). " - "Not provided by the user." - ) - content.append( - BinaryContent(data=data, media_type="image/png", identifier=self_ref) + f"Page image {position} of {total}, retrieved from the knowledge base " + f"for search result [{chunk_id}] ({self_ref}). " + f"Not provided by the user. {RETRIEVED_IMAGE_TAG}" ) + content.append(picture) return content diff --git a/tests/capabilities/test_capabilities.py b/tests/capabilities/test_capabilities.py index 2e140f6f..f4f7943f 100644 --- a/tests/capabilities/test_capabilities.py +++ b/tests/capabilities/test_capabilities.py @@ -6,7 +6,6 @@ from unittest.mock import AsyncMock, patch import pytest from pydantic_ai import Agent, DeferredToolResults, ModelRetry, RunContext, ToolFailed from pydantic_ai.messages import ( - BinaryContent, ModelRequest, ModelResponse, TextPart, @@ -20,9 +19,7 @@ from pydantic_ai.usage import RunUsage from haiku.rag.capabilities._base import ( CITATION_GRACE_REQUESTS, - PRIOR_TURN_NOTICE, _called_own_tool, - _compact_old_tool_returns, ) from haiku.rag.capabilities.analysis import AnalysisCapability, AnalysisState from haiku.rag.capabilities.analysis import create_capability as create_analysis @@ -848,182 +845,6 @@ async def test_deferred_capability_loads_native_tools(temp_db_path): assert "rag_search" in loaded_payloads[0] -def test_prior_turn_tool_results_are_compacted_but_current_evidence_is_kept(): - messages = [ - ModelRequest(parts=[UserPromptPart("old question")]), - ModelResponse(parts=[ToolCallPart("rag_search", {}, "old-call")]), - ModelRequest( - parts=[ToolReturnPart("rag_search", "large old evidence", "old-call")] - ), - ModelRequest(parts=[UserPromptPart("current question")]), - ModelResponse(parts=[ToolCallPart("rag_search", {}, "current-call")]), - ModelRequest( - parts=[ToolReturnPart("rag_search", "current evidence", "current-call")] - ), - ] - - compacted = _compact_old_tool_returns( - messages, frozenset({"rag_search"}), turn_start=3 - ) - - old_return = compacted[2].parts[0] - current_return = compacted[5].parts[0] - assert isinstance(old_return, ToolReturnPart) - assert old_return.content == PRIOR_TURN_NOTICE - assert isinstance(current_return, ToolReturnPart) - assert current_return.content == "current evidence" - - -def test_nothing_is_compacted_on_the_first_question(): - messages = [ - ModelResponse(parts=[ToolCallPart("rag_search", {}, "current-call")]), - ModelRequest( - parts=[ToolReturnPart("rag_search", "current evidence", "current-call")] - ), - ] - - compacted = _compact_old_tool_returns( - messages, frozenset({"rag_search"}), turn_start=0 - ) - - assert compacted is messages - current_return = compacted[1].parts[0] - assert isinstance(current_return, ToolReturnPart) - assert current_return.content == "current evidence" - - -PAGE_IMAGE = BinaryContent(data=b"\x89PNG" + b"\x00" * 64, media_type="image/png") - - -def _search_exchange(call_id: str, evidence: str, *, images: bool): - """One search round-trip in the shape pydantic-ai produces. - - Images on a ``ToolReturn`` arrive as a separate ``UserPromptPart`` appended - to the same ``ModelRequest`` as the ``ToolReturnPart``. - """ - parts: list[Any] = [ToolReturnPart("rag_search", evidence, call_id)] - if images: - parts.append(UserPromptPart(content=[PAGE_IMAGE])) - return [ - ModelResponse(parts=[ToolCallPart("rag_search", {}, call_id)]), - ModelRequest(parts=parts), - ] - - -@pytest.mark.parametrize( - ("label", "trailing"), - [ - ("page images on a tool return", [UserPromptPart(content=[PAGE_IMAGE])]), - ("a notice injected mid-run", [UserPromptPart("You answered without citing")]), - ], -) -def test_current_turn_evidence_survives_later_user_prompt_parts(label, trailing): - """Nothing that arrives mid-question may be read as the next question. - - Page images on a tool return and a notice this capability injects both - appear as a ``UserPromptPart`` after the question, so deriving the turn from - message shape stripped evidence the model was still answering from. - """ - messages = [ - ModelRequest(parts=[UserPromptPart("current question")]), - *_search_exchange("first-call", "first evidence", images=False), - ModelRequest(parts=trailing), - ] - - compacted = _compact_old_tool_returns( - messages, frozenset({"rag_search"}), turn_start=0 - ) - - first_return = compacted[2].parts[0] - assert isinstance(first_return, ToolReturnPart) - assert first_return.content == "first evidence", label - - -def test_prior_turn_images_outlive_their_tool_return(): - """A follow-up about a figure cannot retrieve it again, so keep the image. - - "What colour is that box?" has no terms the search can use, so dropping the - image with its text turns an answerable question into a refusal. - """ - messages = [ - ModelRequest(parts=[UserPromptPart("old question")]), - *_search_exchange("old-call", "old evidence", images=True), - ModelResponse(parts=[TextPart("an answer")]), - ModelRequest(parts=[UserPromptPart("current question")]), - ] - - compacted = _compact_old_tool_returns( - messages, frozenset({"rag_search"}), turn_start=4 - ) - - old_return = compacted[2].parts[0] - assert isinstance(old_return, ToolReturnPart) - assert old_return.content == PRIOR_TURN_NOTICE - assert any( - isinstance(part, UserPromptPart) and not isinstance(part.content, str) - for message in compacted - for part in message.parts - ) - - -def test_user_attached_image_is_never_dropped(): - messages = [ - ModelRequest(parts=[UserPromptPart("old question")]), - *_search_exchange("old-call", "old evidence", images=False), - ModelResponse(parts=[TextPart("an answer")]), - ModelRequest(parts=[UserPromptPart(content=[PAGE_IMAGE, "what is this?"])]), - ] - - compacted = _compact_old_tool_returns( - messages, frozenset({"rag_search"}), turn_start=4 - ) - - old_return = compacted[2].parts[0] - assert isinstance(old_return, ToolReturnPart) - assert old_return.content == PRIOR_TURN_NOTICE - attached = compacted[-1].parts[0] - assert isinstance(attached, UserPromptPart) - assert attached.content == [PAGE_IMAGE, "what is this?"] - - -async def test_compaction_never_reaches_the_stored_message_history(temp_db_path): - """Trimming is for the wire; hosts keep the evidence they gathered.""" - capability = create_rag( - db_path=temp_db_path, config=AppConfig(), defer_loading=False - ) - turns = iter( - [ - ModelResponse(parts=[ToolCallPart("rag_search", {"query": "q"}, "call-1")]), - ModelResponse(parts=[TextPart("first answer")]), - ModelResponse(parts=[TextPart("second answer")]), - ] - ) - - async def model(_messages, _info): - return next(turns) - - agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[capability]) - deps = Deps(state={"rag": RAGState().model_dump(mode="json")}) - - with patch.object( - RAGCapability, "_search", AsyncMock(return_value="REAL EVIDENCE") - ): - first = await agent.run("old question", deps=deps) - second = await agent.run( - "current question", deps=deps, message_history=first.all_messages() - ) - - returns = [ - str(part.content) - for message in second.all_messages() - if isinstance(message, ModelRequest) - for part in message.parts - if isinstance(part, ToolReturnPart) - ] - assert "REAL EVIDENCE" in returns - assert PRIOR_TURN_NOTICE not in returns - - def _resuming_deps() -> Deps: """State as a resumption always finds it: the question already identified. @@ -1051,105 +872,9 @@ def _in_flight_history() -> list[Any]: ] -def _wire_returns(sent: list[Any]) -> list[str]: - return [ - str(part.content) - for message in sent - if isinstance(message, ModelRequest) - for part in message.parts - if isinstance(part, ToolReturnPart) - ] - - @pytest.mark.asyncio -async def test_a_new_question_compacts_the_previous_one(temp_db_path): - """The baseline the resume cases are contrasted against. - - A settled history — the previous question answered — is what a genuinely new - question follows. An unfinished tail is ambiguous instead, and treated as a - continuation. - """ - capability = create_rag( - db_path=temp_db_path, config=AppConfig(), defer_loading=False - ) - wire: list[list[Any]] = [] - - async def model(messages, _info): - wire.append(list(messages)) - return ModelResponse(parts=[TextPart("answer")]) - - agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[capability]) - settled = [*_in_flight_history(), ModelResponse(parts=[TextPart("first answer")])] - - await agent.run("a different question", deps=Deps(), message_history=settled) - - assert _wire_returns(wire[-1]) == [PRIOR_TURN_NOTICE] - - @pytest.mark.asyncio -async def test_a_prompt_on_an_unanswered_tail_is_treated_as_a_continuation( - temp_db_path, -): - """Ambiguous shape, resolved the safe way. - - A history ending in a request the model never answered, plus a new prompt, - could be a fresh question or a continuation. Compacting would cost the answer - if it is a continuation; not compacting only costs a larger request. - """ - capability = create_rag( - db_path=temp_db_path, config=AppConfig(), defer_loading=False - ) - wire: list[list[Any]] = [] - - async def model(messages, _info): - wire.append(list(messages)) - return ModelResponse(parts=[TextPart("answer")]) - - agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[capability]) - - await agent.run( - "a different question", - deps=_resuming_deps(), - message_history=_in_flight_history(), - ) - - assert _wire_returns(wire[-1]) == ["EVIDENCE FOR THE LIVE TURN"] - - @pytest.mark.asyncio -async def test_a_resume_carrying_a_prompt_keeps_the_active_evidence(temp_db_path): - """Deferred results may arrive with a prompt, and that is still a continuation. - - ``ctx.prompt`` is non-null here, so the absence of a prompt cannot be the only - signal: the history tail is unfinished — a response whose tool call has no - return yet — and the question it belongs to is still being answered. - """ - capability = create_rag( - db_path=temp_db_path, config=AppConfig(), defer_loading=False - ) - wire: list[list[Any]] = [] - - async def model(messages, _info): - wire.append(list(messages)) - return ModelResponse(parts=[TextPart("answer")]) - - agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[capability]) - history = [ - *_in_flight_history(), - ModelResponse(parts=[ToolCallPart("external_tool", {}, "call-2")]), - ] - - await agent.run( - "carry on", - deferred_tool_results=DeferredToolResults(calls={"call-2": "external result"}), - message_history=history, - deps=_resuming_deps(), - ) - - assert "EVIDENCE FOR THE LIVE TURN" in _wire_returns(wire[-1]) - assert PRIOR_TURN_NOTICE not in _wire_returns(wire[-1]) - - @pytest.mark.asyncio @pytest.mark.parametrize( "resume_kwargs", @@ -1160,88 +885,7 @@ async def test_a_resume_carrying_a_prompt_keeps_the_active_evidence(temp_db_path ), ], ) -async def test_a_resumed_run_keeps_the_active_questions_evidence( - temp_db_path, resume_kwargs -): - """A run without a prompt continues a question; nothing in it is prior. - - ``len(ctx.messages)`` cannot tell the two apart — on a resumption it counts - the live question's own messages and marks its evidence as earlier-question - evidence, leaving the model to answer with a notice where its search result - used to be. Deferred, interrupted and suspended resumes all differ in shape, - so the absence of a prompt is the signal rather than the message layout. - """ - capability = create_rag( - db_path=temp_db_path, config=AppConfig(), defer_loading=False - ) - wire: list[list[Any]] = [] - - async def model(messages, _info): - wire.append(list(messages)) - return ModelResponse(parts=[TextPart("answer")]) - - agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[capability]) - - await agent.run( - deps=_resuming_deps(), message_history=_in_flight_history(), **resume_kwargs - ) - - assert _wire_returns(wire[-1]) == ["EVIDENCE FOR THE LIVE TURN"] - - @pytest.mark.asyncio -async def test_prior_turn_search_is_compacted_but_its_cite_receipt_is_not(temp_db_path): - """Only evidence is compacted, and the boundary comes from the run. - - A cite acknowledgement is a receipt, not evidence: replacing it lengthened - the request and erased the record that citations had been registered. - """ - capability = create_rag( - db_path=temp_db_path, config=AppConfig(), defer_loading=False - ) - turns = iter( - [ - ModelResponse(parts=[ToolCallPart("rag_search", {"query": "q"}, "call-1")]), - ModelResponse( - parts=[ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-2")] - ), - ModelResponse(parts=[TextPart("first answer")]), - ModelResponse(parts=[TextPart("second answer")]), - ] - ) - wire: list[list[Any]] = [] - - async def model(messages, _info): - wire.append(list(messages)) - return next(turns) - - agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[capability]) - deps = Deps(state={"rag": RAGState().model_dump(mode="json")}) - - async def search(self, query: str, _limit: int | None) -> str: - """Record a result the way the real search does, so citing resolves.""" - cast(Any, self.state).searches[query] = [ - SearchResult(content="evidence", score=1.0, chunk_id="chunk-1") - ] - return "REAL EVIDENCE" - - with patch.object(RAGCapability, "_search", search): - first = await agent.run("old question", deps=deps) - await agent.run( - "current question", deps=deps, message_history=first.all_messages() - ) - - prior = { - part.tool_name: str(part.content) - for message in wire[-1] - if isinstance(message, ModelRequest) - for part in message.parts - if isinstance(part, ToolReturnPart) - } - assert prior["rag_search"] == PRIOR_TURN_NOTICE - assert prior["rag_cite"] != PRIOR_TURN_NOTICE - - def _record(deps: Deps, namespace: str) -> CapabilityEvidenceRecord: return CapabilityEvidenceRecord.model_validate(deps.state[namespace]["evidence"]) @@ -1596,3 +1240,19 @@ async def test_a_resumption_keeps_the_evidence_the_question_already_gathered( assert record.question == identity assert record.occurrences["chunk-1"].retrieved_in_questions == [identity] assert citation_status([record], question=identity) == "grounded" + + +@pytest.mark.asyncio +async def test_a_capability_fetches_its_own_evidences_pictures(temp_db_path): + """Compaction rehydrates through the owner, which already holds the connection.""" + capability = create_rag(db_path=temp_db_path, config=AppConfig()) + client = AsyncMock() + client.document_item_repository.get_picture_bytes.return_value = b"picture-bytes" + capability.rag = client + + data = await capability.get_picture_bytes("doc-1", "#/pictures/0") + + assert data == b"picture-bytes" + client.document_item_repository.get_picture_bytes.assert_awaited_once_with( + "doc-1", "#/pictures/0" + ) diff --git a/tests/capabilities/test_evidence_wire.py b/tests/capabilities/test_evidence_wire.py new file mode 100644 index 00000000..bcce85f5 --- /dev/null +++ b/tests/capabilities/test_evidence_wire.py @@ -0,0 +1,655 @@ +import base64 +from dataclasses import dataclass, field, replace +from typing import Any, cast +from unittest.mock import AsyncMock, patch + +import pytest +from pydantic_ai import Agent, RunContext +from pydantic_ai.messages import ( + BinaryContent, + ModelRequest, + ModelResponse, + TextPart, + ToolCallPart, + ToolReturnPart, + UserPromptPart, +) +from pydantic_ai.models import ModelRequestContext, ModelRequestParameters +from pydantic_ai.models.function import FunctionModel +from pydantic_ai.models.test import TestModel +from pydantic_ai.usage import RunUsage + +from haiku.rag.capabilities.compaction import ( + RECEIPT, + Capsule, + compact_history, + picture_label, +) +from haiku.rag.capabilities.compaction import create_capability as create_compaction +from haiku.rag.capabilities.ledger import CapabilityEvidenceRecord +from haiku.rag.capabilities.rag import RAGCapability, RAGState +from haiku.rag.capabilities.rag import create_capability as create_rag +from haiku.rag.config.models import AppConfig +from haiku.rag.store.models.chunk import SearchResult + +OWNED = frozenset({"rag_search"}) +PNG = BinaryContent(data=b"fake-image-bytes", media_type="image/png") + + +def retrieved_image(chunk_id: str = "chunk-1", self_ref: str = "#/pictures/0"): + """A page image on the wire, labelled the way a search result attaches it.""" + return UserPromptPart(content=[picture_label(chunk_id, self_ref), PNG]) + + +def answered_question(question: str, *, evidence: str, images: bool = False): + """One settled question: prompt, search, result, answer.""" + returned: list[Any] = [ToolReturnPart("rag_search", evidence, "call-1")] + if images: + returned.append(retrieved_image()) + return [ + ModelRequest(parts=[UserPromptPart(question)]), + ModelResponse(parts=[ToolCallPart("rag_search", {"query": "q"}, "call-1")]), + ModelRequest(parts=returned), + ModelResponse(parts=[TextPart("an answer")]), + ] + + +def returns_of(messages) -> list[str]: + return [ + str(part.content) + for message in messages + if isinstance(message, ModelRequest) + for part in message.parts + if isinstance(part, ToolReturnPart) + ] + + +def images_of(messages) -> list[BinaryContent]: + return [ + item + for message in messages + if isinstance(message, ModelRequest) + for part in message.parts + if isinstance(part, UserPromptPart) and not isinstance(part.content, str) + for item in part.content + if isinstance(item, BinaryContent) + ] + + +def texts_of(messages) -> list[str]: + return [ + item + for message in messages + if isinstance(message, ModelRequest) + for part in message.parts + if isinstance(part, UserPromptPart) and not isinstance(part.content, str) + for item in part.content + if isinstance(item, str) + ] + + +def test_nothing_before_the_first_question_is_compacted(): + messages = answered_question("first", evidence="EVIDENCE") + + compacted = compact_history( + messages, boundary=0, owned_tools=OWNED, capsule_text="CAPSULE" + ) + + assert compacted == messages + + +def test_the_newest_earlier_return_carries_the_capsule(): + messages = [ + *answered_question("first", evidence="OLD EVIDENCE"), + ModelRequest(parts=[UserPromptPart("second")]), + ] + + compacted = compact_history( + messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE" + ) + + assert returns_of(compacted) == ["CAPSULE"] + assert "OLD EVIDENCE" not in returns_of(compacted) + + +def test_older_returns_become_receipts_and_only_the_newest_carries_the_capsule(): + messages = [ + *answered_question("first", evidence="OLDEST"), + *answered_question("second", evidence="NEWER"), + ModelRequest(parts=[UserPromptPart("third")]), + ] + + compacted = compact_history( + messages, boundary=8, owned_tools=OWNED, capsule_text="CAPSULE" + ) + + assert returns_of(compacted) == [RECEIPT, "CAPSULE"] + + +def test_the_current_question_keeps_its_own_evidence(): + messages = [ + *answered_question("first", evidence="OLD EVIDENCE"), + ModelRequest(parts=[UserPromptPart("second")]), + ModelResponse(parts=[ToolCallPart("rag_search", {"query": "q"}, "call-2")]), + ModelRequest(parts=[ToolReturnPart("rag_search", "LIVE EVIDENCE", "call-2")]), + ] + + compacted = compact_history( + messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE" + ) + + assert returns_of(compacted) == ["CAPSULE", "LIVE EVIDENCE"] + + +def test_another_capabilitys_return_is_left_alone(): + messages = [ + ModelRequest(parts=[UserPromptPart("first")]), + ModelResponse(parts=[ToolCallPart("other_tool", {}, "call-1")]), + ModelRequest(parts=[ToolReturnPart("other_tool", "NOT OURS", "call-1")]), + ModelResponse(parts=[TextPart("an answer")]), + ModelRequest(parts=[UserPromptPart("second")]), + ] + + compacted = compact_history( + messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE" + ) + + assert returns_of(compacted) == ["NOT OURS"] + + +def test_a_cite_acknowledgement_survives(): + """A receipt of the model's own action, not evidence.""" + messages = [ + ModelRequest(parts=[UserPromptPart("first")]), + ModelResponse( + parts=[ToolCallPart("rag_cite", {"chunk_ids": ["c1"]}, "call-1")] + ), + ModelRequest( + parts=[ToolReturnPart("rag_cite", "Registered 1 citation.", "c1")] + ), + ModelResponse(parts=[TextPart("an answer")]), + ModelRequest(parts=[UserPromptPart("second")]), + ] + + compacted = compact_history( + messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE" + ) + + assert returns_of(compacted) == ["Registered 1 citation."] + + +def test_an_uncited_earlier_image_is_dropped_with_its_label(): + messages = [ + *answered_question("first", evidence="OLD", images=True), + ModelRequest(parts=[UserPromptPart("second")]), + ] + + compacted = compact_history( + messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE" + ) + + assert images_of(compacted) == [] + assert texts_of(compacted) == [] + + +def test_cited_pictures_are_attached_beside_the_capsule(): + messages = [ + *answered_question("first", evidence="OLD", images=True), + ModelRequest(parts=[UserPromptPart("second")]), + ] + fresh = BinaryContent(data=b"cited-bytes", media_type="image/png") + + compacted = compact_history( + messages, + boundary=4, + owned_tools=OWNED, + capsule_text="CAPSULE", + capsule_images=[picture_label("cited-chunk", "#/pictures/3"), fresh], + ) + + assert images_of(compacted) == [fresh] + assert texts_of(compacted) == [picture_label("cited-chunk", "#/pictures/3")] + carrier = [ + index + for index, message in enumerate(compacted) + if isinstance(message, ModelRequest) + and any( + isinstance(part, ToolReturnPart) and part.content == "CAPSULE" + for part in message.parts + ) + ] + attached = [ + index + for index, message in enumerate(compacted) + if isinstance(message, ModelRequest) + and any( + isinstance(part, UserPromptPart) and not isinstance(part.content, str) + for part in message.parts + ) + ] + assert carrier == attached + + +def test_a_user_attached_image_is_never_dropped(): + """The user's own picture is not ours to remove, even in an earlier question.""" + mine = UserPromptPart(content=["look at this", PNG]) + messages = [ + ModelRequest(parts=[mine]), + ModelResponse(parts=[ToolCallPart("rag_search", {"query": "q"}, "call-1")]), + ModelRequest(parts=[ToolReturnPart("rag_search", "OLD", "call-1")]), + ModelResponse(parts=[TextPart("an answer")]), + ModelRequest(parts=[UserPromptPart("second")]), + ] + + compacted = compact_history( + messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE" + ) + + assert images_of(compacted) == [PNG] + assert texts_of(compacted) == ["look at this"] + + +def test_the_stored_messages_are_never_mutated(): + messages = [ + *answered_question("first", evidence="OLD EVIDENCE", images=True), + ModelRequest(parts=[UserPromptPart("second")]), + ] + before = [list(message.parts) for message in messages] + + compact_history(messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE") + + assert [list(message.parts) for message in messages] == before + assert "OLD EVIDENCE" in returns_of(messages) + + +def test_nothing_cited_leaves_only_receipts(): + messages = [ + *answered_question("first", evidence="OLD"), + ModelRequest(parts=[UserPromptPart("second")]), + ] + + compacted = compact_history( + messages, boundary=4, owned_tools=OWNED, capsule_text="" + ) + + assert returns_of(compacted) == [RECEIPT] + + +@dataclass +class Deps: + state: dict[str, Any] = field(default_factory=dict) + + +def rag_and_compactor(temp_db_path): + return ( + create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False), + create_compaction(), + ) + + +def in_flight_history() -> list[Any]: + """A question already asked and searched, still awaiting its answer.""" + return [ + ModelRequest(parts=[UserPromptPart("what does the supervisor do?")]), + ModelResponse(parts=[ToolCallPart("rag_search", {"query": "s"}, "call-1")]), + ModelRequest( + parts=[ToolReturnPart("rag_search", "EVIDENCE FOR THE LIVE TURN", "call-1")] + ), + ] + + +def resuming_deps(question: int = 0) -> Deps: + """State as a resumption always finds it: the question already identified.""" + return Deps( + state={ + "rag": RAGState( + evidence=CapabilityEvidenceRecord(question=question) + ).model_dump(mode="json") + } + ) + + +@pytest.mark.asyncio +async def test_without_the_compactor_the_history_is_untouched(temp_db_path): + """Omission is the switch: there is no flag to test, only absence.""" + rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False) + wire: list[list[Any]] = [] + + async def model(messages, _info): + wire.append(list(messages)) + return ModelResponse(parts=[TextPart("answer")]) + + agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag]) + settled = [*in_flight_history(), ModelResponse(parts=[TextPart("first answer")])] + + await agent.run("a different question", deps=Deps(), message_history=settled) + + assert returns_of(wire[-1]) == ["EVIDENCE FOR THE LIVE TURN"] + + +@pytest.mark.asyncio +async def test_with_the_compactor_a_new_question_compacts_the_previous_one( + temp_db_path, +): + rag, compactor = rag_and_compactor(temp_db_path) + wire: list[list[Any]] = [] + + async def model(messages, _info): + wire.append(list(messages)) + return ModelResponse(parts=[TextPart("answer")]) + + agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag, compactor]) + settled = [*in_flight_history(), ModelResponse(parts=[TextPart("first answer")])] + + await agent.run("a different question", deps=Deps(), message_history=settled) + + assert returns_of(wire[-1]) == [RECEIPT] + + +@pytest.mark.asyncio +async def test_a_resumed_question_keeps_the_evidence_it_is_answering_from( + temp_db_path, +): + """The boundary is the stored identity of the question in progress. + + An earlier question below it is compacted; the evidence the model is still + answering from sits above it and survives. Deriving the boundary from message + shape instead would put the live evidence below it and answer with a receipt + where the search result should be. + """ + rag, compactor = rag_and_compactor(temp_db_path) + wire: list[list[Any]] = [] + + async def model(messages, _info): + wire.append(list(messages)) + return ModelResponse(parts=[TextPart("answer")]) + + agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag, compactor]) + history = [ + *answered_question("an earlier question", evidence="EVIDENCE FOR THE OLD TURN"), + *in_flight_history(), + ] + + await agent.run(deps=resuming_deps(question=4), message_history=history) + + assert returns_of(wire[-1]) == [RECEIPT, "EVIDENCE FOR THE LIVE TURN"] + + +@pytest.mark.asyncio +async def test_compaction_never_reaches_the_stored_message_history(temp_db_path): + """Rewriting is for the wire; hosts keep the evidence they gathered.""" + rag, compactor = rag_and_compactor(temp_db_path) + turns = iter( + [ + ModelResponse(parts=[ToolCallPart("rag_search", {"query": "q"}, "call-1")]), + ModelResponse(parts=[TextPart("first answer")]), + ModelResponse(parts=[TextPart("second answer")]), + ] + ) + + async def model(_messages, _info): + return next(turns) + + agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag, compactor]) + deps = Deps() + + with patch.object( + RAGCapability, "_search", AsyncMock(return_value="REAL EVIDENCE") + ): + first = await agent.run("old question", deps=deps) + second = await agent.run( + "current question", deps=deps, message_history=first.all_messages() + ) + + assert "REAL EVIDENCE" in returns_of(second.all_messages()) + assert RECEIPT not in returns_of(second.all_messages()) + + +REAL_PNG = base64.b64decode( + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8DwHwAFBQIAX8jx0gAAAABJRU5ErkJggg==" +) + + +async def _search_with_a_picture(self, query: str, _limit: int | None) -> str: + """Record a result carrying a page image, the way a real search does.""" + cast(Any, self.state).searches[query] = [ + SearchResult( + content="evidence", + score=1.0, + chunk_id="chunk-1", + document_id="doc-1", + doc_item_refs=["#/pictures/0"], + ) + ] + self._note_evidence() + return "EVIDENCE" + + +async def _cite_a_picture_chunk(temp_db_path, fetched: bytes | None): + """Two questions: cite a picture chunk, then ask something else.""" + rag, compactor = rag_and_compactor(temp_db_path) + calls = iter( + [ + [ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")], + [ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-2")], + [TextPart("first answer")], + [TextPart("second answer")], + ] + ) + wire: list[list[Any]] = [] + + async def model(messages, _info): + wire.append(list(messages)) + return ModelResponse(parts=next(calls)) + + agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag, compactor]) + deps = Deps() + + with ( + patch.object(RAGCapability, "_search", _search_with_a_picture), + patch.object( + RAGCapability, "get_picture_bytes", AsyncMock(return_value=fetched) + ), + ): + first = await agent.run("what does the supervisor do?", deps=deps) + await agent.run( + "and what else?", deps=deps, message_history=first.all_messages() + ) + return wire + + +@pytest.mark.asyncio +async def test_a_cited_picture_is_fetched_and_attached_with_its_label(temp_db_path): + wire = await _cite_a_picture_chunk(temp_db_path, REAL_PNG) + + assert [picture.data for picture in images_of(wire[-1])] == [REAL_PNG] + assert texts_of(wire[-1]) == [picture_label("chunk-1", "#/pictures/0")] + + +@pytest.mark.asyncio +async def test_a_picture_that_cannot_be_fetched_emits_neither_image_nor_label( + temp_db_path, +): + """A label without its picture tells the model a figure is there when it is not.""" + wire = await _cite_a_picture_chunk(temp_db_path, None) + + assert images_of(wire[-1]) == [] + assert texts_of(wire[-1]) == [] + + +@pytest.mark.asyncio +async def test_a_picture_that_will_not_decode_emits_neither_image_nor_label( + temp_db_path, +): + """One vision placeholder is rendered per attachment, so a corrupt one miscounts.""" + wire = await _cite_a_picture_chunk(temp_db_path, b"not-an-image") + + assert images_of(wire[-1]) == [] + assert texts_of(wire[-1]) == [] + + +@pytest.mark.asyncio +async def test_the_capsule_is_built_once_per_request_and_again_for_the_next( + temp_db_path, +): + """Two hook passes for one request must not rebuild; the next request must.""" + rag, compactor = rag_and_compactor(temp_db_path) + deps = Deps() + ctx = RunContext( + deps=deps, model=TestModel(), usage=RunUsage(), run_id="run-1", run_step=1 + ) + run_rag = await rag.for_run(ctx) + run_compactor = await compactor.for_run(ctx) + cast(Any, run_rag.state).evidence.begin_question(4) + ctx = replace(ctx, capabilities={"rag": run_rag, "compaction": run_compactor}) + builds = 0 + + def counting_build(evidence): + nonlocal builds + builds += 1 + return Capsule(text="CAPSULE") + + async def handler(_request_context): + return ModelResponse(parts=[TextPart("answer")]) + + request = ModelRequestContext( + messages=[*answered_question("first", evidence="OLD")], + model=TestModel(), + model_request_parameters=ModelRequestParameters(), + model_settings=None, + ) + + with patch("haiku.rag.capabilities.compaction.build_capsule", counting_build): + await run_compactor.wrap_model_request( + ctx, request_context=request, handler=handler + ) + await run_compactor.wrap_model_request( + ctx, request_context=request, handler=handler + ) + assert builds == 1 + + await run_compactor.wrap_model_request( + replace(ctx, run_step=2), request_context=request, handler=handler + ) + + assert builds == 2 + + +def test_a_user_quoting_our_wording_keeps_their_picture_and_their_text(): + """Prose is not proof of ownership: a user can write any phrase. + + Recognising our own pictures by a natural-language substring removed a user's + image, its text, and with it the whole message part. + """ + quoted = UserPromptPart( + content=[ + "Here is a page image retrieved from the knowledge base for my report", + PNG, + ] + ) + messages = [ + ModelRequest(parts=[quoted]), + ModelResponse(parts=[ToolCallPart("rag_search", {"query": "q"}, "call-1")]), + ModelRequest(parts=[ToolReturnPart("rag_search", "OLD", "call-1")]), + ModelResponse(parts=[TextPart("an answer")]), + ModelRequest(parts=[UserPromptPart("second")]), + ] + + compacted = compact_history( + messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE" + ) + + assert images_of(compacted) == [PNG] + assert texts_of(compacted) == [ + "Here is a page image retrieved from the knowledge base for my report" + ] + assert all(message.parts for message in compacted) + + +def test_a_label_of_ours_with_no_picture_after_it_is_kept(): + """Only a genuine pair is ours to remove; a lone label is someone else's text.""" + lonely = UserPromptPart( + content=[picture_label("chunk-1", "#/pictures/0"), "and more"] + ) + messages = [ + ModelRequest(parts=[lonely]), + ModelResponse(parts=[ToolCallPart("rag_search", {"query": "q"}, "call-1")]), + ModelRequest(parts=[ToolReturnPart("rag_search", "OLD", "call-1")]), + ModelResponse(parts=[TextPart("an answer")]), + ModelRequest(parts=[UserPromptPart("second")]), + ] + + compacted = compact_history( + messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE" + ) + + assert texts_of(compacted) == [ + picture_label("chunk-1", "#/pictures/0"), + "and more", + ] + + +@pytest.mark.asyncio +async def test_a_picture_whose_fetch_raises_costs_the_picture_not_the_answer( + temp_db_path, +): + rag, compactor = rag_and_compactor(temp_db_path) + calls = iter( + [ + [ToolCallPart("rag_search", {"query": "supervisor"}, "call-1")], + [ToolCallPart("rag_cite", {"chunk_ids": ["chunk-1"]}, "call-2")], + [TextPart("first answer")], + [TextPart("second answer")], + ] + ) + wire: list[list[Any]] = [] + + async def model(messages, _info): + wire.append(list(messages)) + return ModelResponse(parts=next(calls)) + + agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag, compactor]) + deps = Deps() + + with ( + patch.object(RAGCapability, "_search", _search_with_a_picture), + patch.object( + RAGCapability, + "get_picture_bytes", + AsyncMock(side_effect=OSError("the read failed")), + ), + ): + first = await agent.run("what does the supervisor do?", deps=deps) + second = await agent.run( + "and what else?", deps=deps, message_history=first.all_messages() + ) + + assert second.output == "second answer" + assert images_of(wire[-1]) == [] + assert texts_of(wire[-1]) == [] + + +def test_a_request_is_never_left_with_no_parts(): + """Emptying a message would leave something that is not a message. + + Our own pictures always travel with the tool return in their request, so this + shape does not come from us — but a rewritten history can hold it, and a + partless request is invalid whatever produced it. + """ + ours_alone = ModelRequest( + parts=[UserPromptPart(content=[picture_label("chunk-1", "#/pictures/0"), PNG])] + ) + messages = [ + ours_alone, + ModelResponse(parts=[ToolCallPart("rag_search", {"query": "q"}, "call-1")]), + ModelRequest(parts=[ToolReturnPart("rag_search", "OLD", "call-1")]), + ModelResponse(parts=[TextPart("an answer")]), + ModelRequest(parts=[UserPromptPart("second")]), + ] + + compacted = compact_history( + messages, boundary=4, owned_tools=OWNED, capsule_text="CAPSULE" + ) + + assert all(message.parts for message in compacted) + assert compacted[0] is ours_alone From a69f3a8a98bc9ab72560f7b69771e1b0cad24b21 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 11 Aug 2026 14:00:17 +0300 Subject: [PATCH 10/13] Document the evidence record and the compaction capability The capability pages said tool results from earlier turns are replaced before every model request. That is now the compaction capability's job, and only when a host registers it, so both pages point at it instead of describing it as automatic. `RAGState` gains its `evidence` field, and the note about per-run resets now says what a resumption keeps. --- docs/capabilities/analysis.md | 4 +++- docs/capabilities/rag.md | 7 +++++-- 2 files changed, 8 insertions(+), 3 deletions(-) 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/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 From 11b1bfbc94b74060f079c331cb2b40499f922e4d Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Tue, 11 Aug 2026 16:52:57 +0300 Subject: [PATCH 11/13] Give the capsule to one return, and the run its own copy of state MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A request can carry several of this capability's returns — a model can call search twice in one response — and the carrier was identified by message alone, so each of them received the whole capsule. It is selected by message and part now, so exactly one carries it however many share the request. The chat TUI passed its persisted state into the run, so tool synchronisation mutated it in place while the message history was promoted only on success. A cancelled or failed run therefore kept the evidence the tools had recorded and discarded the messages that justified it, and the next question derived its identity from the shorter history: behind the recorded epoch, refused as non-append-only, the conversation unusable until cleared. The run gets a copy, promoted with the messages or not at all. Five decorators had been left attached to a helper by an earlier extraction, which pytest does not collect, so the resume case they carried was silently untested. The wire test covers both resume shapes again, no prompt and deferred results. --- .../haiku/rag/capabilities/compaction.py | 27 ++++--- haiku_rag_slim/haiku/rag/chat/app.py | 8 +- tests/capabilities/test_capabilities.py | 14 ---- tests/capabilities/test_evidence_wire.py | 80 ++++++++++++++++++- tests/chat/test_chat_app.py | 53 ++++++++++++ 5 files changed, 154 insertions(+), 28 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/capabilities/compaction.py b/haiku_rag_slim/haiku/rag/capabilities/compaction.py index e9b9e805..e2c051f0 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/compaction.py +++ b/haiku_rag_slim/haiku/rag/capabilities/compaction.py @@ -253,16 +253,17 @@ def compact_history( if not isinstance(message, ModelRequest): continue parts: list[Any] = [] - for part in message.parts: + for position, part in enumerate(message.parts): if isinstance(part, ToolReturnPart) and part.tool_name in owned_tools: - body = capsule_text or RECEIPT if index == carrier else RECEIPT + 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 index == carrier and capsule_images: + 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 @@ -275,15 +276,21 @@ def compact_history( def _newest_owned_return( messages: list[ModelMessage], boundary: int, owned_tools: frozenset[str] -) -> int | None: - """Index of the last earlier request holding one of our evidence returns.""" +) -> 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 isinstance(message, ModelRequest) and any( - isinstance(part, ToolReturnPart) and part.tool_name in owned_tools - for part in message.parts - ): - return 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 diff --git a/haiku_rag_slim/haiku/rag/chat/app.py b/haiku_rag_slim/haiku/rag/chat/app.py index 6aa42799..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 @@ -208,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/tests/capabilities/test_capabilities.py b/tests/capabilities/test_capabilities.py index f4f7943f..e8af1d14 100644 --- a/tests/capabilities/test_capabilities.py +++ b/tests/capabilities/test_capabilities.py @@ -872,20 +872,6 @@ def _in_flight_history() -> list[Any]: ] -@pytest.mark.asyncio -@pytest.mark.asyncio -@pytest.mark.asyncio -@pytest.mark.asyncio -@pytest.mark.parametrize( - "resume_kwargs", - [ - pytest.param({}, id="no prompt"), - pytest.param( - {"deferred_tool_results": DeferredToolResults()}, id="deferred results" - ), - ], -) -@pytest.mark.asyncio def _record(deps: Deps, namespace: str) -> CapabilityEvidenceRecord: return CapabilityEvidenceRecord.model_validate(deps.state[namespace]["evidence"]) diff --git a/tests/capabilities/test_evidence_wire.py b/tests/capabilities/test_evidence_wire.py index bcce85f5..3915fe17 100644 --- a/tests/capabilities/test_evidence_wire.py +++ b/tests/capabilities/test_evidence_wire.py @@ -4,7 +4,7 @@ from typing import Any, cast from unittest.mock import AsyncMock, patch import pytest -from pydantic_ai import Agent, RunContext +from pydantic_ai import Agent, DeferredToolResults, RunContext from pydantic_ai.messages import ( BinaryContent, ModelRequest, @@ -346,9 +346,18 @@ async def test_with_the_compactor_a_new_question_compacts_the_previous_one( 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, + temp_db_path, resume_kwargs ): """The boundary is the stored identity of the question in progress. @@ -370,7 +379,9 @@ async def test_a_resumed_question_keeps_the_evidence_it_is_answering_from( *in_flight_history(), ] - await agent.run(deps=resuming_deps(question=4), message_history=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"] @@ -653,3 +664,66 @@ def test_a_request_is_never_left_with_no_parts(): 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 == [] From 97a29d168fd242865cd2073afdcb267e5fba0486 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Wed, 12 Aug 2026 15:21:11 +0300 Subject: [PATCH 12/13] Confine the ledger's epochs to the question that recorded them A question's evidence epoch and declaration describe that question and end with it, so `begin_question` clears them. Held past it, they let the last question's citations ground the next one and hold a horizon its own declarations cannot pass, and they force every question's message count to be comparable to counts taken over a history that a host may since have rebuilt: a thread reconstructed without a run that never finished shifts every later position, and the next question was refused where answering it is correct. Identities still separate one question from the next. Occurrences outlive the question that recorded them and carry the identities that cited them, which the capsule is grouped and ordered by, so reuse merges two questions into one group and a lower identity renders later evidence as earlier. --- .../haiku/rag/capabilities/ledger.py | 38 +++++++++++--- tests/capabilities/test_evidence_ledger.py | 49 +++++++++++++++++-- 2 files changed, 75 insertions(+), 12 deletions(-) diff --git a/haiku_rag_slim/haiku/rag/capabilities/ledger.py b/haiku_rag_slim/haiku/rag/capabilities/ledger.py index 4d837f79..2750bd6c 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/ledger.py +++ b/haiku_rag_slim/haiku/rag/capabilities/ledger.py @@ -65,14 +65,15 @@ class CapabilityEvidenceRecord(BaseModel): declaration: CitationDeclaration | None = None def _reject_regression(self, count: int, what: str) -> None: - """Refuse a message count below one already recorded. + """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. 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. + 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, @@ -87,9 +88,30 @@ class CapabilityEvidenceRecord(BaseModel): ) def begin_question(self, identity: int) -> None: - """Take the identity of a question that has just arrived.""" - self._reject_regression(identity, "A question") + """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. diff --git a/tests/capabilities/test_evidence_ledger.py b/tests/capabilities/test_evidence_ledger.py index fa849c64..c9f94353 100644 --- a/tests/capabilities/test_evidence_ledger.py +++ b/tests/capabilities/test_evidence_ledger.py @@ -47,14 +47,13 @@ def test_refs_make_it_grounded_and_no_refs_make_it_ungrounded(): def test_an_earlier_questions_declaration_is_never_current(): - """Epochs outlive a question, so the epoch alone would inherit it.""" + """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 record.declaration is not None assert citation_status([record], question=8) == "missing" @@ -193,9 +192,51 @@ def test_evidence_cannot_predate_a_recorded_declaration(): record.note_evidence(3) -def test_a_question_cannot_start_before_what_is_already_recorded(): +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) - with pytest.raises(ValueError, match="append-only"): + 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" From f4d8e744a2178b7fa7e07963fde4afad154ed832 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Thu, 13 Aug 2026 11:21:31 +0300 Subject: [PATCH 13/13] Record this release's ORB multimodal numbers The nemotron rows for retrieval and for both capabilities are re-measured on this release with no reranker, judged by Qwen3.6-35B with thinking on. Rows for other embedders and older versions keep their own attribution. --- docs/benchmarks.md | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) 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