Add dspy.Tool.from_mcp_tool (#8130)

* Add util for using mcp tools as DSPy tools

Co-authored-by: ThanabordeeN <[email protected]>

* rename

* update dependency

* add from_mcp_tool

* fix circular import

* test fix

---------

Co-authored-by: ThanabordeeN <[email protected]>

Author: TomeHirata

dspy/primitives/tool.py

@@ -1,12 +1,14 @@
 import asyncio
 import inspect
-from typing import Any, Callable, Optional, get_origin, get_type_hints
+from typing import Any, Callable, Optional, get_origin, get_type_hints, TYPE_CHECKING
 
 from jsonschema import ValidationError, validate
 from pydantic import BaseModel, TypeAdapter, create_model
 
 from dspy.utils.callback import with_callbacks
 
+if TYPE_CHECKING:
+    import mcp
 
 class Tool:
     """Tool class.
@@ -176,3 +178,18 @@ async def acall(self, **kwargs):
         if not asyncio.iscoroutine(result):
             raise ValueError("You are calling `acall` on a non-async tool, please use `__call__` instead.")
         return await result
+    
+    @classmethod
+    def from_mcp_tool(cls, session: "mcp.client.session.ClientSession", tool: "mcp.types.Tool") -> "Tool":
+        """
+        Build a DSPy tool from an MCP tool and a ClientSession.
+
+        Args:
+            session: The MCP session to use.
+            tool: The MCP tool to convert.
+
+        Returns:
+            A Tool object.
+        """
+        from dspy.utils.mcp import convert_mcp_tool
+        return convert_mcp_tool(session, tool)

dspy/utils/mcp.py

@@ -0,0 +1,79 @@
+from typing import Any, Tuple, Type, Union, TYPE_CHECKING
+from dspy.primitives.tool import Tool
+
+if TYPE_CHECKING:
+    import mcp
+
+TYPE_MAPPING = {"string": str, "integer": int, "number": float, "boolean": bool, "array": list, "object": dict}
+
+
+def _convert_input_schema_to_tool_args(
+    schema: dict[str, Any],
+) -> Tuple[dict[str, Any], dict[str, Type], dict[str, str]]:
+    """Convert an input schema to tool arguments compatible with DSPy Tool.
+
+    Args:
+        schema: An input schema describing the tool's input parameters
+
+    Returns:
+        A tuple of (args, arg_types, arg_desc) for DSPy Tool definition.
+    """
+    args, arg_types, arg_desc = {}, {}, {}
+    properties = schema.get("properties", None)
+    if properties is None:
+        return args, arg_types, arg_desc
+
+    required = schema.get("required", [])
+
+    for name, prop in properties.items():
+        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)
+        arg_desc[name] = prop.get("description", "No description provided.")
+        if name in required:
+            arg_desc[name] += " (Required)"
+
+    return args, arg_types, arg_desc
+
+
+def _convert_mcp_tool_result(call_tool_result: "mcp.types.CallToolResult") -> Union[str, list[Any]]:
+    from mcp.types import (
+        TextContent,
+    )
+
+    text_contents: list[TextContent] = []
+    non_text_contents = []
+    for content in call_tool_result.content:
+        if isinstance(content, TextContent):
+            text_contents.append(content)
+        else:
+            non_text_contents.append(content)
+
+    tool_content = [content.text for content in text_contents]
+    if len(text_contents) == 1:
+        tool_content = tool_content[0]
+
+    if call_tool_result.isError:
+        raise RuntimeError(f"Failed to call a MCP tool: {tool_content}")
+
+    return tool_content or non_text_contents
+
+
+def convert_mcp_tool(session: "mcp.client.session.ClientSession", tool: "mcp.types.Tool") -> Tool:
+    """Build a DSPy tool from an MCP tool.
+
+    Args:
+        session: The MCP session to use.
+        tool: The MCP tool to convert.
+
+    Returns:
+        A dspy Tool object.
+    """
+    args, arg_types, arg_desc = _convert_input_schema_to_tool_args(tool.inputSchema)
+
+    # Convert the MCP tool and Session to a single async method
+    async def func(*args, **kwargs):
+        result = await session.call_tool(tool.name, arguments=kwargs)
+        return _convert_mcp_tool_result(result)
+
+    return Tool(func=func, name=tool.name, desc=tool.description, args=args, arg_types=arg_types, arg_desc=arg_desc)

pyproject.toml

@@ -62,7 +62,9 @@ dev = [
     "build>=1.0.3",
     "litellm>=1.60.3; sys_platform == 'win32'",
     "litellm[proxy]>=1.60.3; sys_platform != 'win32'",
+    "mcp>=1.5.0; python_version >= '3.10'",
 ]
+mcp = ["mcp; python_version >= '3.10'"]
 
 [tool.setuptools.packages.find]
 where = ["."]

tests/utils/resources/mcp_server.py

@@ -0,0 +1,23 @@
+from mcp.server.fastmcp import FastMCP
+
+mcp = FastMCP("test")
+
+
+@mcp.tool()
+def add(a: int, b: int) -> int:
+    """Add two numbers"""
+    return a + b
+
+
+@mcp.tool()
+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!")
+
+if __name__ == "__main__":
+    mcp.run()

tests/utils/test_mcp.py

@@ -0,0 +1,50 @@
+import pytest
+import asyncio
+
+from mcp import ClientSession, StdioServerParameters
+from mcp.client.stdio import stdio_client
+
+from dspy.utils.mcp import convert_mcp_tool
+
+
+@pytest.mark.asyncio
+async def test_convert_mcp_tool():
+    server_params = StdioServerParameters(
+        command="python",
+        args=["tests/utils/resources/mcp_server.py"],
+        env=None,
+    )
+    async with stdio_client(server_params) as (read, write):
+        async with ClientSession(read, write) as session:
+            await asyncio.wait_for(session.initialize(), timeout=5)
+            response = await session.list_tools()
+
+            # Check add
+            add_tool = convert_mcp_tool(session, response.tools[0])
+            assert add_tool.name == "add"
+            assert add_tool.desc == "Add two numbers"
+            assert add_tool.args == {"a": {"title": "A", "type": "integer"}, "b": {"title": "B", "type": "integer"}}
+            assert add_tool.arg_types == {"a": int, "b": int}
+            assert add_tool.arg_desc == {
+                "a": "No description provided. (Required)",
+                "b": "No description provided. (Required)",
+            }
+            assert await add_tool.acall(a=1, b=2) == "3"
+
+            # Check hello
+            hello_tool = convert_mcp_tool(session, response.tools[1])
+            assert hello_tool.name == "hello"
+            assert hello_tool.desc == "Greet people"
+            assert hello_tool.args == {"names": {"title": "Names", "type": "array", "items": {"type": "string"}}}
+            assert hello_tool.arg_types == {"names": list}
+            assert hello_tool.arg_desc == {"names": "No description provided. (Required)"}
+            assert await hello_tool.acall(names=["Bob", "Tom"]) == ["Hello, Bob!", "Hello, Tom!"]
+
+            # Check error handling
+            error_tool = convert_mcp_tool(session, response.tools[2])
+            assert error_tool.name == "wrong_tool"
+            assert error_tool.desc == "This tool raises an error"
+            with pytest.raises(
+                RuntimeError, match="Failed to call a MCP tool: Error executing tool wrong_tool: error!"
+            ):
+                await error_tool.acall()