mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
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
This commit is contained in:
committed by
Copybara-Service
parent
54680edf3c
commit
ebd726f1f5
@@ -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):
|
||||
|
||||
@@ -16,6 +16,8 @@ from __future__ import annotations
|
||||
|
||||
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 @@ class BaseToolset(ABC):
|
||||
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.
|
||||
|
||||
|
||||
@@ -23,17 +23,29 @@ from google.adk.models.llm_request import LlmRequest
|
||||
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 test_process_llm_request_can_be_overridden():
|
||||
|
||||
# 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]
|
||||
|
||||
Reference in New Issue
Block a user