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