You've already forked adk-python
mirror of
https://github.com/encounter/adk-python.git
synced 2026-03-30 10:57:20 -07:00
fix(mcp): Initialize tool_name_prefix in MCPToolse
Merge https://github.com/google/adk-python/pull/2823 Description This change introduces a tool_name_prefix attribute to McpToolset and McpToolsetConfig. This allows for adding a prefix to the names of all tools within the toolset, which can help avoid naming collisions and provide better organization. The implementation involves updating the McpToolset's __init__ and from_config methods to handle the new tool_name_prefix and adding the corresponding field to McpToolsetConfig. Testing Plan A new unit test file has been added to ensure the functionality works as expected. - `tests/unittests/tools/test_mcp_toolset.py`: - The test_mcp_toolset_with_prefix test case verifies that the tool_name_prefix is correctly applied to the tool names retrieved from the toolset. - All tests were run via pytest and passed. Related Issue - Closes #2814 COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/2823 from shsha4:fix/issue-2814 e8e5b0d6d5f406d3875faf2229a96701725b7a5e PiperOrigin-RevId: 810500616
This commit is contained in:
@@ -136,7 +136,7 @@ class McpTool(BaseAuthenticatedTool):
|
||||
# Get the session from the session manager
|
||||
session = await self._mcp_session_manager.create_session(headers=headers)
|
||||
|
||||
response = await session.call_tool(self.name, arguments=args)
|
||||
response = await session.call_tool(self._mcp_tool.name, arguments=args)
|
||||
return response
|
||||
|
||||
async def _get_headers(
|
||||
|
||||
@@ -100,6 +100,7 @@ class McpToolset(BaseToolset):
|
||||
StreamableHTTPConnectionParams,
|
||||
],
|
||||
tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
|
||||
tool_name_prefix: Optional[str] = None,
|
||||
errlog: TextIO = sys.stderr,
|
||||
auth_scheme: Optional[AuthScheme] = None,
|
||||
auth_credential: Optional[AuthCredential] = None,
|
||||
@@ -118,11 +119,13 @@ class McpToolset(BaseToolset):
|
||||
tool_filter: Optional filter to select specific tools. Can be either: - A
|
||||
list of tool names to include - A ToolPredicate function for custom
|
||||
filtering logic
|
||||
tool_name_prefix: A prefix to be added to the name of each tool in this
|
||||
toolset.
|
||||
errlog: TextIO stream for error logging.
|
||||
auth_scheme: The auth scheme of the tool for tool calling
|
||||
auth_credential: The auth credential of the tool for tool calling
|
||||
"""
|
||||
super().__init__(tool_filter=tool_filter)
|
||||
super().__init__(tool_filter=tool_filter, tool_name_prefix=tool_name_prefix)
|
||||
|
||||
if not connection_params:
|
||||
raise ValueError("Missing connection params in MCPToolset.")
|
||||
@@ -207,6 +210,7 @@ class McpToolset(BaseToolset):
|
||||
return cls(
|
||||
connection_params=connection_params,
|
||||
tool_filter=mcp_toolset_config.tool_filter,
|
||||
tool_name_prefix=mcp_toolset_config.tool_name_prefix,
|
||||
auth_scheme=mcp_toolset_config.auth_scheme,
|
||||
auth_credential=mcp_toolset_config.auth_credential,
|
||||
)
|
||||
@@ -239,6 +243,8 @@ class McpToolsetConfig(BaseToolConfig):
|
||||
|
||||
tool_filter: Optional[List[str]] = None
|
||||
|
||||
tool_name_prefix: Optional[str] = None
|
||||
|
||||
auth_scheme: Optional[AuthScheme] = None
|
||||
|
||||
auth_credential: Optional[AuthCredential] = None
|
||||
|
||||
@@ -0,0 +1,89 @@
|
||||
# Copyright 2025 Google LLC
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
"""Unit tests for McpToolset."""
|
||||
|
||||
import sys
|
||||
from unittest.mock import AsyncMock
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
# Skip all tests in this module if Python version is less than 3.10
|
||||
pytestmark = pytest.mark.skipif(
|
||||
sys.version_info < (3, 10), reason="MCP tool requires Python 3.10+"
|
||||
)
|
||||
|
||||
# Import dependencies with version checking
|
||||
try:
|
||||
from google.adk.tools.mcp_tool.mcp_toolset import McpToolset
|
||||
except ImportError as e:
|
||||
if sys.version_info < (3, 10):
|
||||
# Create dummy classes to prevent NameError during test collection
|
||||
# Tests will be skipped anyway due to pytestmark
|
||||
class DummyClass:
|
||||
pass
|
||||
|
||||
McpToolset = DummyClass
|
||||
else:
|
||||
raise e
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_toolset_with_prefix():
|
||||
"""Test that McpToolset correctly applies the tool_name_prefix."""
|
||||
# Mock the connection parameters
|
||||
mock_connection_params = MagicMock()
|
||||
|
||||
# Mock the MCPSessionManager and its create_session method
|
||||
mock_session_manager = MagicMock()
|
||||
mock_session = MagicMock()
|
||||
|
||||
# Mock the list_tools response from the MCP server
|
||||
mock_tool1 = MagicMock()
|
||||
mock_tool1.name = "tool1"
|
||||
mock_tool1.description = "tool 1 desc"
|
||||
mock_tool2 = MagicMock()
|
||||
mock_tool2.name = "tool2"
|
||||
mock_tool2.description = "tool 2 desc"
|
||||
list_tools_result = MagicMock()
|
||||
list_tools_result.tools = [mock_tool1, mock_tool2]
|
||||
mock_session.list_tools = AsyncMock(return_value=list_tools_result)
|
||||
mock_session_manager.create_session = AsyncMock(return_value=mock_session)
|
||||
|
||||
# Create an instance of McpToolset with a prefix
|
||||
toolset = McpToolset(
|
||||
connection_params=mock_connection_params,
|
||||
tool_name_prefix="my_prefix",
|
||||
)
|
||||
|
||||
# Replace the internal session manager with our mock
|
||||
toolset._mcp_session_manager = mock_session_manager
|
||||
|
||||
# Get the tools from the toolset
|
||||
tools = await toolset.get_tools()
|
||||
|
||||
# The get_tools method in McpToolset returns MCPTool objects, which are
|
||||
# instances of BaseTool. The prefixing is handled by the BaseToolset,
|
||||
# so we need to call get_tools_with_prefix to get the prefixed tools.
|
||||
prefixed_tools = await toolset.get_tools_with_prefix()
|
||||
|
||||
# Assert that the tools are prefixed correctly
|
||||
assert len(prefixed_tools) == 2
|
||||
assert prefixed_tools[0].name == "my_prefix_tool1"
|
||||
assert prefixed_tools[1].name == "my_prefix_tool2"
|
||||
|
||||
# Assert that the original tools are not modified
|
||||
assert tools[0].name == "tool1"
|
||||
assert tools[1].name == "tool2"
|
||||
Reference in New Issue
Block a user