From ebd726f1f5e0a76f383192cace4a80a83204325b Mon Sep 17 00:00:00 2001 From: "Xiang (Sean) Zhou" Date: Tue, 12 Aug 2025 16:17:21 -0700 Subject: [PATCH] feat: Support adding prefix to tool names returned by toolset This is to address the name conflict issue of tools returned by different toolset. Mainly it's to give each toolset a namespace. We have a flag `add_tool_name_prefix` to decide whether to apply this behavior We have a `tool_name_prefix` to let client specify a custom prefix, if not set , toolset name will be used as prefix. PiperOrigin-RevId: 794306796 --- src/google/adk/agents/llm_agent.py | 5 +- src/google/adk/tools/base_toolset.py | 67 ++++- tests/unittests/tools/test_base_toolset.py | 281 ++++++++++++++++++++- 3 files changed, 349 insertions(+), 4 deletions(-) diff --git a/src/google/adk/agents/llm_agent.py b/src/google/adk/agents/llm_agent.py index 7f753521bfa..ae55dd1e464 100644 --- a/src/google/adk/agents/llm_agent.py +++ b/src/google/adk/agents/llm_agent.py @@ -113,10 +113,11 @@ async def _convert_tool_union_to_tools( ) -> list[BaseTool]: if isinstance(tool_union, BaseTool): return [tool_union] - if isinstance(tool_union, Callable): + if callable(tool_union): return [FunctionTool(func=tool_union)] - return await tool_union.get_tools(ctx) + # At this point, tool_union must be a BaseToolset + return await tool_union.get_tools_with_prefix(ctx) class LlmAgent(BaseAgent): diff --git a/src/google/adk/tools/base_toolset.py b/src/google/adk/tools/base_toolset.py index 41ca1f118f8..3400be40f2d 100644 --- a/src/google/adk/tools/base_toolset.py +++ b/src/google/adk/tools/base_toolset.py @@ -16,6 +16,8 @@ from abc import ABC from abc import abstractmethod +import copy +from typing import final from typing import List from typing import Optional from typing import Protocol @@ -58,9 +60,19 @@ class BaseToolset(ABC): """ def __init__( - self, *, tool_filter: Optional[Union[ToolPredicate, List[str]]] = None + self, + *, + tool_filter: Optional[Union[ToolPredicate, List[str]]] = None, + tool_name_prefix: Optional[str] = None, ): + """Initialize the toolset. + + Args: + tool_filter: Filter to apply to tools. + tool_name_prefix: The prefix to prepend to the names of the tools returned by the toolset. + """ self.tool_filter = tool_filter + self.tool_name_prefix = tool_name_prefix @abstractmethod async def get_tools( @@ -77,6 +89,59 @@ async def get_tools( list[BaseTool]: A list of tools available under the specified context. """ + @final + async def get_tools_with_prefix( + self, + readonly_context: Optional[ReadonlyContext] = None, + ) -> list[BaseTool]: + """Return all tools with optional prefix applied to tool names. + + This method calls get_tools() and applies prefixing if tool_name_prefix is provided. + + Args: + readonly_context (ReadonlyContext, optional): Context used to filter tools + available to the agent. If None, all tools in the toolset are returned. + + Returns: + list[BaseTool]: A list of tools with prefixed names if tool_name_prefix is provided. + """ + tools = await self.get_tools(readonly_context) + + if not self.tool_name_prefix: + return tools + + prefix = self.tool_name_prefix + + # Create copies of tools to avoid modifying original instances + prefixed_tools = [] + for tool in tools: + # Create a shallow copy of the tool + tool_copy = copy.copy(tool) + + # Apply prefix to the copied tool + prefixed_name = f"{prefix}_{tool.name}" + tool_copy.name = prefixed_name + + # Also update the function declaration name if the tool has one + # Use default parameters to capture the current values in the closure + def _create_prefixed_declaration( + original_get_declaration=tool._get_declaration, + prefixed_name=prefixed_name, + ): + def _get_prefixed_declaration(): + declaration = original_get_declaration() + if declaration is not None: + declaration.name = prefixed_name + return declaration + return None + + return _get_prefixed_declaration + + tool_copy._get_declaration = _create_prefixed_declaration() + prefixed_tools.append(tool_copy) + + return prefixed_tools + async def close(self) -> None: """Performs cleanup and releases resources held by the toolset. diff --git a/tests/unittests/tools/test_base_toolset.py b/tests/unittests/tools/test_base_toolset.py index 5414bb3c899..20d7f9d8259 100644 --- a/tests/unittests/tools/test_base_toolset.py +++ b/tests/unittests/tools/test_base_toolset.py @@ -23,17 +23,29 @@ from google.adk.sessions.in_memory_session_service import InMemorySessionService from google.adk.tools.base_tool import BaseTool from google.adk.tools.base_toolset import BaseToolset +from google.adk.tools.function_tool import FunctionTool from google.adk.tools.tool_context import ToolContext import pytest +class _TestingTool(BaseTool): + """A test implementation of BaseTool.""" + + async def run_async(self, *, args, tool_context): + return 'test result' + + class _TestingToolset(BaseToolset): """A test implementation of BaseToolset.""" + def __init__(self, *args, tools: Optional[list[BaseTool]] = None, **kwargs): + super().__init__(*args, **kwargs) + self._tools = tools or [] + async def get_tools( self, readonly_context: Optional[ReadonlyContext] = None ) -> list[BaseTool]: - return [] + return self._tools async def close(self) -> None: pass @@ -107,3 +119,270 @@ async def process_llm_request( # Verify the custom processing was applied assert llm_request.contents == ['Custom processing applied'] + + +@pytest.mark.asyncio +async def test_prefix_functionality_disabled_by_default(): + """Test that prefix functionality is disabled by default.""" + tool1 = _TestingTool(name='tool1', description='Test tool 1') + tool2 = _TestingTool(name='tool2', description='Test tool 2') + toolset = _TestingToolset(tools=[tool1, tool2]) + + # When tool_name_prefix is None (default), get_tools_with_prefix should return original tools + prefixed_tools = await toolset.get_tools_with_prefix() + + assert len(prefixed_tools) == 2 + assert prefixed_tools[0].name == 'tool1' + assert prefixed_tools[1].name == 'tool2' + assert toolset.tool_name_prefix is None + + +@pytest.mark.asyncio +async def test_prefix_functionality_with_custom_prefix(): + """Test prefix functionality with custom prefix.""" + tool1 = _TestingTool(name='tool1', description='Test tool 1') + tool2 = _TestingTool(name='tool2', description='Test tool 2') + toolset = _TestingToolset(tools=[tool1, tool2], tool_name_prefix='custom') + + # Should use the provided prefix + prefixed_tools = await toolset.get_tools_with_prefix() + + assert len(prefixed_tools) == 2 + assert prefixed_tools[0].name == 'custom_tool1' + assert prefixed_tools[1].name == 'custom_tool2' + assert toolset.tool_name_prefix == 'custom' + + +@pytest.mark.asyncio +async def test_prefix_with_none_has_no_effect(): + """Test that when prefix is None, tools are returned unchanged.""" + tool1 = _TestingTool(name='tool1', description='Test tool 1') + tool2 = _TestingTool(name='tool2', description='Test tool 2') + toolset = _TestingToolset(tools=[tool1, tool2], tool_name_prefix=None) + + prefixed_tools = await toolset.get_tools_with_prefix() + + assert len(prefixed_tools) == 2 + assert prefixed_tools[0].name == 'tool1' + assert prefixed_tools[1].name == 'tool2' + assert toolset.tool_name_prefix is None + + +@pytest.mark.asyncio +async def test_prefix_with_empty_string(): + """Test prefix functionality with empty string prefix.""" + tool1 = _TestingTool(name='tool1', description='Test tool 1') + toolset = _TestingToolset(tools=[tool1], tool_name_prefix='') + + prefixed_tools = await toolset.get_tools_with_prefix() + + # Empty prefix should be treated as no prefix + assert len(prefixed_tools) == 1 + assert prefixed_tools[0].name == 'tool1' + assert toolset.tool_name_prefix == '' + + +@pytest.mark.asyncio +async def test_prefix_assignment(): + """Test that prefix is properly assigned.""" + toolset = _TestingToolset(tool_name_prefix='explicit') + assert toolset.tool_name_prefix == 'explicit' + + # Test None assignment + toolset_none = _TestingToolset(tool_name_prefix=None) + assert toolset_none.tool_name_prefix is None + + +@pytest.mark.asyncio +async def test_prefix_creates_tool_copies(): + """Test that prefixing creates copies and preserves original tools.""" + original_tool = _TestingTool( + name='original', description='Original description' + ) + original_tool.is_long_running = True + original_tool.custom_attribute = 'custom_value' + + toolset = _TestingToolset(tools=[original_tool], tool_name_prefix='test') + prefixed_tools = await toolset.get_tools_with_prefix() + + prefixed_tool = prefixed_tools[0] + + # Name should be prefixed in the copy + assert prefixed_tool.name == 'test_original' + + # Other attributes should be preserved + assert prefixed_tool.description == 'Original description' + assert prefixed_tool.is_long_running == True + assert prefixed_tool.custom_attribute == 'custom_value' + + # Original tool should remain unchanged + assert original_tool.name == 'original' + assert original_tool is not prefixed_tool + + +@pytest.mark.asyncio +async def test_get_tools_vs_get_tools_with_prefix(): + """Test that get_tools returns tools without prefixing.""" + tool1 = _TestingTool(name='test_tool1', description='Test tool 1') + tool2 = _TestingTool(name='test_tool2', description='Test tool 2') + toolset = _TestingToolset(tools=[tool1, tool2], tool_name_prefix='prefix') + + # get_tools should return original tools (unmodified) + original_tools = await toolset.get_tools() + assert len(original_tools) == 2 + assert original_tools[0].name == 'test_tool1' + assert original_tools[1].name == 'test_tool2' + + # Now calling get_tools_with_prefix should return prefixed copies + prefixed_tools = await toolset.get_tools_with_prefix() + assert len(prefixed_tools) == 2 + assert prefixed_tools[0].name == 'prefix_test_tool1' + assert prefixed_tools[1].name == 'prefix_test_tool2' + + # Original tools should remain unchanged + assert original_tools[0].name == 'test_tool1' + assert original_tools[1].name == 'test_tool2' + + # The prefixed tools should be different instances + assert prefixed_tools[0] is not original_tools[0] + assert prefixed_tools[1] is not original_tools[1] + + +@pytest.mark.asyncio +async def test_empty_toolset_with_prefix(): + """Test prefix functionality with empty toolset.""" + toolset = _TestingToolset(tools=[], tool_name_prefix='test') + + prefixed_tools = await toolset.get_tools_with_prefix() + assert len(prefixed_tools) == 0 + + +@pytest.mark.asyncio +async def test_function_declarations_are_prefixed(): + """Test that function declarations have prefixed names.""" + + def test_function(param1: str, param2: int) -> str: + """A test function for checking prefixes.""" + return f'{param1}_{param2}' + + function_tool = FunctionTool(test_function) + toolset = _TestingToolset( + tools=[function_tool], + tool_name_prefix='prefix', + ) + + prefixed_tools = await toolset.get_tools_with_prefix() + prefixed_tool = prefixed_tools[0] + + # Tool name should be prefixed + assert prefixed_tool.name == 'prefix_test_function' + + # Function declaration should also have prefixed name + declaration = prefixed_tool._get_declaration() + assert declaration is not None + assert declaration.name == 'prefix_test_function' + + # Description should remain unchanged + assert 'A test function for checking prefixes.' in declaration.description + + +@pytest.mark.asyncio +async def test_prefixed_tools_in_llm_request(): + """Test that prefixed tools are properly added to LLM request.""" + + def test_function(param: str) -> str: + """A test function.""" + return f'result: {param}' + + function_tool = FunctionTool(test_function) + toolset = _TestingToolset(tools=[function_tool], tool_name_prefix='test') + + prefixed_tools = await toolset.get_tools_with_prefix() + prefixed_tool = prefixed_tools[0] + + # Create LLM request and tool context + session_service = InMemorySessionService() + session = await session_service.create_session( + app_name='test_app', user_id='test_user' + ) + agent = SequentialAgent(name='test_agent') + invocation_context = InvocationContext( + invocation_id='test_id', + agent=agent, + session=session, + session_service=session_service, + ) + tool_context = ToolContext(invocation_context) + llm_request = LlmRequest() + + # Process the LLM request with the prefixed tool + await prefixed_tool.process_llm_request( + tool_context=tool_context, llm_request=llm_request + ) + + # Verify the tool is registered with prefixed name in tools_dict + assert 'test_test_function' in llm_request.tools_dict + assert llm_request.tools_dict['test_test_function'] == prefixed_tool + + # Verify the function declaration has prefixed name + assert llm_request.config is not None + assert llm_request.config.tools is not None + assert len(llm_request.config.tools) == 1 + tool_config = llm_request.config.tools[0] + assert len(tool_config.function_declarations) == 1 + func_decl = tool_config.function_declarations[0] + assert func_decl.name == 'test_test_function' + + +@pytest.mark.asyncio +async def test_multiple_tools_have_correct_declarations(): + """Test that each tool maintains its own function declaration after prefixing.""" + + def tool_one(param: str) -> str: + """Function one.""" + return f'one: {param}' + + def tool_two(param: int) -> str: + """Function two.""" + return f'two: {param}' + + tool1 = FunctionTool(tool_one) + tool2 = FunctionTool(tool_two) + toolset = _TestingToolset(tools=[tool1, tool2], tool_name_prefix='test') + + prefixed_tools = await toolset.get_tools_with_prefix() + + # Verify each tool has its own correct declaration + decl1 = prefixed_tools[0]._get_declaration() + decl2 = prefixed_tools[1]._get_declaration() + + assert decl1.name == 'test_tool_one' + assert decl2.name == 'test_tool_two' + + assert 'Function one.' in decl1.description + assert 'Function two.' in decl2.description + + +@pytest.mark.asyncio +async def test_no_duplicate_prefixing(): + """Test that multiple calls to get_tools_with_prefix don't cause duplicate prefixing.""" + original_tool = _TestingTool(name='original', description='Original tool') + toolset = _TestingToolset(tools=[original_tool], tool_name_prefix='test') + + # First call + prefixed_tools_1 = await toolset.get_tools_with_prefix() + assert len(prefixed_tools_1) == 1 + assert prefixed_tools_1[0].name == 'test_original' + + # Second call - should not double-prefix + prefixed_tools_2 = await toolset.get_tools_with_prefix() + assert len(prefixed_tools_2) == 1 + assert prefixed_tools_2[0].name == 'test_original' # Not 'test_test_original' + + # Original tool should remain unchanged + original_tools = await toolset.get_tools() + assert original_tools[0].name == 'original' + + # The prefixed tools should be different instances + assert prefixed_tools_1[0] is not prefixed_tools_2[0] + assert prefixed_tools_1[0] is not original_tools[0]