feat: Implement the ToolboxTool wrapper and core logic

Author: anubhav756

packages/toolbox-adk/src/toolbox_adk/tool.py

@@ -0,0 +1,176 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, Dict, Optional, Callable, Awaitable
+from typing_extensions import override
+
+import toolbox_core
+from toolbox_core.tool import ToolboxTool as CoreToolboxTool
+from google.adk.tools.base_tool import BaseTool
+from google.adk.agents.readonly_context import ReadonlyContext
+
+
+class ToolboxContext:
+    """Context object passed to pre/post hooks."""
+    def __init__(self, arguments: Dict[str, Any], tool_context: ReadonlyContext):
+        self.arguments = arguments
+        self.tool_context = tool_context
+        self.result: Optional[Any] = None
+        self.error: Optional[Exception] = None
+
+
+class ToolboxTool(BaseTool):
+    """
+    A tool that delegates to a remote Toolbox tool, integrated with ADK.
+    """
+
+    def __init__(
+        self,
+        core_tool: CoreToolboxTool,
+        pre_hook: Optional[Callable[[ToolboxContext], Awaitable[None]]] = None,
+        post_hook: Optional[Callable[[ToolboxContext], Awaitable[None]]] = None,
+    ):
+        """
+        Args:
+            core_tool: The underlying toolbox_core.py tool instance.
+            pre_hook: Async function called before execution. Can modify ctx.arguments.
+            post_hook: Async function called after execution (finally block). Can inspect ctx.result/error.
+        """
+        # We act as a proxy.
+        # We need to extract metadata from the core tool to satisfy BaseTool's contract.
+        # core_tool name, description etc. are available.
+        
+        # Note: BaseTool in adk-python typically requires tool_name, tool_description etc.
+        # We map them from the core tool.
+        # core_tool.tool_def is usually available or we use the wrapper's attributes if exposed.
+        # In toolbox-core 0.1.0+, tool instance is callable but metadata is also stored.
+        # Assuming core_tool has typical attributes or we introspect.
+        # For now, we rely on core_tool's string representation or assume it has name/doc.
+        
+        name = getattr(core_tool, "__name__", "unknown_tool")
+        description = getattr(core_tool, "__doc__", "No description provided.") or "No description provided."
+        
+        # toolbox_core tools might not expose explicit args schema in a way BaseTool expects 
+        # (dict of name -> type). ADK tools usually define `tool_args` as a dict.
+        # We might need to inspect the signature if not explicitly available.
+        # For this implementation, we allow dynamic arguments similar to FunctionTool.
+        # If BaseTool enforces providing `tool_args`, we might need to introspect.
+        # However, BaseTool logic in adk-python often derives it. 
+        # We'll call super with minimal known info or defaults.
+        
+        super().__init__(
+            name=name,
+            description=description,
+            # We pass empty custom_metadata or whatever is needed
+            custom_metadata={} 
+        )
+        self._core_tool = core_tool
+        self._pre_hook = pre_hook
+        self._post_hook = post_hook
+
+    @override
+    async def run_async(
+        self,
+        args: Dict[str, Any],
+        tool_context: ReadonlyContext,
+    ) -> Any:
+        # Create context
+        ctx = ToolboxContext(arguments=args, tool_context=tool_context)
+        
+        # 1. Pre-hook
+        if self._pre_hook:
+            await self._pre_hook(ctx)
+            
+        # 2. ADK Auth Integration (3LO)
+        # We check if we need to inject a user token managed by ADK.
+        # This typically happens if the toolset was configured with USER_IDENTITY.
+        # We don't know the configuration here explicitly unless passed, 
+        # BUT the convention is: if we can get a token from tool_context, use it.
+        # Or better: The client (ToolboxClient) didn't set auth headers for USER_IDENTITY,
+        # so we must do it here. 
+        
+        # How do we know if we logic requires it? 
+        # We can try to fetch; if it exists, we assume we should use it?
+        # Or we rely on the fact that `core_tool` instance itself might have bound auth getters?
+        # Actually, toolbox-core tools support `add_auth_token_getters`.
+        # We need to bridge ADK's `get_auth_response` to toolbox-core's getter.
+        
+        # ADK 3LO flow:
+        # response = tool_context.get_auth_response()
+        # if not response:
+        #    tool_context.request_credential(...)
+        #    return "Please authenticate..."
+        # else:
+        #    token = response.token
+        
+        # If we have a way to detect 3LO requirement, we trigger it.
+        # For now, we'll try to retrieve an auth response if one is potentially available.
+        # But `toolbox-core` tool will fail if it needs auth and doesn't get it.
+        # The proper pattern is: wrap the core tool with a dynamically injected getter 
+        # that calls into ADK logic.
+        
+        # However, `toolbox-core`'s `auth_token_getters` expects a callable.
+        # Be careful: `tool_context` is available here in `run_async`, but `auth_token_getters`
+        # are usually bound locally.
+        
+        # We'll define a token getter that closes over `tool_context`.
+        # We'll define a token getter that closes over `tool_context` if needed in future.
+        # For now, unused.
+
+        # Let's try to get a token if available, but fundamentally, ADK expects the tool to return 
+        # a request for credentials if missing.
+        auth_response = tool_context.get_auth_response()
+        
+        # NOTE: We can't easily tell if this specific tool *needs* user auth without metadata.
+        # Strategy: We assume that if the tool fails with an auth error, we might fallback?
+        # OR we rely on configuration. 
+        # Ideally, `ToolboxToolset` would have configured this tool with a special "trigger" if needed.
+        
+        # For now, we will proceed to execute. 
+        # If `USER_IDENTITY` was configured, we might have injected a special marker or getter.
+        # Revisit this logic if explicit "Force Auth" flag is needed on the tool.
+        
+        try:
+            # Execute the core tool
+            # mapping: toolbox-core expects (arg1=..., arg2=...)
+            # We unpack ctx.arguments
+            
+            # For 3LO support: We really need to know if we should be requesting credentials.
+            # For this MVP, we will rely on static auth or pre-existing headers.
+            # If we want to support 3LO properly, we need to handle the `tool_context.request_credential` flow.
+            # We'll check if `auth_token` argument is present or if we have a bound getter.
+            
+            # We execute:
+            result = await self._core_tool(**ctx.arguments)
+            
+            ctx.result = result
+            return result
+            
+        except Exception as e:
+            # TODO: Inspect e for Auth errors to trigger 3LO if configured?
+            ctx.error = e
+            raise e
+        finally:
+            if self._post_hook:
+                await self._post_hook(ctx)
+
+    def bind_params(self, bounded_params: Dict[str, Any]) -> 'ToolboxTool':
+        """Allows runtime binding of parameters, delegating to core tool."""
+        new_core_tool = self._core_tool.bind_params(bounded_params)
+        # Return a new wrapper
+        return ToolboxTool(
+            core_tool=new_core_tool, 
+            pre_hook=self._pre_hook, 
+            post_hook=self._post_hook
+        )

