NeoAnd/adk-python
1
0
Fork 0
mirror of https://github.com/EvolutionAPI/adk-python.git synced 2025-07-16 04:02:55 -06:00
Code Issues Actions Packages Projects Releases Wiki Activity

enhance toolbox toolset:

Browse Source 
1. use async client
2. expose add headers and bind parameters functionality of toolbox tool

PiperOrigin-RevId: 760845486
This commit is contained in:
Xiang (Sean) Zhou 2025-05-19 18:07:29 -07:00 committed by Copybara-Service
parent 9eef4464c6
commit de7c9c6509
1 changed files with 29 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

32
src/google/adk/tools/toolbox_toolset.py
View File

@ -12,8 +12,12 @@
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
from typing import Any
from typing import Callable
from typing import List from typing import List
from typing import Mapping
from typing import Optional from typing import Optional
from typing import Union
import toolbox_core as toolbox import toolbox_core as toolbox
from typing_extensions import override from typing_extensions import override
@ -40,12 +44,22 @@ class ToolboxToolset(BaseToolset):
server_url: str, server_url: str,
toolset_name: Optional[str] = None, toolset_name: Optional[str] = None,
tool_names: Optional[List[str]] = None, tool_names: Optional[List[str]] = None,
auth_token_getters: dict[str, Callable[[], str]] = {},
bound_params: Mapping[str, Union[Callable[[], Any], Any]] = {},
): ):
"""Args: """Args:
server_url: The URL of the toolbox server. server_url: The URL of the toolbox server.
toolset_name: The name of the toolbox toolset to load. toolset_name: The name of the toolbox toolset to load.
tool_names: The names of the tools to load. tool_names: The names of the tools to load.
auth_token_getters: A mapping of authentication service names to
callables that return the corresponding authentication token. see:
https://github.com/googleapis/mcp-toolbox-sdk-python/tree/main/packages/toolbox-core#authenticating-tools
for details.
bound_params: A mapping of parameter names to bind to specific values or
callables that are called to produce values as needed. see:
https://github.com/googleapis/mcp-toolbox-sdk-python/tree/main/packages/toolbox-core#binding-parameter-values
for details.
The resulting ToolboxToolset will contain both tools loaded by tool_names The resulting ToolboxToolset will contain both tools loaded by tool_names
and toolset_name. and toolset_name.
""" """
@ -53,9 +67,11 @@ class ToolboxToolset(BaseToolset):
raise ValueError("tool_names and toolset_name cannot both be None") raise ValueError("tool_names and toolset_name cannot both be None")
super().__init__() super().__init__()
self._server_url = server_url self._server_url = server_url
self._toolbox_client = toolbox.ToolboxSyncClient(server_url) self._toolbox_client = toolbox.ToolboxClient(server_url)
self._toolset_name = toolset_name self._toolset_name = toolset_name
self._tool_names = tool_names self._tool_names = tool_names
self._auth_token_getters = auth_token_getters
self._bound_params = bound_params
@override @override
async def get_tools( async def get_tools(
@ -65,11 +81,21 @@ class ToolboxToolset(BaseToolset):
if self._toolset_name: if self._toolset_name:
tools.extend([ tools.extend([
FunctionTool(tool) FunctionTool(tool)
for tool in self._toolbox_client.load_toolset(self._toolset_name) for tool in await self._toolbox_client.load_toolset(
self._toolset_name,
auth_token_getters=self._auth_token_getters,
bound_params=self._bound_params,
)
]) ])
if self._tool_names: if self._tool_names:
tools.extend([ tools.extend([
FunctionTool(self._toolbox_client.load_tool(tool_name)) FunctionTool(
await self._toolbox_client.load_tool(
tool_name,
auth_token_getters=self._auth_token_getters,
bound_params=self._bound_params,
)
)
for tool_name in self._tool_names for tool_name in self._tool_names
]) ])
return tools return tools