diff --git a/CHANGELOG.md b/CHANGELOG.md index cb1d2e1f..8d8bc2b2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,7 +6,7 @@ - `CitationPolicyCapability` (`haiku.rag.capabilities.policy.create_capability`): registering it requires every answer to declare its grounding, in any conversation that has something to declare — this question retrieved evidence, or something was cited earlier. A question that ends undeclared is sent back once to record what grounded the answer already given, and is recorded in `CitationPolicyState.violations` if it finishes undeclared regardless. A conversation with neither a current-question evidence outcome nor any earlier citation is not enforced. - `haiku.rag.capabilities.evidence.discover_evidence()` and `DiscoveredEvidence`, moved out of `compaction` so both optional capabilities share them. `RAGCapabilityBase.cite_available`. - `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. +- `RAGState.evidence` / `AnalysisState.evidence` (`CapabilityEvidenceRecord`): which evidence a capability retrieved and cited, per question, keyed by message-count question identities and epochs, and whether that question is still being answered. `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()`. @@ -14,7 +14,7 @@ - `rag_cite` / `analysis_cite` accept an empty `chunk_ids`, recording the answer as ungrounded rather than failing the call, and the instructions no longer exempt a refusal or a corpus-level computation from citing. - `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. +- Resuming a run (deferred tool results, an unfinished history tail) raises `RuntimeError` unless the host carries the capability state from the run being resumed. ### Removed diff --git a/haiku_rag_slim/haiku/rag/capabilities/_base.py b/haiku_rag_slim/haiku/rag/capabilities/_base.py index fe13d323..ab844ab8 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/_base.py +++ b/haiku_rag_slim/haiku/rag/capabilities/_base.py @@ -6,13 +6,19 @@ from pathlib import Path from typing import Any, cast from pydantic import BaseModel -from pydantic_ai import ModelRetry, RunContext, ToolFailed +from pydantic_ai import ( + DeferredToolRequests, + ModelRetry, + RunContext, + ToolFailed, +) from pydantic_ai.capabilities import AbstractCapability from pydantic_ai.messages import ( InstructionPart, ModelMessage, ModelRequest, ModelResponse, + RetryPromptPart, ToolCallPart, ToolReturn, ) @@ -82,29 +88,21 @@ def _clear_invocation_state(state: BaseModel) -> None: value.clear() -def _is_resumption(prompt: Any, messages: list[ModelMessage]) -> bool: - """Whether this run continues a question rather than asking a new one. +def _awaits_the_model(messages: list[ModelMessage]) -> bool: + """Whether the history unmistakably leaves the model something to answer. - 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. + Used to validate what the record already says, never to decide it. Only two + shapes are unambiguous: a response whose tool calls have no returns, and a + retry the model has not answered. A trailing tool return is not one of them, + being both how a settled structured answer ends and how results reach a + question still in progress. """ - 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) + if isinstance(last, ModelResponse): + return any(isinstance(part, ToolCallPart) for part in last.parts) + return any(isinstance(part, RetryPromptPart) for part in last.parts) def _called_own_tool(messages: list[ModelMessage], tool_names: frozenset[str]) -> bool: @@ -154,11 +152,10 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): 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) - 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: + continuing = record.in_progress + if not continuing and _awaits_the_model(ctx.messages): raise RuntimeError( f"The {self.state_namespace} capability is resuming a question with " "no stored question identity. Capabilities cannot be added, removed " @@ -321,6 +318,16 @@ class RAGCapabilityBase[StateT: BaseModel](AbstractCapability[Any]): async def after_run( self, ctx: RunContext[Any], *, result: AgentRunResult[Any] ) -> AgentRunResult[Any]: + """Close the question, unless the run is only pausing for deferred results. + + A run that raised never arrives here, which is what leaves an interrupted + question in progress for the resumption to claim. + """ + if self.state is not None and not isinstance( + result.output, DeferredToolRequests + ): + self._evidence_record().end_question() + self._sync_state() await self._close() return result diff --git a/haiku_rag_slim/haiku/rag/capabilities/ledger.py b/haiku_rag_slim/haiku/rag/capabilities/ledger.py index 778ae7ec..4599f672 100644 --- a/haiku_rag_slim/haiku/rag/capabilities/ledger.py +++ b/haiku_rag_slim/haiku/rag/capabilities/ledger.py @@ -51,6 +51,11 @@ class CapabilityEvidenceRecord(BaseModel): overwritten by whichever of them synced its state last; merging happens in the transient views built by ``citation_status`` and the optional capabilities. + ``in_progress`` is whether the question this record names is still being + answered. It is the only authority on that: a transcript ending in a tool + return is a settled structured answer and a question awaiting its model + equally, so the shape of the history cannot decide it. + ``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 @@ -61,6 +66,7 @@ class CapabilityEvidenceRecord(BaseModel): occurrences: dict[str, EvidenceOccurrence] = Field(default_factory=dict) question: int | None = None + in_progress: bool = False latest_evidence_epoch: int = 0 declaration: CitationDeclaration | None = None @@ -110,9 +116,14 @@ class CapabilityEvidenceRecord(BaseModel): "one question from the next and are compared as recency." ) self.question = identity + self.in_progress = True self.latest_evidence_epoch = 0 self.declaration = None + def end_question(self) -> None: + """Mark the question answered, so the next run knows it is a new one.""" + self.in_progress = False + def note_evidence(self, epoch: int) -> None: """Record that the model has seen an evidence outcome. diff --git a/tests/capabilities/test_capabilities.py b/tests/capabilities/test_capabilities.py index cb58b2c7..f73fb7b6 100644 --- a/tests/capabilities/test_capabilities.py +++ b/tests/capabilities/test_capabilities.py @@ -4,7 +4,16 @@ from typing import Any, cast from unittest.mock import AsyncMock, patch import pytest -from pydantic_ai import Agent, DeferredToolResults, ModelRetry, RunContext, ToolFailed +from pydantic import BaseModel +from pydantic_ai import ( + Agent, + CallDeferred, + DeferredToolRequests, + DeferredToolResults, + ModelRetry, + RunContext, + ToolFailed, +) from pydantic_ai.messages import ( ModelRequest, ModelResponse, @@ -920,9 +929,9 @@ async def test_a_resumption_keeps_the_identity_of_the_question_in_progress( agent = Agent(FunctionModel(model), deps_type=Deps, capabilities=[rag]) deps = Deps( state={ - "rag": RAGState(evidence=CapabilityEvidenceRecord(question=7)).model_dump( - mode="json" - ) + "rag": RAGState( + evidence=CapabilityEvidenceRecord(question=7, in_progress=True) + ).model_dump(mode="json") } ) history = [ @@ -1159,10 +1168,12 @@ async def test_citing_without_searching_grounds_the_question(temp_db_path): @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. + """A seeded record says nothing about a question, so the history has to. - 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. + Seeding one is what a host does when it has no state to send. Its flag is + unset, so a history that unmistakably awaits the model means the host dropped + the state of a question in progress, and answering as question zero would + silently relabel it. """ rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False) @@ -1174,7 +1185,10 @@ async def test_a_host_seeded_record_does_not_pass_for_a_resumption(temp_db_path) with pytest.raises(RuntimeError, match="no stored question identity"): await agent.run( "carry on", - message_history=_in_flight_history(), + message_history=[ + *_in_flight_history(), + ModelResponse(parts=[ToolCallPart("external_tool", {}, "call-2")]), + ], deps=Deps(state={"rag": RAGState().model_dump(mode="json")}), ) @@ -1208,6 +1222,9 @@ async def test_a_resumption_keeps_the_evidence_the_question_already_gathered( interrupted = await agent.run("what does the supervisor do?", deps=deps) identity = _record(deps, "rag").question assert identity is not None + # A run that ends awaiting external work leaves the question in progress, + # which is what the resumption claims. See the deferred-request test. + deps.state["rag"]["evidence"]["in_progress"] = True await agent.run( deferred_tool_results=DeferredToolResults( calls={"call-2": "external result"} @@ -1280,3 +1297,147 @@ async def test_citing_nothing_after_citing_something_keeps_it_grounded(temp_db_p record = capability.state.evidence assert citation_status([record], question=0) == "grounded" + + +@pytest.mark.asyncio +async def test_a_promptless_run_on_a_settled_history_is_a_new_question(temp_db_path): + """AG-UI hosts never pass a prompt: the client's message is the history. + + Pydantic AI's UI adapter builds `message_history` from the frontend messages + and calls the agent without a prompt, so the run has no prompt *and* the + history ends with the user's own request. Reading either as a continuation + fails every AG-UI host on its first message. + """ + 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().model_dump(mode="json")}) + history: list[Any] = [ + ModelRequest(parts=[UserPromptPart("what does the manual say about masks?")]) + ] + + await agent.run(message_history=history, deps=deps) + + assert _record(deps, "rag").question == len(history) + + +@pytest.mark.asyncio +async def test_a_promptless_run_on_an_unfinished_tail_is_still_a_continuation( + temp_db_path, +): + """A suspended run resumes without a prompt, and must keep its question.""" + 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=3, in_progress=True) + ).model_dump(mode="json") + } + ) + + await agent.run(message_history=_in_flight_history(), deps=deps) + + assert _record(deps, "rag").question == 3 + + +@pytest.mark.asyncio +async def test_a_structured_answer_does_not_leave_the_question_in_progress( + temp_db_path, +): + """A settled run ends with a tool return, which says nothing about progress. + + Pydantic AI answers a structured `output_type` by calling an output tool, so the + history ends with a request carrying that tool's return. Reading the transcript + shape alone, that is indistinguishable from tool results delivered to a question + still being answered, and every following question inherited the first one's + identity. + """ + + class Answer(BaseModel): + text: str + + rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False) + + async def model(_messages, info): + return ModelResponse( + parts=[ToolCallPart(info.output_tools[0].name, {"text": "answer"})] + ) + + agent = Agent( + FunctionModel(model), + deps_type=Deps, + capabilities=[rag], + output_type=Answer, + ) + deps = Deps() + + first = await agent.run("first question", deps=deps) + first_identity = _record(deps, "rag").question + assert _record(deps, "rag").in_progress is False + + await agent.run( + "second question", deps=deps, message_history=list(first.all_messages()) + ) + + second_identity = _record(deps, "rag").question + assert first_identity == 0 + assert second_identity is not None and second_identity > 0 + + +@pytest.mark.asyncio +async def test_a_run_pausing_for_deferred_work_leaves_the_question_in_progress( + temp_db_path, +): + """The question is unfinished, so its resumption must find it claimable. + + A deferred tool call ends the run with `DeferredToolRequests` rather than an + answer. Closing the question here would let the resumption relabel it. + """ + rag = create_rag(db_path=temp_db_path, config=AppConfig(), defer_loading=False) + + async def model(_messages, _info): + return ModelResponse(parts=[ToolCallPart("external_tool", {})]) + + agent = Agent( + FunctionModel(model), + deps_type=Deps, + capabilities=[rag], + output_type=[str, DeferredToolRequests], + ) + + @agent.tool_plain + def external_tool() -> str: + raise CallDeferred + + deps = Deps() + + result = await agent.run("a question needing external work", deps=deps) + + assert isinstance(result.output, DeferredToolRequests) + assert _record(deps, "rag").in_progress is True + + +@pytest.mark.asyncio +async def test_an_answered_question_is_no_longer_in_progress(temp_db_path): + """The flag is what tells the next run it is asking something new.""" + 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("a question", deps=deps) + + record = _record(deps, "rag") + assert record.question == 0 + assert record.in_progress is False diff --git a/tests/capabilities/test_citation_policy.py b/tests/capabilities/test_citation_policy.py index b800f8c6..bcc3fde9 100644 --- a/tests/capabilities/test_citation_policy.py +++ b/tests/capabilities/test_citation_policy.py @@ -365,7 +365,9 @@ async def test_a_resumed_question_is_not_redirected_twice(temp_db_path): with patch.object(RAGCapability, "_search", stub_search): first = await agent.run("what does the supervisor do?", deps=deps) - # The same question again, continued rather than asked anew. + # The same question again, continued rather than asked anew: a run that + # ends awaiting external work leaves the question in progress. + deps.state["rag"]["evidence"]["in_progress"] = True await agent.run( deps=deps, message_history=[ diff --git a/tests/capabilities/test_evidence_wire.py b/tests/capabilities/test_evidence_wire.py index 3915fe17..6ad32aa4 100644 --- a/tests/capabilities/test_evidence_wire.py +++ b/tests/capabilities/test_evidence_wire.py @@ -299,11 +299,11 @@ def in_flight_history() -> list[Any]: def resuming_deps(question: int = 0) -> Deps: - """State as a resumption always finds it: the question already identified.""" + """State as a resumption always finds it: the question identified and unfinished.""" return Deps( state={ "rag": RAGState( - evidence=CapabilityEvidenceRecord(question=question) + evidence=CapabilityEvidenceRecord(question=question, in_progress=True) ).model_dump(mode="json") } )