refactor: use MCP protocol to access mcp.run via SSE/stdio (#4)

zach · web-flow · commit 4258a39a4b2d · 2025-05-08T15:15:15.000-07:00
* refactor: only use MCP for mcp.run tool calling
diff --git a/.python-version b/.python-version
diff --git a/example.py b/example.py
@@ -31,8 +31,9 @@ async def run(agent, msg):
     Send a message to an agent and return the result
     """
     global history
-    async with agent.run_stream(msg, message_history=history) as result:
-        return await result.get_data()
+    async with agent.run_mcp_servers():
+        async with agent.run_stream(msg, message_history=history) as result:
+            return await result.get_output()
 
 
 types = ImageList | VowelCount
diff --git a/mcpx_pydantic_ai.py b/mcpx_pydantic_ai.py
@@ -4,12 +4,21 @@
 from pydantic import BaseModel, Field
 from pydantic_ai.models.openai import OpenAIModel
 from pydantic_ai.providers.openai import OpenAIProvider
-from mcp_run.client import _convert_type
-
-from typing import TypedDict, List, Set, AsyncIterator, Any
-import traceback
-
-__all__ = ["BaseModel", "Field", "Agent", "mcp_run", "pydantic_ai", "pydantic"]
+from pydantic_ai.mcp import MCPServerHTTP, MCPServerStdio
+from datetime import timedelta
+from mcp_run import MCPClient, SSEClientConfig, StdioClientConfig
+
+__all__ = [
+    "BaseModel",
+    "Field",
+    "Agent",
+    "mcp_run",
+    "pydantic_ai",
+    "pydantic",
+    "MCPClient",
+    "SSEClientConfig",
+    "StdioClientConfig",
+]
 
 
 def openai_compatible_model(url: str, model: str, api_key: str | None = None):
@@ -26,103 +35,30 @@ class Agent(pydantic_ai.Agent):
     """
 
     client: mcp_run.Client
-    ignore_tools: Set[str]
-    _original_tools: list
-    _registered_tools: List[str]
 
     def __init__(
         self,
         *args,
         client: mcp_run.Client | None = None,
-        ignore_tools: List[str] | None = None,
+        mcp_client: MCPClient | None = None,
+        expires_in: timedelta | None = None,
         **kw,
     ):
         self.client = client or mcp_run.Client()
-        self._original_tools = kw.get("tools", [])
-        self._registered_tools = []
-        self.ignore_tools = set(ignore_tools or [])
-        super().__init__(*args, **kw)
-        self._update_tools()
-
-        for t in self._original_tools:
-            self._registered_tools.append(t.name)
-
-    def set_profile(self, profile: str):
-        self.client.set_profile(profile)
-        self._update_tools()
-
-    def register_tool(self, tool: mcp_run.Tool, f=None):
-        if tool.name in self.ignore_tools:
-            return
-
-        def wrap(tool, inner):
-            if inner is not None:
-                props = tool.input_schema["properties"]
-                t = {k: _convert_type(v["type"]) for k, v in props.items()}
-                InputType = TypedDict("Input", t)
-
-                def f(input: InputType):
-                    try:
-                        return inner(input)
-                    except Exception as exc:
-                        return f"ERROR call to tool {tool.name} failed: {traceback.format_exception(exc)}"
-
-                return f
-            else:
-                return self.client._make_pydantic_function(tool)
-
-        self._register_tool(
-            pydantic_ai.Tool(
-                wrap(tool, f),
-                name=tool.name,
-                description=tool.description,
-            )
+        mcp = mcp_client or self.client.mcp_sse(
+            profile=self.client.config.profile, expires_in=expires_in
         )
-
-        if f is not None:
-            self._registered_tools.append(tool.name)
-
-    def reset_tools(self):
-        for k in list(self._function_tools.keys()):
-            if k not in self._registered_tools:
-                del self._function_tools[k]
-
-    def _update_tools(self):
-        self.reset_tools()
-        for tool in self.client.tools.values():
-            self.register_tool(tool)
-
-    async def run(self, *args, update_tools: bool = True, **kw):
-        if update_tools:
-            self._update_tools()
-        return await super().run(*args, **kw)
-
-    def run_sync(self, *args, update_tools: bool = True, **kw):
-        if update_tools:
-            self._update_tools()
-        return super().run_sync(*args, **kw)
-
-    async def run_async(self, *args, update_tools: bool = True, **kw):
-        if update_tools:
-            self._update_tools()
-        return await super().run_async(*args, **kw)
-
-    def run_stream(
-        self,
-        *args,
-        update_tools: bool = True,
-        **kw,
-    ) -> AsyncIterator[Any]:
-        if update_tools:
-            self._update_tools()
-        return super().run_stream(*args, **kw)
-
-    def iter(
-        self,
-        *args,
-        update_tools: bool = True,
-        **kw,
-    ) -> AsyncIterator[Any]:
-        if update_tools:
-            self._update_tools()
-        return super().iter(*args, **kw)
+        mcp_servers = kw.get("mcp_servers", [])
+        if mcp.is_sse:
+            mcp_servers.append(MCPServerHTTP(url=mcp.config.url))
+        elif mcp.is_stdio:
+            mcp_servers.append(
+                MCPServerStdio(
+                    command=mcp.config.command,
+                    args=mcp.config.args,
+                    env=mcp.config.env,
+                    cwd=mcp.config.cwd,
+                )
+            )
+        kw["mcp_servers"] = mcp_servers
+        super().__init__(*args, **kw)
diff --git a/pyproject.toml b/pyproject.toml
@@ -1,11 +1,11 @@
 [project]
 name = "mcpx-pydantic-ai"
-version = "0.6.0"
+version = "0.7.0"
 description = "Pydantic Agent with mcp.run tools"
 readme = "README.md"
 requires-python = ">=3.12"
 dependencies = [
-    "mcp-run>=0.4.0",
+    "mcp-run>=0.5.0",
     "pydantic>=2.10.4",
     "pydantic-ai>=0.0.35",
 ]
diff --git a/tests/test_mcpx.py b/tests/test_mcpx.py
@@ -1,9 +1,11 @@
 import unittest
-from unittest.mock import Mock, patch
-from mcpx_pydantic_ai import Agent, _convert_type
+from unittest.mock import Mock
 from typing import Dict, Any
 import os
 
+
+from mcpx_pydantic_ai import Agent
+
 os.environ["ANTHROPIC_API_KEY"] = "something"
 
 
@@ -35,33 +37,28 @@ def __init__(self):
         }
         self.called_tool = None
         self.called_params = None
+        self.config = Mock(profile="default")
 
     def call_tool(self, tool: str, params: Dict[str, Any]) -> MockResponse:
         self.called_tool = tool
         self.called_params = params
         return MockResponse("mock response")
 
-    def set_profile(self, profile: str):
-        self.profile = profile
-
     def _make_pydantic_function(self, tool):
         def test(input: dict):
             return self.call_tool(tool.name, input).content[0].text
-        return test
 
+        return test
 
-class TestTypeConversion(unittest.TestCase):
-    def test_convert_basic_types(self):
-        self.assertEqual(_convert_type("string"), str)
-        self.assertEqual(_convert_type("boolean"), bool)
-        self.assertEqual(_convert_type("number"), float)
-        self.assertEqual(_convert_type("integer"), int)
-        self.assertEqual(_convert_type("object"), dict)
-        self.assertEqual(_convert_type("array"), list)
+    def mcp_sse(self, profile=None, expires_in=None):
+        mock_mcp = Mock()
+        mock_mcp.is_sse = True
+        mock_mcp.is_stdio = False
+        mock_mcp.config = Mock(url="http://mock-url.com")
+        return mock_mcp
 
-    def test_convert_invalid_type(self):
-        with self.assertRaises(TypeError):
-            _convert_type("invalid_type")
+    def _fix_profile(self, profile):
+        return profile
 
 
 class TestAgent(unittest.TestCase):
@@ -76,102 +73,6 @@ def setUp(self):
     def test_init_with_custom_client(self):
         """Test agent initialization with custom client"""
         self.assertEqual(self.agent.client, self.mock_client)
-        self.assertEqual(
-            len(self.agent._function_tools), 1
-        )  # Should have our mock tool
-
-    def test_init_with_ignore_tools(self):
-        """Test agent initialization with ignored tools"""
-        agent = Agent(
-            model="claude-3-5-sonnet-latest",
-            client=self.mock_client,
-            ignore_tools=["test_tool"],
-            system_prompt="test prompt",
-        )
-        self.assertEqual(
-            len(agent._function_tools), 0
-        )  # Should have no tools due to ignore
-
-    def test_set_profile(self):
-        """Test setting profile updates client profile"""
-        self.agent.set_profile("test_profile")
-        self.assertEqual(self.mock_client.profile, "test_profile")
-
-    def test_register_custom_tool(self):
-        """Test registering a custom tool with custom function"""
-        custom_mock = Mock(return_value="custom response")
-
-        self.agent.register_tool(
-            MockTool(
-                "custom_tool",
-                "A custom tool",
-                {"properties": {"param": {"type": "string"}}},
-            ),
-            custom_mock,
-        )
-
-        # Verify tool was registered
-        self.assertIn("custom_tool", self.agent._function_tools)
-
-        # Test tool execution
-        tool_func = self.agent._function_tools["custom_tool"].function
-        result = tool_func({"param": "test"})
-
-        custom_mock.assert_called_once_with({"param": "test"})
-        self.assertEqual(result, "custom response")
-
-    def test_tool_execution(self):
-        """Test executing a registered tool"""
-        # Our mock tool should be registered automatically
-        tool_func = self.agent._function_tools["test_tool"].function
-
-        result = tool_func({"param1": "test", "param2": 123})
-
-        self.assertEqual(self.mock_client.called_tool, "test_tool")
-        self.assertEqual(
-            self.mock_client.called_params, {"param1": "test", "param2": 123}
-        )
-        self.assertEqual(result, "mock response")
-
-    def test_reset_tools(self):
-        """Test resetting tools"""
-        # Add a custom tool
-        self.agent.register_tool(
-            MockTool(
-                "custom_tool",
-                "A custom tool",
-                {"properties": {"param": {"type": "string"}}},
-            ),
-            Mock(),
-        )
-
-        # Reset tools
-        self.agent.reset_tools()
-
-        # Only custom tool should remain
-        self.assertEqual(len(self.agent._function_tools), 1)
-        self.assertIn("custom_tool", self.agent._function_tools)
-        self.assertNotIn("test_tool", self.agent._function_tools)
-
-    @patch("mcpx_pydantic_ai.pydantic_ai.Agent.run_sync")
-    def test_run_sync_updates_tools(self, mock_run_sync):
-        """Test that run_sync updates tools by default"""
-        mock_run_sync.return_value = "test response"
-
-        result = self.agent.run_sync("test prompt")
-
-        self.assertEqual(result, "test response")
-        mock_run_sync.assert_called_once()
-
-    @patch("mcpx_pydantic_ai.pydantic_ai.Agent.run_sync")
-    def test_run_sync_without_tool_update(self, mock_run_sync):
-        """Test that run_sync can skip tool updates"""
-        mock_run_sync.return_value = "test response"
-
-        result = self.agent.run_sync("test prompt", update_tools=False)
-
-        self.assertEqual(result, "test response")
-        mock_run_sync.assert_called_once()
 
 
 if __name__ == "__main__":
diff --git a/uv.lock b/uv.lock