packages/toolbox-adk/src/toolbox_adk/toolset.py

@@ -0,0 +1,107 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, Dict, List, Optional, Mapping, Union, Callable, Awaitable
+from typing_extensions import override
+
+from google.adk.tools.base_tool import BaseTool
+from google.adk.tools.base_toolset import BaseToolset
+from google.adk.agents.readonly_context import ReadonlyContext
+
+from .client import ToolboxClient
+from .credentials import CredentialConfig
+from .tool import ToolboxTool, ToolboxContext
+
+class ToolboxToolset(BaseToolset):
+    """
+    A Toolset that provides tools from a remote Toolbox server.
+    """
+
+    def __init__(
+        self,
+        server_url: str,
+        toolset_name: Optional[str] = None,
+        tool_names: Optional[List[str]] = None,
+        credentials: Optional[CredentialConfig] = None,
+        additional_headers: Optional[Dict[str, str]] = None,
+        bound_params: Optional[Mapping[str, Union[Callable[[], Any], Any]]] = None,
+        pre_hook: Optional[Callable[[ToolboxContext], Awaitable[None]]] = None,
+        post_hook: Optional[Callable[[ToolboxContext], Awaitable[None]]] = None,
+        **kwargs
+    ):
+        """
+        Args:
+            server_url: The URL of the Toolbox server.
+            toolset_name: The name of the remote toolset to load.
+            tool_names: Specific tool names to load (alternative to toolset_name).
+            credentials: Authentication configuration.
+            additional_headers: Extra static headers (e.g. for proxy auth).
+            bound_params: Parameters to bind globally to loaded tools.
+            pre_hook: Hook to run before every tool execution.
+            post_hook: Hook to run after every tool execution.
+        """
+        super().__init__()
+        self._client = ToolboxClient(
+            server_url=server_url,
+            credentials=credentials,
+            additional_headers=additional_headers,
+            **kwargs
+        )
+        self._toolset_name = toolset_name
+        self._tool_names = tool_names
+        self._bound_params = bound_params
+        self._pre_hook = pre_hook
+        self._post_hook = post_hook
+
+    @override
+    async def get_tools(
+        self, readonly_context: Optional[ReadonlyContext] = None
+    ) -> List[BaseTool]:
+        """Loads tools from the toolbox server and wraps them."""
+        # Note: We don't close the client after get_tools because tools might need it 
+        # (though currently tools are self-contained http-wise if they don't share session state).
+        # Actually toolbox-core tools have their own client reference or use a shared one?
+        # toolbox_core.ToolboxClient.load_tool returns a tool which... 
+        # wait, toolbox_core tools make HTTP calls. Do they hold a reference to the client session?
+        # Yes, toolbox_core 0.1.0+ tools usually have an internal `client` or `session`.
+        # So we must keep self._client alive or ensuring it's not closed prematurely.
+        
+        tools = []
+        if self._toolset_name:
+            core_tools = await self._client.load_toolset(
+                self._toolset_name,
+                bound_params=self._bound_params
+            )
+            tools.extend(core_tools)
+
+        if self._tool_names:
+            for name in self._tool_names:
+                core_tool = await self._client.load_tool(
+                    name,
+                    bound_params=self._bound_params
+                )
+                tools.append(core_tool)
+        
+        # Wrap all core tools in ToolboxTool
+        return [
+            ToolboxTool(
+                core_tool=t, 
+                pre_hook=self._pre_hook, 
+                post_hook=self._post_hook
+            ) 
+            for t in tools
+        ]
+
+    async def close(self):
+        await self._client.close()