Skip to content

Commit 6f407d0

Browse files
Rehan Ahmedclaude
andcommitted
feat(mcp): add stdio_mcp_tools and http_mcp_tools convenience adapters
Add async context managers that handle the full MCP session lifecycle (connect, initialize, list/filter tools) and yield ready-to-use dspy.Tool lists. Supports include_tools/exclude_tools filtering, unknown tool name validation, and read_timeout propagation via ClientSession. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1 parent 90a9dd8 commit 6f407d0

2 files changed

Lines changed: 272 additions & 4 deletions

File tree

‎dspy/utils/mcp.py‎

Lines changed: 125 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,15 @@
1-
from typing import TYPE_CHECKING, Any
1+
from __future__ import annotations
2+
3+
from contextlib import asynccontextmanager
4+
from typing import TYPE_CHECKING, Any, AsyncIterator
25

36
from dspy.adapters.types.tool import Tool, convert_input_schema_to_tool_args
47

58
if TYPE_CHECKING:
69
import mcp
710

811

9-
def _convert_mcp_tool_result(call_tool_result: "mcp.types.CallToolResult") -> str | list[Any]:
12+
def _convert_mcp_tool_result(call_tool_result: mcp.types.CallToolResult) -> str | list[Any]:
1013
from mcp.types import TextContent
1114

1215
text_contents: list[TextContent] = []
@@ -27,7 +30,7 @@ def _convert_mcp_tool_result(call_tool_result: "mcp.types.CallToolResult") -> st
2730
return tool_content or non_text_contents
2831

2932

