diff --git a/python/restate/ext/pydantic/_agent.py b/python/restate/ext/pydantic/_agent.py index cc7888c..0308ed7 100644 --- a/python/restate/ext/pydantic/_agent.py +++ b/python/restate/ext/pydantic/_agent.py @@ -27,7 +27,7 @@ from pydantic_ai.models import Model from pydantic_ai.output import OutputDataT, OutputSpec from pydantic_ai.result import StreamedRunResult -from pydantic_ai.run import AgentRunResult +from pydantic_ai.run import AgentRun, AgentRunResult from pydantic_ai.tools import DeferredToolResults, RunContext from pydantic_ai.toolsets.abstract import AbstractToolset from pydantic_ai.toolsets.function import FunctionToolset @@ -308,6 +308,48 @@ async def main(): spec=spec, ) + @overload + def iter( + self, + user_prompt: str | Sequence[UserContent] | None = None, + *, + output_type: None = None, + model: models.Model | models.KnownModelName | str | None = None, + **kwargs: Any, + ) -> AbstractAsyncContextManager[AgentRun[AgentDepsT, OutputDataT]]: ... + + @overload + def iter( + self, + user_prompt: str | Sequence[UserContent] | None = None, + *, + output_type: OutputSpec[RunOutputDataT], + model: models.Model | models.KnownModelName | str | None = None, + **kwargs: Any, + ) -> AbstractAsyncContextManager[AgentRun[AgentDepsT, RunOutputDataT]]: ... + + @asynccontextmanager + async def iter( + self, + user_prompt: str | Sequence[UserContent] | None = None, + *, + output_type: OutputSpec[RunOutputDataT] | None = None, + model: models.Model | models.KnownModelName | str | None = None, + **kwargs: Any, + ) -> AsyncIterator[AgentRun[AgentDepsT, Any]]: + """Iterate over the agent graph's nodes with the same durable model and toolsets as `run()`. + + The remaining keyword arguments are forwarded unchanged so that this override accepts every + argument `AbstractAgent.run()` passes to `iter()` across supported pydantic-ai versions. + """ + if model is not None: + raise TerminalError( + "An agent needs to have a `model` in order to be used with Restate, it cannot be set at agent run time." + ) + with self._restate_overrides(): + async with super().iter(user_prompt, output_type=output_type, model=model, **kwargs) as agent_run: + yield agent_run + @overload def run_stream( self, diff --git a/python/restate/ext/pydantic/_model.py b/python/restate/ext/pydantic/_model.py index fd162da..3c5812a 100644 --- a/python/restate/ext/pydantic/_model.py +++ b/python/restate/ext/pydantic/_model.py @@ -120,6 +120,7 @@ async def request_stream_run(): ) try: response = await context.run_typed("Model stream call", request_stream_run, self._options) + current_state().turnstile = Turnstile([c.tool_call_id for c in response.tool_calls]) yield RestateStreamedResponse(model_request_parameters, response) except SdkInternalBaseException as e: raise Exception("Internal error during model stream call") from e diff --git a/python/restate/ext/pydantic/_toolset.py b/python/restate/ext/pydantic/_toolset.py index 2f62db7..4d4e153 100644 --- a/python/restate/ext/pydantic/_toolset.py +++ b/python/restate/ext/pydantic/_toolset.py @@ -11,6 +11,7 @@ from pydantic_ai import ToolDefinition from pydantic_ai._run_context import AgentDepsT from pydantic_ai.exceptions import ApprovalRequired, CallDeferred, ModelRetry, UserError +from pydantic_ai.messages import ToolReturn from pydantic_ai.mcp import MCPToolset from pydantic_ai.tools import RunContext from pydantic_ai.toolsets.abstract import AbstractToolset, ToolsetTool @@ -27,6 +28,9 @@ class RestateContextRunResult: kind: Literal["output", "call_deferred", "approval_required", "model_retry"] output: Any error: str | None = None + # `output: Any` round-trips a ToolReturn as a plain dict, so it is journaled in its own typed field. + tool_return: ToolReturn | None = None + metadata: dict[str, Any] | None = None CONTEXT_RUN_SERDE = PydanticTypeAdapter(RestateContextRunResult) @@ -69,13 +73,15 @@ async def action() -> RestateContextRunResult: # Since, restate ctx.run() will retry this exception we need to convert these exceptions # to a return value and handle them outside of the ctx.run(). output = await self.wrapped.call_tool(name, tool_args, ctx, tool) + if isinstance(output, ToolReturn): + return RestateContextRunResult(kind="output", output=None, tool_return=output) return RestateContextRunResult(kind="output", output=output, error=None) except ModelRetry as e: return RestateContextRunResult(kind="model_retry", output=None, error=e.message) - except CallDeferred: - return RestateContextRunResult(kind="call_deferred", output=None, error=None) - except ApprovalRequired: - return RestateContextRunResult(kind="approval_required", output=None, error=None) + except CallDeferred as e: + return RestateContextRunResult(kind="call_deferred", output=None, metadata=e.metadata) + except ApprovalRequired as e: + return RestateContextRunResult(kind="approval_required", output=None, metadata=e.metadata) except UserError as e: raise TerminalError(str(e)) from e @@ -93,16 +99,25 @@ async def action() -> RestateContextRunResult: res = await context.run_typed(f"Calling {name}", action, self.options) if res.kind == "call_deferred": - raise CallDeferred() + # Deferring, approval and retry concern only this call; parallel siblings still + # need to run so the agent can continue with their results. + turnstile.allow_next_after(id) + raise CallDeferred(metadata=res.metadata) elif res.kind == "approval_required": - raise ApprovalRequired() + turnstile.allow_next_after(id) + raise ApprovalRequired(metadata=res.metadata) elif res.kind == "model_retry": assert res.error is not None + turnstile.allow_next_after(id) raise ModelRetry(res.error) else: assert res.kind == "output" turnstile.allow_next_after(id) + if res.tool_return is not None: + return res.tool_return return res.output + except (CallDeferred, ApprovalRequired, ModelRetry): + raise except Exception as e: turnstile.cancel_all_after(id) raise e from None diff --git a/python/restate/ext/turnstile.py b/python/restate/ext/turnstile.py index 08981c6..f00e3f2 100644 --- a/python/restate/ext/turnstile.py +++ b/python/restate/ext/turnstile.py @@ -22,12 +22,21 @@ def __init__(self, ids: list[str]): event.set() async def wait_for(self, id: str) -> None: - event = self.events[id] + event = self.events.get(id) + if event is None: + # Untracked ids (resumed deferred calls, nested calls with synthesized ids) are not part of + # this turn's ordering, so they proceed unless the turn was already canceled. + if self.canceled: + raise SdkInternalException() from None + return await event.wait() if self.canceled: raise SdkInternalException() from None def cancel_all_after(self, id: str) -> None: + if id not in self.events: + # A failing untracked call must not cancel the tracked calls of this turn. + return self.canceled = True next_id = self.turns.get(id) while next_id is not None: diff --git a/tests/test_pydantic_ext.py b/tests/test_pydantic_ext.py new file mode 100644 index 0000000..5e92e79 --- /dev/null +++ b/tests/test_pydantic_ext.py @@ -0,0 +1,345 @@ +# +# Copyright (c) 2023-2025 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +"""Unit tests for the pydantic-ai integration, driven by a fake Restate context. + +The fake context executes each `run_typed` action and round-trips its result through the configured +serde, which is what a journaled `ctx.run` does on replay. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import AsyncIterable, AsyncIterator, Iterator +from typing import Any + +import pytest + +pytest.importorskip("pydantic_ai") + +from pydantic_ai import Agent # noqa: E402 +from pydantic_ai.exceptions import ApprovalRequired, CallDeferred, ModelRetry # noqa: E402 +from pydantic_ai.messages import ( # noqa: E402 + ModelMessage, + ModelResponse, + RetryPromptPart, + TextPart, + ToolCallPart, + ToolReturn, + ToolReturnPart, + UserPromptPart, +) +from pydantic_ai.models.function import AgentInfo, DeltaToolCall, DeltaToolCalls, FunctionModel # noqa: E402 +from pydantic_ai.tools import DeferredToolRequests, DeferredToolResults # noqa: E402 + +from restate import RunOptions, TerminalError # noqa: E402 +from restate.exceptions import SdkInternalException # noqa: E402 +from restate.ext.pydantic import RestateAgent # noqa: E402 +from restate.ext.pydantic._model import RestateModelWrapper # noqa: E402 +from restate.ext.turnstile import Turnstile # noqa: E402 +from restate.server_context import _restate_context_var # noqa: E402 + + +@pytest.fixture(scope="session") +def anyio_backend(): + return "asyncio" + + +pytestmark = [ + pytest.mark.anyio, +] + + +class FakeContext: + def __init__(self) -> None: + self.run_names: list[str] = [] + + async def run_typed(self, name: str, action: Any, options: RunOptions | None = None, *args: Any, **kwargs: Any): + self.run_names.append(name) + result = action(*args, **kwargs) + if asyncio.iscoroutine(result): + result = await result + if options is None or options.serde is None: + return result + return options.serde.deserialize(options.serde.serialize(result)) + + +@pytest.fixture +def fake_context() -> Iterator[FakeContext]: + ctx = FakeContext() + token = _restate_context_var.set(ctx) # type: ignore[arg-type] + try: + yield ctx + finally: + _restate_context_var.reset(token) + + +def _has_tool_return(messages: list[ModelMessage]) -> bool: + return any(isinstance(part, ToolReturnPart) for message in messages for part in message.parts) + + +def _tool_calls_then_text(*calls: ToolCallPart) -> FunctionModel: + def model_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: + if _has_tool_return(messages): + return ModelResponse(parts=[TextPart("done")]) + return ModelResponse(parts=list(calls)) + + return FunctionModel(model_fn) + + +def _tool_returns(messages: list[ModelMessage]) -> dict[str, ToolReturnPart]: + return { + part.tool_call_id: part for message in messages for part in message.parts if isinstance(part, ToolReturnPart) + } + + +async def test_tool_return_survives_journal_round_trip(fake_context: FakeContext): + agent = Agent(_tool_calls_then_text(ToolCallPart("lookup", {}, tool_call_id="c1"))) + + @agent.tool_plain + def lookup() -> ToolReturn: + return ToolReturn(return_value={"temp": 21}, content="extra context", metadata={"source": "sensor"}) + + result = await RestateAgent(agent, auto_wrap_tools=True).run("go") + + messages = result.all_messages() + part = _tool_returns(messages)["c1"] + assert part.content == {"temp": 21} + assert part.metadata == {"source": "sensor"} + user_prompts = [p.content for m in messages for p in m.parts if isinstance(p, UserPromptPart)] + assert "extra context" in user_prompts + assert "Calling lookup" in fake_context.run_names + + +async def test_plain_tool_output_still_round_trips(fake_context: FakeContext): + agent = Agent(_tool_calls_then_text(ToolCallPart("lookup", {}, tool_call_id="c1"))) + + @agent.tool_plain + def lookup() -> dict[str, int]: + return {"temp": 21} + + result = await RestateAgent(agent, auto_wrap_tools=True).run("go") + + part = _tool_returns(result.all_messages())["c1"] + assert part.content == {"temp": 21} + assert part.metadata is None + + +async def test_call_deferred_metadata_is_preserved(fake_context: FakeContext): + agent = Agent( + _tool_calls_then_text(ToolCallPart("slow", {}, tool_call_id="c1")), + output_type=[str, DeferredToolRequests], + ) + + @agent.tool_plain + def slow() -> str: + raise CallDeferred(metadata={"job": "42"}) + + result = await RestateAgent(agent, auto_wrap_tools=True).run("go") + + assert isinstance(result.output, DeferredToolRequests) + assert [c.tool_call_id for c in result.output.calls] == ["c1"] + assert result.output.metadata == {"c1": {"job": "42"}} + + +async def test_parallel_approval_required_does_not_abort_sibling(fake_context: FakeContext): + agent = Agent( + _tool_calls_then_text( + ToolCallPart("dangerous", {}, tool_call_id="c1"), + ToolCallPart("safe", {}, tool_call_id="c2"), + ), + output_type=[str, DeferredToolRequests], + ) + + @agent.tool_plain + def dangerous() -> str: + raise ApprovalRequired(metadata={"reason": "destructive"}) + + @agent.tool_plain + def safe() -> str: + return "ok" + + result = await RestateAgent(agent, auto_wrap_tools=True).run("go") + + assert isinstance(result.output, DeferredToolRequests) + assert [c.tool_call_id for c in result.output.approvals] == ["c1"] + assert result.output.metadata == {"c1": {"reason": "destructive"}} + assert _tool_returns(result.all_messages())["c2"].content == "ok" + + +async def test_parallel_call_deferred_does_not_abort_sibling(fake_context: FakeContext): + agent = Agent( + _tool_calls_then_text( + ToolCallPart("slow", {}, tool_call_id="c1"), + ToolCallPart("safe", {}, tool_call_id="c2"), + ), + output_type=[str, DeferredToolRequests], + ) + + @agent.tool_plain + def slow() -> str: + raise CallDeferred(metadata={"job": "42"}) + + @agent.tool_plain + def safe() -> str: + return "ok" + + result = await RestateAgent(agent, auto_wrap_tools=True).run("go") + + assert isinstance(result.output, DeferredToolRequests) + assert [c.tool_call_id for c in result.output.calls] == ["c1"] + assert result.output.metadata == {"c1": {"job": "42"}} + assert _tool_returns(result.all_messages())["c2"].content == "ok" + + +async def test_parallel_model_retry_does_not_abort_sibling(fake_context: FakeContext): + seen: list[list[ModelMessage]] = [] + + def model_fn(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: + seen.append(list(messages)) + if len(seen) > 1: + return ModelResponse(parts=[TextPart("done")]) + return ModelResponse( + parts=[ + ToolCallPart("flaky", {}, tool_call_id="c1"), + ToolCallPart("safe", {}, tool_call_id="c2"), + ] + ) + + agent = Agent(FunctionModel(model_fn)) + calls: list[str] = [] + + @agent.tool_plain + def flaky() -> str: + calls.append("flaky") + raise ModelRetry("try again later") + + @agent.tool_plain + def safe() -> str: + calls.append("safe") + return "ok" + + result = await RestateAgent(agent, auto_wrap_tools=True).run("go") + + assert result.output == "done" + assert calls == ["flaky", "safe"] + retry_request = seen[1][-1] + retries = [p for p in retry_request.parts if isinstance(p, RetryPromptPart)] + returns = [p for p in retry_request.parts if isinstance(p, ToolReturnPart)] + assert [(p.tool_call_id, p.content) for p in retries] == [("c1", "try again later")] + assert [(p.tool_call_id, p.content) for p in returns] == [("c2", "ok")] + + +async def test_request_stream_initializes_turnstile(fake_context: FakeContext): + async def stream_fn(messages: list[ModelMessage], info: AgentInfo) -> AsyncIterator[str | DeltaToolCalls]: + if _has_tool_return(messages): + yield "done" + else: + yield {0: DeltaToolCall(name="lookup", json_args="{}", tool_call_id="c1")} + yield {1: DeltaToolCall(name="lookup", json_args="{}", tool_call_id="c2")} + + async def handler(ctx: Any, events: AsyncIterable[Any]) -> None: + async for _ in events: + pass + + agent = Agent(FunctionModel(stream_function=stream_fn)) + calls: list[str] = [] + + @agent.tool_plain + def lookup() -> str: + calls.append("lookup") + return "ok" + + result = await RestateAgent(agent, auto_wrap_tools=True, event_stream_handler=handler).run("go") + + assert result.output == "done" + assert calls == ["lookup", "lookup"] + assert "Model stream call" in fake_context.run_names + assert set(_tool_returns(result.all_messages())) == {"c1", "c2"} + + +async def test_resumed_deferred_call_runs_without_model_request(fake_context: FakeContext): + agent = Agent( + _tool_calls_then_text(ToolCallPart("dangerous", {}, tool_call_id="c1")), + output_type=[str, DeferredToolRequests], + ) + + @agent.tool_plain(requires_approval=True) + def dangerous() -> str: + return "deleted" + + restate_agent = RestateAgent(agent, auto_wrap_tools=True) + first = await restate_agent.run("go") + assert isinstance(first.output, DeferredToolRequests) + + resumed = await restate_agent.run( + message_history=first.all_messages(), + deferred_tool_results=DeferredToolResults(approvals={"c1": True}), + ) + + assert resumed.output == "done" + assert _tool_returns(resumed.all_messages())["c1"].content == "deleted" + + +async def test_turnstile_lets_untracked_id_through(): + turnstile = Turnstile(["a", "b"]) + + await asyncio.wait_for(turnstile.wait_for("nested"), timeout=1) + turnstile.allow_next_after("nested") + turnstile.cancel_all_after("nested") + + assert not turnstile.canceled + await asyncio.wait_for(turnstile.wait_for("a"), timeout=1) + + +async def test_turnstile_untracked_id_raises_after_cancel(): + turnstile = Turnstile(["a", "b"]) + turnstile.cancel_all_after("a") + + with pytest.raises(SdkInternalException): + await turnstile.wait_for("nested") + + +async def test_turnstile_allow_next_after_deferred_unblocks_sibling(): + turnstile = Turnstile(["a", "b"]) + sibling = asyncio.create_task(turnstile.wait_for("b")) + await asyncio.sleep(0) + assert not sibling.done() + + turnstile.allow_next_after("a") + + await asyncio.wait_for(sibling, timeout=1) + + +async def test_iter_uses_durable_model_and_toolsets(fake_context: FakeContext): + agent = Agent(_tool_calls_then_text(ToolCallPart("lookup", {}, tool_call_id="c1"))) + + @agent.tool_plain + def lookup() -> str: + return "ok" + + restate_agent = RestateAgent(agent, auto_wrap_tools=True) + async with restate_agent.iter("go") as agent_run: + async for _ in agent_run: + pass + + assert agent_run.result is not None + assert agent_run.result.output == "done" + assert fake_context.run_names.count("Model call") == 2 + assert "Calling lookup" in fake_context.run_names + assert isinstance(restate_agent.model, RestateModelWrapper) + + +async def test_iter_rejects_run_time_model(fake_context: FakeContext): + agent = Agent(_tool_calls_then_text()) + + with pytest.raises(TerminalError): + async with RestateAgent(agent).iter("go", model=_tool_calls_then_text()): + pass