diff --git a/src/google/adk/tools/toolbox_toolset.py b/src/google/adk/tools/toolbox_toolset.py index 67d3bdb..f4ce579 100644 --- a/src/google/adk/tools/toolbox_toolset.py +++ b/src/google/adk/tools/toolbox_toolset.py @@ -12,8 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +from typing import Any +from typing import Callable from typing import List +from typing import Mapping from typing import Optional +from typing import Union import toolbox_core as toolbox from typing_extensions import override @@ -40,12 +44,22 @@ class ToolboxToolset(BaseToolset): server_url: str, toolset_name: Optional[str] = None, tool_names: Optional[List[str]] = None, + auth_token_getters: dict[str, Callable[[], str]] = {}, + bound_params: Mapping[str, Union[Callable[[], Any], Any]] = {}, ): """Args: server_url: The URL of the toolbox server. toolset_name: The name of the toolbox toolset 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 and toolset_name. """ @@ -53,9 +67,11 @@ class ToolboxToolset(BaseToolset): raise ValueError("tool_names and toolset_name cannot both be None") super().__init__() 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._tool_names = tool_names + self._auth_token_getters = auth_token_getters + self._bound_params = bound_params @override async def get_tools( @@ -65,11 +81,21 @@ class ToolboxToolset(BaseToolset): if self._toolset_name: tools.extend([ 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: 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 ]) return tools