30-
def convert_mcp_tool(session: "mcp.ClientSession", tool: "mcp.types.Tool") -> Tool:
33+
def convert_mcp_tool(session: mcp.ClientSession, tool: mcp.types.Tool) -> Tool:
3134
"""Build a DSPy tool from an MCP tool.
3235
3336
Args:
@@ -45,3 +48,122 @@ async def func(*args, **kwargs):
4548
return _convert_mcp_tool_result(result)
4649

4750
return Tool(func=func, name=tool.name, desc=tool.description, args=args, arg_types=arg_types, arg_desc=arg_desc)
51+
52+
53+
def _validate_tool_filters(include_tools: list[str] | None, exclude_tools: list[str] | None) -> None:
54+
if include_tools and exclude_tools:
55+
raise ValueError("Cannot specify both include_tools and exclude_tools.")
56+
57+
58+
def _validate_tool_names(requested: list[str] | None, available: list[str]) -> None:
59+
if not requested:
60+
return
61+
unknown = set(requested) - set(available)
62+
if unknown:
63+
raise ValueError(f"Tool names not found on MCP server: {unknown}. Available tools: {set(available)}")
64+
65+
66+
def _filter_mcp_tools(
67+
tools: list[mcp.types.Tool],
68+
include_tools: list[str] | None,
69+
exclude_tools: list[str] | None,
70+
) -> list[mcp.types.Tool]:
71+
if include_tools is not None:
72+
tools = [t for t in tools if t.name in set(include_tools)]
73+
elif exclude_tools is not None:
74+
excluded = set(exclude_tools)
75+
tools = [t for t in tools if t.name not in excluded]
76+
return tools
77+
78+
79+
@asynccontextmanager
80+
async def stdio_mcp_tools(
81+
server_params: mcp.StdioServerParameters,
82+
*,
83+
include_tools: list[str] | None = None,
84+
exclude_tools: list[str] | None = None,
85+
timeout: float | None = None,
86+
) -> AsyncIterator[list[Tool]]:
87+
"""Yield DSPy tools backed by a live MCP stdio session.
88+
89+
Usage::
90+
91+
params = StdioServerParameters(command="docker", args=[...])
92+
async with stdio_mcp_tools(params, exclude_tools=["dangerous"], timeout=30) as tools:
93+
react = dspy.ReAct(MySignature, tools=tools)
94+
result = await react.acall(...)
95+
96+
Args:
97+
server_params: MCP stdio server parameters.
98+
include_tools: If set, only include tools whose names are in this list.
99+
exclude_tools: If set, exclude tools whose names are in this list.
100+
timeout: Read timeout in seconds for the MCP session. None means no timeout.
101+
102+
Yields:
103+
A list of DSPy Tool objects.
104+
105+
Raises:
106+
ValueError: If both include_tools and exclude_tools are specified, or if any
107+
requested tool names are not found on the MCP server.
108+
"""
109+
_validate_tool_filters(include_tools, exclude_tools)
110+
111+
from datetime import timedelta
112+
113+
from mcp import ClientSession
114+
from mcp.client.stdio import stdio_client
115+
116+
read_timeout = timedelta(seconds=timeout) if timeout is not None else None
117+
async with stdio_client(server_params) as (read, write):
118+
async with ClientSession(read, write, read_timeout_seconds=read_timeout) as session:
119+
await session.initialize()
120+
listed = await session.list_tools()
121+
_validate_tool_names(include_tools or exclude_tools, [t.name for t in listed.tools])
122+
filtered = _filter_mcp_tools(listed.tools, include_tools, exclude_tools)
123+
yield [convert_mcp_tool(session, t) for t in filtered]
124+
125+
126+
@asynccontextmanager
127+
async def http_mcp_tools(
128+
url: str,
129+
*,
130+
include_tools: list[str] | None = None,
131+
exclude_tools: list[str] | None = None,
132+
timeout: float | None = None,
133+
) -> AsyncIterator[list[Tool]]:
134+
"""Yield DSPy tools backed by a live MCP streamable-HTTP session.
135+
136+
Usage::
137+
138+
async with http_mcp_tools("http://127.0.0.1:8003/mcp", include_tools=["search"], timeout=30) as tools:
139+
react = dspy.ReAct(MySignature, tools=tools)
140+
result = await react.acall(...)
141+
142+
Args:
143+
url: The URL of the MCP streamable-HTTP server.
144+
include_tools: If set, only include tools whose names are in this list.
145+
exclude_tools: If set, exclude tools whose names are in this list.
146+
timeout: Read timeout in seconds for the MCP session. None means no timeout.
147+
148+
Yields:
149+
A list of DSPy Tool objects.
150+
151+
Raises:
152+
ValueError: If both include_tools and exclude_tools are specified, or if any
153+
requested tool names are not found on the MCP server.
154+
"""
155+
_validate_tool_filters(include_tools, exclude_tools)
156+
157+
from datetime import timedelta
158+
159+
from mcp import ClientSession
160+
from mcp.client.streamable_http import streamable_http_client
161+
162+
read_timeout = timedelta(seconds=timeout) if timeout is not None else None
163+
async with streamable_http_client(url) as (read, write, _):
164+
async with ClientSession(read, write, read_timeout_seconds=read_timeout) as session:
165+
await session.initialize()
166+
listed = await session.list_tools()
167+
_validate_tool_names(include_tools or exclude_tools, [t.name for t in listed.tools])
168+
filtered = _filter_mcp_tools(listed.tools, include_tools, exclude_tools)
169+
yield [convert_mcp_tool(session, t) for t in filtered]

‎tests/utils/test_mcp.py‎

Lines changed: 147 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@
33

44
import pytest
55

6-
from dspy.utils.mcp import convert_mcp_tool
6+
from dspy.utils.mcp import convert_mcp_tool, http_mcp_tools, stdio_mcp_tools
77

88
if importlib.util.find_spec("mcp") is None:
99
pytest.skip(reason="mcp is not installed", allow_module_level=True)
@@ -14,6 +14,7 @@
1414
async def test_convert_mcp_tool():
1515
from mcp import ClientSession, StdioServerParameters
1616
from mcp.client.stdio import stdio_client
17+
1718
server_params = StdioServerParameters(
1819
command="python",
1920
args=["tests/utils/resources/mcp_server.py"],
@@ -96,3 +97,148 @@ async def test_convert_mcp_tool():
9697
assert current_datetime_tool.arg_types == {}
9798
assert current_datetime_tool.arg_desc == {}
9899
assert await current_datetime_tool.acall() == "2025-07-23T09:10:10.0+00:00"
100+
101+
102+
def _find_nested_exception(eg, exc_type):
103+
"""Walk a (possibly nested) ExceptionGroup to find an exception of the given type."""
104+
if isinstance(eg, exc_type):
105+
return eg
106+
if isinstance(eg, BaseException) and hasattr(eg, "exceptions"):
107+
for exc in eg.exceptions:
108+
found = _find_nested_exception(exc, exc_type)
109+
if found:
110+
return found
111+
return None
112+
113+
114+
@pytest.fixture()
115+
def server_params():
116+
from mcp import StdioServerParameters
117+
118+
return StdioServerParameters(
119+
command="python",
120+
args=["tests/utils/resources/mcp_server.py"],
121+
env=None,
122+
)
123+
124+
125+
@pytest.mark.asyncio
126+
@pytest.mark.extra
127+
async def test_stdio_mcp_tools(server_params):
128+
# Basic usage: all tools available and callable
129+
async with stdio_mcp_tools(server_params) as tools:
130+
names = [t.name for t in tools]
131+
assert "add" in names
132+
assert "hello" in names
133+
assert await tools[0].acall(a=1, b=2) == "3"
134+
135+
# include_tools: only the specified tool is returned
136+
async with stdio_mcp_tools(server_params, include_tools=["add"]) as tools:
137+
assert len(tools) == 1
138+
assert tools[0].name == "add"
139+
140+
# exclude_tools: the specified tool is filtered out
141+
async with stdio_mcp_tools(server_params, exclude_tools=["wrong_tool"]) as tools:
142+
names = [t.name for t in tools]
143+
assert "wrong_tool" not in names
144+
assert "add" in names
145+
146+
# include_tools and exclude_tools together raises ValueError
147+
with pytest.raises(ValueError, match="Cannot specify both include_tools and exclude_tools"):
148+
async with stdio_mcp_tools(server_params, include_tools=["add"], exclude_tools=["hello"]) as _:
149+
pass
150+
151+
# timeout: session works with a timeout set
152+
async with stdio_mcp_tools(server_params, timeout=10) as tools:
153+
assert len(tools) > 0
154+
155+
# include_tools with unknown name raises ValueError
156+
try:
157+
async with stdio_mcp_tools(server_params, include_tools=["nonexistent"]) as _:
158+
pass
159+
pytest.fail("Expected ValueError")
160+
except BaseException as eg:
161+
err = _find_nested_exception(eg, ValueError)
162+
assert err and "Tool names not found on MCP server" in str(err)
163+
164+
# exclude_tools with unknown name raises ValueError
165+
try:
166+
async with stdio_mcp_tools(server_params, exclude_tools=["nonexistent"]) as _:
167+
pass
168+
pytest.fail("Expected ValueError")
169+
except BaseException as eg:
170+
err = _find_nested_exception(eg, ValueError)
171+
assert err and "Tool names not found on MCP server" in str(err)
172+
173+
174+
@pytest.fixture()
175+
def mcp_proxy_url():
176+
"""Start mcp-proxy as a subprocess, yield the URL, and tear down on exit."""
177+
import socket
178+
import subprocess
179+
import sys
180+
import time
181+
182+
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
183+
s.bind(("", 0))
184+
port = s.getsockname()[1]
185+
186+
url = f"http://127.0.0.1:{port}/mcp"
187+
proc = subprocess.Popen(
188+
["mcp-proxy", "--port", str(port), "--", sys.executable, "tests/utils/resources/mcp_server.py"],
189+
stdout=subprocess.PIPE,
190+
stderr=subprocess.PIPE,
191+
)
192+
try:
193+
for _ in range(50):
194+
try:
195+
with socket.create_connection(("127.0.0.1", port), timeout=0.1):
196+
break
197+
except OSError:
198+
time.sleep(0.1)
199+
else:
200+
raise TimeoutError(f"mcp-proxy did not start on port {port}")
201+
yield url
202+
finally:
203+
proc.terminate()
204+
proc.wait(timeout=5)
205+
206+
207+
@pytest.mark.asyncio
208+
@pytest.mark.extra
209+
async def test_http_mcp_tools(mcp_proxy_url):
210+
# Basic usage: all tools available and callable
211+
async with http_mcp_tools(mcp_proxy_url) as tools:
212+
names = [t.name for t in tools]
213+
assert "add" in names
214+
assert "hello" in names
215+
assert await tools[0].acall(a=1, b=2) == "3"
216+
217+
# include_tools: only the specified tool is returned
218+
async with http_mcp_tools(mcp_proxy_url, include_tools=["add"]) as tools:
219+
assert len(tools) == 1
220+
assert tools[0].name == "add"
221+
222+
# exclude_tools: the specified tool is filtered out
223+
async with http_mcp_tools(mcp_proxy_url, exclude_tools=["wrong_tool"]) as tools:
224+
names = [t.name for t in tools]
225+
assert "wrong_tool" not in names
226+
assert "add" in names
227+
228+
# include_tools with unknown name raises ValueError
229+
try:
230+
async with http_mcp_tools(mcp_proxy_url, include_tools=["nonexistent"]) as _:
231+
pass
232+
pytest.fail("Expected ValueError")
233+
except BaseException as eg:
234+
err = _find_nested_exception(eg, ValueError)
235+
assert err and "Tool names not found on MCP server" in str(err)
236+
237+
# exclude_tools with unknown name raises ValueError
238+
try:
239+
async with http_mcp_tools(mcp_proxy_url, exclude_tools=["nonexistent"]) as _:
240+
pass
241+
pytest.fail("Expected ValueError")
242+
except BaseException as eg:
243+
err = _find_nested_exception(eg, ValueError)
244+
assert err and "Tool names not found on MCP server" in str(err)

0 commit comments

Comments
 (0)