feat: Modify the ADK ToolboxToolset

Author: anubhav756

src/google/adk/tools/toolbox_toolset.py

@@ -19,22 +19,27 @@
 from typing import Optional
 from typing import Union
 
-import toolbox_core as toolbox
 from typing_extensions import override
 
 from ..agents.readonly_context import ReadonlyContext
 from .base_tool import BaseTool
 from .base_toolset import BaseToolset
-from .function_tool import FunctionTool
 
 
 class ToolboxToolset(BaseToolset):
   """A class that provides access to toolbox toolsets.
 
+  This class acts as a bridge to the `toolbox-adk` package.
+  You must install `toolbox-adk` to use this class.
+
   Example:
   ```python
-  toolbox_toolset = ToolboxToolset("http://127.0.0.1:5000",
-  toolset_name="my-toolset")
+  from toolbox_adk import CredentialStrategy
+
+  toolbox_toolset = ToolboxToolset(
+      server_url="http://127.0.0.1:5000",
+      # toolset_name and tool_names are optional. If omitted, all tools are loaded.
+      credentials=CredentialStrategy.toolbox_identity()
   )
   ```
   """
@@ -48,60 +53,51 @@ def __init__(
       bound_params: Optional[
           Mapping[str, Union[Callable[[], Any], Any]]
       ] = None,
+      # New args mapping to toolbox-adk
+      credentials: Optional[Any] = None,
+      additional_headers: Optional[dict[str, str]] = None,
+      **kwargs,
   ):
     """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.
+      auth_token_getters: (Deprecated) Map of auth token getters.
+      bound_params: Parameters to bind to the tools.
+      credentials: (Optional) toolbox_adk.CredentialConfig object.
+      additional_headers: (Optional) Static headers dictionary.
+      **kwargs: Additional arguments passed to the underlying toolbox_adk.ToolboxToolset.
     """
-    if not tool_names and not toolset_name:
-      raise ValueError("tool_names and toolset_name cannot both be None")
+    try:
+      from toolbox_adk import ToolboxToolset as RealToolboxToolset  # pylint: disable=import-outside-toplevel
+    except ImportError as exc:
+      raise ImportError(
+          "ToolboxToolset requires the 'toolbox-adk' package. "
+          "Please install it using `pip install toolbox-adk`."
+      ) from exc
+      
     super().__init__()
-    self._server_url = 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 or {}
-    self._bound_params = bound_params or {}
+    
+    # Delegate to the real implementation in toolbox-adk.
+    
+    self._delegate = RealToolboxToolset(
+        server_url=server_url,
+        toolset_name=toolset_name,
+        tool_names=tool_names,
+        credentials=credentials,
+        additional_headers=additional_headers,
+        bound_params=bound_params,
+        # Backward compatibility: Pass legacy args via kwargs.
+        auth_token_getters=auth_token_getters, 
+        **kwargs
+    )
 
   @override
   async def get_tools(
       self, readonly_context: Optional[ReadonlyContext] = None
   ) -> list[BaseTool]:
-    tools = []
-    if self._toolset_name:
-      tools.extend([
-          FunctionTool(tool)
-          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(
-              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
+    return await self._delegate.get_tools(readonly_context)
 
   @override
   async def close(self):
-    self._toolbox_client.close()
+    await self._delegate.close()