support complex args in dspy mcp (#8142)

Author: chenmoneygithub

dspy/primitives/tool.py

@@ -65,33 +65,6 @@ def foo(x: int, y: str = "hello"):
 
         self._parse_function(func, arg_desc)
 
-    def _resolve_pydantic_schema(self, model: type[BaseModel]) -> dict:
-        """Recursively resolve Pydantic model schema, expanding all references."""
-        schema = model.model_json_schema()
-
-        # If there are no definitions to resolve, return the main schema
-        if "$defs" not in schema and "definitions" not in schema:
-            return schema
-
-        def resolve_refs(obj: Any) -> Any:
-            if not isinstance(obj, (dict, list)):
-                return obj
-
-            if isinstance(obj, dict):
-                if "$ref" in obj:
-                    ref_path = obj["$ref"].split("/")[-1]
-                    return resolve_refs(schema["$defs"][ref_path])
-                return {k: resolve_refs(v) for k, v in obj.items()}
-
-            # Must be a list
-            return [resolve_refs(item) for item in obj]
-
-        # Resolve all references in the main schema
-        resolved_schema = resolve_refs(schema)
-        # Remove the $defs key as it's no longer needed
-        resolved_schema.pop("$defs", None)
-        return resolved_schema
-
     def _parse_function(self, func: Callable, arg_desc: dict[str, str] = None):
         """Helper method that parses a function to extract the name, description, and args.
 
@@ -121,7 +94,7 @@ def _parse_function(self, func: Callable, arg_desc: dict[str, str] = None):
             origin = get_origin(v) or v
             if isinstance(origin, type) and issubclass(origin, BaseModel):
                 # Get json schema, and replace $ref with the actual schema
-                v_json_schema = self._resolve_pydantic_schema(v)
+                v_json_schema = resolve_json_schema_reference(v.model_json_schema())
                 args[k] = v_json_schema
             else:
                 args[k] = TypeAdapter(v).json_schema()
@@ -197,3 +170,29 @@ def from_mcp_tool(cls, session: "mcp.client.session.ClientSession", tool: "mcp.t
         from dspy.utils.mcp import convert_mcp_tool
 
         return convert_mcp_tool(session, tool)
+
+
+def resolve_json_schema_reference(schema: dict) -> dict:
+    """Recursively resolve json model schema, expanding all references."""
+
+    # If there are no definitions to resolve, return the main schema
+    if "$defs" not in schema and "definitions" not in schema:
+        return schema
+
+    def resolve_refs(obj: Any) -> Any:
+        if not isinstance(obj, (dict, list)):
+            return obj
+        if isinstance(obj, dict):
+            if "$ref" in obj:
+                ref_path = obj["$ref"].split("/")[-1]
+                return resolve_refs(schema["$defs"][ref_path])
+            return {k: resolve_refs(v) for k, v in obj.items()}
+
+        # Must be a list
+        return [resolve_refs(item) for item in obj]
+
+    # Resolve all references in the main schema
+    resolved_schema = resolve_refs(schema)
+    # Remove the $defs key as it's no longer needed
+    resolved_schema.pop("$defs", None)
+    return resolved_schema

dspy/utils/mcp.py

@@ -1,5 +1,6 @@
-from typing import Any, Tuple, Type, Union, TYPE_CHECKING
-from dspy.primitives.tool import Tool
+from typing import TYPE_CHECKING, Any, Tuple, Type, Union
+
+from dspy.primitives.tool import Tool, resolve_json_schema_reference
 
 if TYPE_CHECKING:
     import mcp
@@ -25,7 +26,11 @@ def _convert_input_schema_to_tool_args(
 
     required = schema.get("required", [])
 
+    defs = schema.get("$defs", {})
+
     for name, prop in properties.items():
+        if len(defs) > 0:
+            prop = resolve_json_schema_reference({"$defs": defs, **prop})
         args[name] = prop
         # MCP tools are validated through jsonschema using args, so arg_types are not strictly required.
         arg_types[name] = TYPE_MAPPING.get(prop.get("type"), Any)
@@ -37,9 +42,7 @@ def _convert_input_schema_to_tool_args(
 
 
 def _convert_mcp_tool_result(call_tool_result: "mcp.types.CallToolResult") -> Union[str, list[Any]]:
-    from mcp.types import (
-        TextContent,
-    )
+    from mcp.types import TextContent
 
     text_contents: list[TextContent] = []
     non_text_contents = []

tests/utils/resources/mcp_server.py

@@ -1,8 +1,19 @@
 from mcp.server.fastmcp import FastMCP
+from pydantic import BaseModel
 
 mcp = FastMCP("test")
 
 
+class Profile(BaseModel):
+    name: str
+    age: int
+
+
+class Account(BaseModel):
+    profile: Profile
+    account_id: str
+
+
 @mcp.tool()
 def add(a: int, b: int) -> int:
     """Add two numbers"""
@@ -14,10 +25,18 @@ def hello(names: list[str]) -> str:
     """Greet people"""
     return [f"Hello, {name}!" for name in names]
 
+
 @mcp.tool()
 def wrong_tool():
     """This tool raises an error"""
     raise ValueError("error!")
 
+
+@mcp.tool()
+def get_account_name(account: Account):
+    """This extracts the name from account"""
+    return account.profile.name
+
+
 if __name__ == "__main__":
     mcp.run()

tests/utils/test_mcp.py

@@ -48,3 +48,37 @@ async def test_convert_mcp_tool():
                 RuntimeError, match="Failed to call a MCP tool: Error executing tool wrong_tool: error!"
             ):
                 await error_tool.acall()
+
+            # Check nested Pydantic arg
+            nested_pydantic_tool = convert_mcp_tool(session, response.tools[3])
+
+            assert nested_pydantic_tool.name == "get_account_name"
+            assert nested_pydantic_tool.desc == "This extracts the name from account"
+            assert nested_pydantic_tool.args == {
+                "account": {
+                    "title": "Account",
+                    "type": "object",
+                    "required": ["profile", "account_id"],
+                    "properties": {
+                        "profile": {
+                            "title": "Profile",
+                            "type": "object",
+                            "properties": {
+                                "name": {"title": "Name", "type": "string"},
+                                "age": {"title": "Age", "type": "integer"},
+                            },
+                            "required": ["name", "age"],
+                        },
+                        "account_id": {"title": "Account Id", "type": "string"},
+                    },
+                }
+            }
+            account_in_json = {
+                "profile": {
+                    "name": "Bob",
+                    "age": 20,
+                },
+                "account_id": "123",
+            }
+            result = await nested_pydantic_tool.acall(account=account_in_json)
+            assert result == "Bob"