fix(pydantic): import restate.ext.pydantic without the mcp extra

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:

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

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

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)