8000 feat: Support adding a name prefix to MCPTool and MCPToolset · google/adk-python@c055ecf · GitHub
[go: up one dir, main page]
Skip to content

Commit c055ecf

Browse files
calvingilescalvingiles
committed
feat: Support adding a name prefix to MCPTool and MCPToolset
1 parent f46b73b commit c055ecf

File tree

4 files changed

+103
-6
lines changed

4 files changed

+103
-6
lines changed

‎src/google/adk/tools/mcp_tool/mcp_tool.py

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,7 @@ def __init__(
6262
mcp_session_manager: MCPSessionManager,
6363
auth_scheme: Optional[AuthScheme] = None,
6464
auth_credential: Optional[AuthCredential] = None,
65+
tool_name_prefix: str = "",
6566
):
6667
"""Initializes a MCPTool.
6768
@@ -73,6 +74,8 @@ def __init__(
7374
mcp_session_manager: The MCP session manager to use for communication.
7475
auth_scheme: The authentication scheme to use.
7576
auth_credential: The authentication credential to use.
77+
tool_name_prefix: string to add to the start of the tool name. For example,
78+
`prefix="ns_"` would name `my_tool` as `ns_my_tool`.
7679
7780
Raises:
7881
ValueError: If mcp_tool or mcp_session_manager is None.
@@ -81,15 +84,19 @@ def __init__(
8184
raise ValueError("mcp_tool cannot be None")
8285
if mcp_session_manager is None:
8386
raise ValueError("mcp_session_manager cannot be None")
87+
raw_name = mcp_tool.name
88+
name = tool_name_prefix + raw_name
8489
super().__init__(
85-
name=mcp_tool.name,
90+
name=name,
8691
description=mcp_tool.description if mcp_tool.description else "",
8792
)
8893
self._mcp_tool = mcp_tool
8994
self._mcp_session_manager = mcp_session_manager
9095
# TODO(cheliu): Support passing auth to MCP Server.
9196
self._auth_scheme = auth_scheme
9297
self._auth_credential = auth_credential
98+
self._tool_name_prefix = tool_name_prefix
99+
self._raw_name = raw_name
93100

94101
@override
95102
def _get_declaration(self) -> FunctionDeclaration:

‎src/google/adk/tools/mcp_tool/mcp_toolset.py

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -36,8 +36,8 @@
3636
# Attempt to import MCP Tool from the MCP library, and hints user to upgrade
3737
# their Python version to 3.10 if it fails.
3838
try:
39-
from mcp import StdioServerParameters
4039
from mcp.types import ListToolsResult
40+
from mcp import StdioServerParameters
4141
except ImportError as e:
4242
import sys
4343

@@ -68,7 +68,8 @@ class MCPToolset(BaseToolset):
6868
command='npx',
6969
args=["-y", "@modelcontextprotocol/server-filesystem"],
7070
),
71-
tool_filter=['read_file', 'list_directory'] # Optional: filter specific tools
71+
tool_filter=['read_file', 'list_directory'], # Optional: filter specific tools
72+
tool_name_prefix="sfs_", # Optional: add_name_prefix
7273
)
7374
7475
# Use in an agent
@@ -96,6 +97,7 @@ def __init__(
9697
],
9798
tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
9899
errlog: TextIO = sys.stderr,
100+
tool_name_prefix: str = "",
99101
):
100102
"""Initializes the MCPToolset.
101103
@@ -108,10 +110,15 @@ def __init__(
108110
mcp server (e.g. using `npx` or `python3` ), but it does not support
109111
timeout, and we recommend to use `StdioConnectionParams` instead when
110112
timeout is needed.
111-
tool_filter: Optional filter to select specific tools. Can be either: - A
112-
list of tool names to include - A ToolPredicate function for custom
113-
filtering logic
113+
tool_filter: Optional filter to select specific tools. Can be either:
114+
- A list of tool names to include
115+
- A ToolPredicate function for custom filtering logic
116+
In both cases, the tool name WILL include the `tool_name_prefix` when
117+
matching.
114118
errlog: TextIO stream for error logging.
119+
tool_name_prefix: string to add to the start of the name of all return tools.
120+
For example, `prefix="ns_"` would change a returned tool name from
121+
`my_tool` to `ns_my_tool`.
115122
"""
116123
super().__init__(tool_filter=tool_filter)
117124

@@ -120,6 +127,7 @@ def __init__(
120127

121128
self._connection_params = connection_params
122129
self._errlog = errlog
130+
self._tool_name_prefix = tool_name_prefix
123131

124132
# Create the session manager that will handle the MCP connection
125133
self._mcp_session_manager = MCPSessionManager(

‎tests/unittests/tools/mcp_tool/__init__.py

Whitespace-only changes.

‎tests/unittests/tools/mcp_tool/test_mcp_tool.py

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
from google.adk import Agent
2+
from google.adk.tools import FunctionTool
3+
from google.adk.tools.mcp_tool import MCPTool
4+
from tests.unittests import testing_utils
5+
from google.genai.types import Part
6+
from mcp.types import Tool as McpBaseTool
7+
8+
9+
class TestMCPSession(object):
10+
11+
def __init__(self, function_tool: FunctionTool):
12+
self._function_tool = function_tool
13+
14+
async def 72A1 call_tool(self, name, arguments):
15+
return self._function_tool.func(**arguments)
16+
17+
18+
class TestMCPSessionManager(object):
19+
20+
def __init__(self, function_tool: FunctionTool):
21+
self._function_tool = function_tool
22+
23+
async def create_session(self):
24+
return TestMCPSession(self._function_tool)
25+
26+
async def close(self):
27+
pass
28+
29+
30+
def mcp_tool(function_tool: FunctionTool, prefix=''):
31+
return MCPTool(
32+
mcp_tool=McpBaseTool(
33+
name=function_tool.name,
34+
description=function_tool.description,
35+
inputSchema=function_tool._get_declaration().parameters.json_schema.model_dump(
36+
exclude_none=True
37+
),
38+
),
39+
mcp_session_manager=TestMCPSessionManager(function_tool),
40+
tool_name_prefix=prefix,
41+
)
42+
43+
44+
def test_mcp_tool():
45+
@FunctionTool
46+
def add(a: int, b: int):
47+
"""Add a and b and retuirn the result"""
48+
return a + b
49+
50+
mcp_add = mcp_tool(add, 'mcp_')
51+
52+
add_call = Part.from_function_call(name='add', args={'a': 1, 'b': 2})
53+
add_response = Part.from_function_response(name='add', response={'result': 3})
54+
55+
mcp_add_call = Part.from_function_call(name='mcp_add', args={'a': 5, 'b': 10})
56+
mcp_add_response = Part.from_function_response(
57+
name='mcp_add', response={'result': 15}
58+
)
59+
60+
mock_model = testing_utils.MockModel.create(
61+
responses=[
62+
add_call,
63+
mcp_add_call,
64+
'response1',
65+
]
66+
)
67+
68+
root_agent = Agent(
69+
name='root_agent',
70+
model=mock_model,
71+
tools=[add, mcp_add],
72+
)
73+
74+
runner = testing_utils.InMemoryRunner(root_agent)
75+
76+
assert testing_utils.simplify_events(runner.run('test1')) == [
77+
('root_agent', add_call),
78+
('root_agent', add_response),
79+
('root_agent', mcp_add_call),
80+
('root_agent', mcp_add_response),
81+
('root_agent', 'response1'),
82+
]

0 commit comments

Comments
 (0)
0