From fa0c4a4f500cd07989d3fde5d918e2de3117a5ba Mon Sep 17 00:00:00 2001 From: Yiorgis Gozadinos Date: Mon, 10 Aug 2026 14:13:19 +0300 Subject: [PATCH] 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