From 5d76461b9f0b9519df0fa933364be54144bc3657 Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Fri, 7 Aug 2026 12:40:05 +0300 Subject: [PATCH] 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