Skip to content

Commit ebd726f

Browse files
seanzhougooglecopybara-github
authored andcommitted
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
1 parent 54680ed commit ebd726f

File tree

3 files changed

+349
-4
lines changed

3 files changed

+349
-4
lines changed

src/google/adk/agents/llm_agent.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -113,10 +113,11 @@ async def _convert_tool_union_to_tools(
113113
) -> list[BaseTool]:
114114
if isinstance(tool_union, BaseTool):
115115
return [tool_union]
116-
if isinstance(tool_union, Callable):
116+
if callable(tool_union):
117117
return [FunctionTool(func=tool_union)]
118118

119-
return await tool_union.get_tools(ctx)
119+
# At this point, tool_union must be a BaseToolset
120+
return await tool_union.get_tools_with_prefix(ctx)
120121

121122

122123
class LlmAgent(BaseAgent):

src/google/adk/tools/base_toolset.py

Lines changed: 66 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,8 @@
1616

1717
from abc import ABC
1818
from abc import abstractmethod
19+
import copy
20+
from typing import final
1921
from typing import List
2022
from typing import Optional
2123
from typing import Protocol
@@ -58,9 +60,19 @@ class BaseToolset(ABC):
5860
"""
5961

6062
def __init__(
61-
self, *, tool_filter: Optional[Union[ToolPredicate, List[str]]] = None
63+
self,
64+
*,
65+
tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
66+
tool_name_prefix: Optional[str] = None,
6267
):
68+
"""Initialize the toolset.
69+
70+
Args:
71+
tool_filter: Filter to apply to tools.
72+
tool_name_prefix: The prefix to prepend to the names of the tools returned by the toolset.
73+
"""
6374
self.tool_filter = tool_filter
75+
self.tool_name_prefix = tool_name_prefix
6476

6577
@abstractmethod
6678
async def get_tools(
@@ -77,6 +89,59 @@ async def get_tools(
7789
list[BaseTool]: A list of tools available under the specified context.
7890
"""
7991

92+
@final
93+
async def get_tools_with_prefix(
94+
self,
95+
readonly_context: Optional[ReadonlyContext] = None,
96+
) -> list[BaseTool]:
97+
"""Return all tools with optional prefix applied to tool names.
98+
99+
This method calls get_tools() and applies prefixing if tool_name_prefix is provided.
100+
101+
Args:
102+
readonly_context (ReadonlyContext, optional): Context used to filter tools
103+
available to the agent. If None, all tools in the toolset are returned.
104+
105+
Returns:
106+
list[BaseTool]: A list of tools with prefixed names if tool_name_prefix is provided.
107+
"""
108+
tools = await self.get_tools(readonly_context)
109+
110+
if not self.tool_name_prefix:
111+
return tools
112+
113+
prefix = self.tool_name_prefix
114+
115+
# Create copies of tools to avoid modifying original instances
116+
prefixed_tools = []
117+
for tool in tools:
118+
# Create a shallow copy of the tool
119+
tool_copy = copy.copy(tool)
120+
121+
# Apply prefix to the copied tool
122+
prefixed_name = f"{prefix}_{tool.name}"
123+
tool_copy.name = prefixed_name
124+
125+
# Also update the function declaration name if the tool has one
126+
# Use default parameters to capture the current values in the closure
127+
def _create_prefixed_declaration(
128+
original_get_declaration=tool._get_declaration,
129+
prefixed_name=prefixed_name,
130+
):
131+
def _get_prefixed_declaration():
132+
declaration = original_get_declaration()
133+
if declaration is not None:
134+
declaration.name = prefixed_name
135+
return declaration
136+
return None
137+
138+
return _get_prefixed_declaration
139+
140+
tool_copy._get_declaration = _create_prefixed_declaration()
141+
prefixed_tools.append(tool_copy)
142+
143+
return prefixed_tools
144+
80145
async def close(self) -> None:
81146
"""Performs cleanup and releases resources held by the toolset.
82147

0 commit comments

Comments
 (0)