cleanup: use tool conversion from mcp-run client (#2)

Author: zach

mcpx_pydantic_ai.py

@@ -4,29 +4,14 @@
 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"]
 
 
-def _convert_type(t):
-    if t == "string":
-        return str
-    elif t == "boolean":
-        return bool
-    elif t == "number":
-        return float
-    elif t == "integer":
-        return int
-    elif t == "object":
-        return dict
-    elif t == "array":
-        return list
-    raise TypeError(f"Unhandled conversion type: {t}")
-
-
 def openai_compatible_model(url: str, model: str, api_key: str | None = None):
     """
     Returns an OpenAI compatible model from the provided `url`, `model` name and optional `api_key`
@@ -71,27 +56,20 @@ def register_tool(self, tool: mcp_run.Tool, f=None):
             return
 
         def wrap(tool, inner):
-            props = tool.input_schema["properties"]
-            t = {k: _convert_type(v["type"]) for k, v in props.items()}
-            InputType = TypedDict("Input", t)
-
             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)}"
-            else:
-
-                def f(input: InputType):
-                    try:
-                        res = self.client.call_tool(tool=tool.name, params=input)
-                        return res.content[0].text
-                    except Exception as exc:
-                        return f"ERROR call to tool {tool.name} failed: {traceback.format_exception(exc)}"
 
-            return f
+                return f
+            else:
+                return self.client._make_pydantic_function(tool)
 
         self._register_tool(
             pydantic_ai.Tool(

pyproject.toml

@@ -1,11 +1,11 @@
 [project]
 name = "mcpx-pydantic-ai"
-version = "0.5.0"
+version = "0.6.0"
 description = "Pydantic Agent with mcp.run tools"
 readme = "README.md"
 requires-python = ">=3.12"
 dependencies = [
-    "mcp-run>=0.3.0",
+    "mcp-run>=0.4.0",
     "pydantic>=2.10.4",
     "pydantic-ai>=0.0.35",
 ]

tests/test_mcpx.py

@@ -2,6 +2,9 @@
 from unittest.mock import Mock, patch
 from mcpx_pydantic_ai import Agent, _convert_type
 from typing import Dict, Any
+import os
+
+os.environ["ANTHROPIC_API_KEY"] = "something"
 
 
 class MockTool:
@@ -41,6 +44,11 @@ def call_tool(self, tool: str, params: Dict[str, Any]) -> MockResponse:
     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
+
 
 class TestTypeConversion(unittest.TestCase):
     def test_convert_basic_types(self):