mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
shsha4andCopybara-Service
86dea5b53a
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
90 lines
2.9 KiB
Python
90 lines
2.9 KiB
Python
# 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"
|