From 8eef0b37344f8123fc5de8faa2ec425233a2ef3a Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 10 Aug 2026 12:15:47 +0300 Subject: [PATCH] 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.