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:
shsha4
2025-09-23 10:52:12 -07:00
committed by Copybara-Service
parent 6ca2aee829
commit 86dea5b53a
3 changed files with 97 additions and 2 deletions
+1 -1
View File
@@ -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(
+7 -1
View File
@@ -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
+89
View File
@@ -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"