diff --git a/python/restate/ext/pydantic/_agent.py b/python/restate/ext/pydantic/_agent.py index cc7888c..dfe159d 100644 --- a/python/restate/ext/pydantic/_agent.py +++ b/python/restate/ext/pydantic/_agent.py @@ -121,7 +121,7 @@ def set_context(toolset: AbstractToolset[AgentDepsT]) -> AbstractToolset[AgentDe try: from pydantic_ai.mcp import MCPToolset - from ._toolset import RestateMCPToolset + from ._mcp_toolset import RestateMCPToolset except ImportError: pass else: diff --git a/python/restate/ext/pydantic/_mcp_toolset.py b/python/restate/ext/pydantic/_mcp_toolset.py new file mode 100644 index 0000000..d303b30 --- /dev/null +++ b/python/restate/ext/pydantic/_mcp_toolset.py @@ -0,0 +1,100 @@ +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import dataclass, replace +from typing import Any + +from restate import RunOptions, SdkInternalBaseException +from restate.extensions import current_context + +from pydantic_ai import ToolDefinition +from pydantic_ai._run_context import AgentDepsT +from pydantic_ai.exceptions import UserError +from pydantic_ai.mcp import MCPToolset +from pydantic_ai.tools import RunContext +from pydantic_ai.toolsets.abstract import AbstractToolset, ToolsetTool +from pydantic_ai.toolsets.wrapper import WrapperToolset + +from ._serde import PydanticTypeAdapter + + +@dataclass +class RestateMCPGetToolsContextRunResult: + """A simple wrapper for tool results to be used with Restate's run_typed.""" + + output: dict[str, ToolDefinition] + + +MCP_GET_TOOLS_SERDE = PydanticTypeAdapter(RestateMCPGetToolsContextRunResult) + + +@dataclass +class RestateMCPToolRunResult: + """A simple wrapper for tool results to be used with Restate's run_typed.""" + + output: Any + + +MCP_RUN_SERDE = PydanticTypeAdapter(RestateMCPToolRunResult) + + +class RestateMCPToolset(WrapperToolset[AgentDepsT]): + """A wrapper for MCPToolset that integrates with Restate.""" + + def __init__(self, wrapped: MCPToolset[AgentDepsT], run_options: RunOptions): + super().__init__(wrapped) + self._wrapped = wrapped + self.get_tools_options = replace(run_options, serde=MCP_GET_TOOLS_SERDE) + self.run_tools_options = replace(run_options, serde=MCP_RUN_SERDE) + + def visit_and_replace( + self, visitor: Callable[[AbstractToolset[AgentDepsT]], AbstractToolset[AgentDepsT]] + ) -> AbstractToolset[AgentDepsT]: + return visitor(self) + + async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[AgentDepsT]]: + async def get_tools_in_context() -> RestateMCPGetToolsContextRunResult: + res = await self._wrapped.get_tools(ctx) + # ToolsetTool is not serializable as it holds a SchemaValidator + # (which is also the same for every MCP tool so unnecessary to pass along the wire every time), + # so we just return the ToolDefinitions and wrap them in ToolsetTool outside of the activity. + return RestateMCPGetToolsContextRunResult(output={name: tool.tool_def for name, tool in res.items()}) + + context = current_context() + if context is None: + raise UserError( + "A toolset cannot be used without a Restate context. Make sure to run it within an agent or a run context." + ) + + try: + tool_defs = await context.run_typed("get mcp tools", get_tools_in_context, self.get_tools_options) + return {name: self.tool_for_tool_def(tool_def) for name, tool_def in tool_defs.output.items()} + except SdkInternalBaseException as e: + raise Exception("Internal error during get_tools call") from e + + def tool_for_tool_def(self, tool_def: ToolDefinition) -> ToolsetTool[AgentDepsT]: + assert isinstance(self.wrapped, MCPToolset) + return self.wrapped.tool_for_tool_def(tool_def) + + async def call_tool( + self, + name: str, + tool_args: dict[str, Any], + ctx: RunContext[AgentDepsT], + tool: ToolsetTool[AgentDepsT], + ) -> Any: + async def call_tool_in_context() -> RestateMCPToolRunResult: + res = await self._wrapped.call_tool(name, tool_args, ctx, tool) + return RestateMCPToolRunResult(output=res) + + context = current_context() + if context is None: + raise UserError( + "A toolset cannot be used without a Restate context. Make sure to run it within an agent or a run context." + ) + try: + res = await context.run_typed(f"Calling mcp tool {name}", call_tool_in_context, self.run_tools_options) + except SdkInternalBaseException as e: + raise Exception("Internal error during tool call") from e + + return res.output diff --git a/python/restate/ext/pydantic/_toolset.py b/python/restate/ext/pydantic/_toolset.py index 2f62db7..7852f4d 100644 --- a/python/restate/ext/pydantic/_toolset.py +++ b/python/restate/ext/pydantic/_toolset.py @@ -4,14 +4,12 @@ from dataclasses import dataclass, replace from typing import Any, Literal -from restate import RunOptions, SdkInternalBaseException, TerminalError +from restate import RunOptions, TerminalError from restate.exceptions import SdkInternalException from restate.extensions import current_context -from pydantic_ai import ToolDefinition from pydantic_ai._run_context import AgentDepsT from pydantic_ai.exceptions import ApprovalRequired, CallDeferred, ModelRetry, UserError -from pydantic_ai.mcp import MCPToolset from pydantic_ai.tools import RunContext from pydantic_ai.toolsets.abstract import AbstractToolset, ToolsetTool from pydantic_ai.toolsets.wrapper import WrapperToolset @@ -32,26 +30,6 @@ class RestateContextRunResult: CONTEXT_RUN_SERDE = PydanticTypeAdapter(RestateContextRunResult) -@dataclass -class RestateMCPGetToolsContextRunResult: - """A simple wrapper for tool results to be used with Restate's run_typed.""" - - output: dict[str, ToolDefinition] - - -MCP_GET_TOOLS_SERDE = PydanticTypeAdapter(RestateMCPGetToolsContextRunResult) - - -@dataclass -class RestateMCPToolRunResult: - """A simple wrapper for tool results to be used with Restate's run_typed.""" - - output: Any - - -MCP_RUN_SERDE = PydanticTypeAdapter(RestateMCPToolRunResult) - - class RestateContextRunToolSet(WrapperToolset[AgentDepsT]): """A toolset that automatically wraps tool calls with restate's `ctx.run_typed()`.""" @@ -114,65 +92,3 @@ def visit_and_replace( self, visitor: Callable[[AbstractToolset[AgentDepsT]], AbstractToolset[AgentDepsT]] ) -> AbstractToolset[AgentDepsT]: return visitor(self) - - -class RestateMCPToolset(WrapperToolset[AgentDepsT]): - """A wrapper for MCPToolset that integrates with Restate.""" - - def __init__(self, wrapped: MCPToolset[AgentDepsT], run_options: RunOptions): - super().__init__(wrapped) - self._wrapped = wrapped - self.get_tools_options = replace(run_options, serde=MCP_GET_TOOLS_SERDE) - self.run_tools_options = replace(run_options, serde=MCP_RUN_SERDE) - - def visit_and_replace( - self, visitor: Callable[[AbstractToolset[AgentDepsT]], AbstractToolset[AgentDepsT]] - ) -> AbstractToolset[AgentDepsT]: - return visitor(self) - - async def get_tools(self, ctx: RunContext[AgentDepsT]) -> dict[str, ToolsetTool[AgentDepsT]]: - async def get_tools_in_context() -> RestateMCPGetToolsContextRunResult: - res = await self._wrapped.get_tools(ctx) - # ToolsetTool is not serializable as it holds a SchemaValidator - # (which is also the same for every MCP tool so unnecessary to pass along the wire every time), - # so we just return the ToolDefinitions and wrap them in ToolsetTool outside of the activity. - return RestateMCPGetToolsContextRunResult(output={name: tool.tool_def for name, tool in res.items()}) - - context = current_context() - if context is None: - raise UserError( - "A toolset cannot be used without a Restate context. Make sure to run it within an agent or a run context." - ) - - try: - tool_defs = await context.run_typed("get mcp tools", get_tools_in_context, self.get_tools_options) - return {name: self.tool_for_tool_def(tool_def) for name, tool_def in tool_defs.output.items()} - except SdkInternalBaseException as e: - raise Exception("Internal error during get_tools call") from e - - def tool_for_tool_def(self, tool_def: ToolDefinition) -> ToolsetTool[AgentDepsT]: - assert isinstance(self.wrapped, MCPToolset) - return self.wrapped.tool_for_tool_def(tool_def) - - async def call_tool( - self, - name: str, - tool_args: dict[str, Any], - ctx: RunContext[AgentDepsT], - tool: ToolsetTool[AgentDepsT], - ) -> Any: - async def call_tool_in_context() -> RestateMCPToolRunResult: - res = await self._wrapped.call_tool(name, tool_args, ctx, tool) - return RestateMCPToolRunResult(output=res) - - context = current_context() - if context is None: - raise UserError( - "A toolset cannot be used without a Restate context. Make sure to run it within an agent or a run context." - ) - try: - res = await context.run_typed(f"Calling mcp tool {name}", call_tool_in_context, self.run_tools_options) - except SdkInternalBaseException as e: - raise Exception("Internal error during tool call") from e - - return res.output diff --git a/tests/ext_pydantic.py b/tests/ext_pydantic.py new file mode 100644 index 0000000..3ab52df --- /dev/null +++ b/tests/ext_pydantic.py @@ -0,0 +1,61 @@ +# +# Copyright (c) 2023-2026 - 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 +# +import subprocess +import sys + +import pytest + +pytest.importorskip("pydantic_ai") + + +def run_without_mcp(code: str) -> subprocess.CompletedProcess[str]: + """Run code in a fresh interpreter where the MCP client packages can't be imported, + as with pydantic-ai-slim installed without its `mcp` extra.""" + prelude = 'import sys\nsys.modules["mcp"] = None\nsys.modules["fastmcp"] = None\n' + return subprocess.run([sys.executable, "-c", prelude + code], capture_output=True, text=True) + + +def test_import_without_mcp(): + result = run_without_mcp("import restate.ext.pydantic") + assert result.returncode == 0, result.stderr + + +def test_agent_with_function_tools_without_mcp(): + result = run_without_mcp( + """ +from pydantic_ai import Agent +from pydantic_ai.models.test import TestModel +from restate.ext.pydantic import RestateAgent, RestateContextRunToolSet + +agent = Agent(TestModel()) + +@agent.tool_plain +def add(a: int, b: int) -> int: + return a + b + +restate_agent = RestateAgent(agent, auto_wrap_tools=True) +assert any(isinstance(t, RestateContextRunToolSet) for t in restate_agent.toolsets), restate_agent.toolsets +""" + ) + assert result.returncode == 0, result.stderr + + +def test_agent_wraps_mcp_toolsets(): + mcp = pytest.importorskip("pydantic_ai.mcp", exc_type=ImportError) + from pydantic_ai import Agent + from pydantic_ai.models.test import TestModel + + from restate.ext.pydantic import RestateAgent + from restate.ext.pydantic._mcp_toolset import RestateMCPToolset + + agent = Agent(TestModel(), toolsets=[mcp.MCPToolset("http://localhost:8000/mcp")]) + restate_agent = RestateAgent(agent) + assert any(isinstance(t, RestateMCPToolset) for t in restate_agent.toolsets)