Skip to content

Commit ca00852

Browse files
anubhav756anubhav756
committed
feat(adk): Implement tool wrapper & auth getters
Implements the ToolboxToolset class in toolbox-adk. - Wraps ToolboxClient to provide tools to ADK agents. - Supports all standard credential strategies. - Adds support for auth_token_getters to handle tools with custom authSources (runtime token resolution). - Includes unit tests for toolset functionality.
1 parent b92c624 commit ca00852

File tree

5 files changed

+644
-5
lines changed

5 files changed

+644
-5
lines changed

‎packages/toolbox-adk/src/toolbox_adk/__init__.py‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -14,16 +14,16 @@
1414

1515
from .client import ToolboxClient
1616
from .credentials import CredentialConfig, CredentialStrategy, CredentialType
17-
# from .tool import ToolboxContext, ToolboxTool
18-
# from .toolset import ToolboxToolset
17+
from .tool import ToolboxContext, ToolboxTool
18+
from .toolset import ToolboxToolset
1919
from .version import __version__
2020

2121
__all__ = [
2222
"CredentialStrategy",
2323
"CredentialConfig",
2424
"CredentialType",
2525
"ToolboxClient",
26-
# "ToolboxTool",
27-
# "ToolboxContext",
28-
# "ToolboxToolset",
26+
"ToolboxTool",
27+
"ToolboxContext",
28+
"ToolboxToolset",
2929
]

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

Lines changed: 179 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,179 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from typing import Any, Awaitable, Callable, Dict, Optional, cast
16+
17+
import toolbox_core
18+
from fastapi.openapi.models import (
19+
OAuth2,
20+
OAuthFlowAuthorizationCode,
21+
OAuthFlows,
22+
SecurityScheme,
23+
)
24+
from google.adk.agents.readonly_context import ReadonlyContext
25+
from google.adk.auth.auth_credential import (
26+
AuthCredential,
27+
AuthCredentialTypes,
28+
OAuth2Auth,
29+
)
30+
from google.adk.auth.auth_tool import AuthConfig
31+
from google.adk.tools.base_tool import BaseTool
32+
from toolbox_core.tool import ToolboxTool as CoreToolboxTool
33+
from typing_extensions import override
34+
35+
from .client import USER_TOKEN_CONTEXT_VAR
36+
from .credentials import CredentialConfig, CredentialType
37+
38+
39+
class ToolboxContext:
40+
"""Context object passed to pre/post hooks."""
41+
42+
def __init__(self, arguments: Dict[str, Any], tool_context: ReadonlyContext):
43+
self.arguments = arguments
44+
self.tool_context = tool_context
45+
self.result: Optional[Any] = None
46+
self.error: Optional[Exception] = None
47+
48+
49+
class ToolboxTool(BaseTool):
50+
"""
51+
A tool that delegates to a remote Toolbox tool, integrated with ADK.
52+
"""
53+
54+
def __init__(
55+
self,
56+
core_tool: CoreToolboxTool,
57+
pre_hook: Optional[Callable[[ToolboxContext], Awaitable[None]]] = None,
58+
post_hook: Optional[Callable[[ToolboxContext], Awaitable[None]]] = None,
59+
auth_config: Optional[CredentialConfig] = None,
60+
):
61+
"""
62+
Args:
63+
core_tool: The underlying toolbox_core.py tool instance.
64+
pre_hook: Async function called before execution. Can modify ctx.arguments.
65+
post_hook: Async function called after execution (finally block). Can inspect ctx.result/error.
66+
auth_config: Credential configuration to handle interactive flows.
67+
"""
68+
# We act as a proxy.
69+
# We need to extract metadata from the core tool to satisfy BaseTool's contract.
70+
71+
name = getattr(core_tool, "__name__", "unknown_tool")
72+
description = (
73+
getattr(core_tool, "__doc__", "No description provided.")
74+
or "No description provided."
75+
)
76+
77+
super().__init__(
78+
name=name,
79+
description=description,
80+
# Pass empty custom_metadata as it is not currently used
81+
custom_metadata={},
82+
)
83+
self._core_tool = core_tool
84+
self._pre_hook = pre_hook
85+
self._post_hook = post_hook
86+
self._auth_config = auth_config
87+
88+
@override
89+
async def run_async(
90+
self,
91+
args: Dict[str, Any],
92+
tool_context: ReadonlyContext,
93+
) -> Any:
94+
# Create context
95+
ctx = ToolboxContext(arguments=args, tool_context=tool_context)
96+
97+
# 1. Pre-hook
98+
if self._pre_hook:
99+
await self._pre_hook(ctx)
100+
101+
# 2. ADK Auth Integration (3LO)
102+
# Check if USER_IDENTITY is configured
103+
reset_token = None
104+
105+
if self._auth_config and self._auth_config.type == CredentialType.USER_IDENTITY:
106+
if not self._auth_config.client_id or not self._auth_config.client_secret:
107+
raise ValueError("USER_IDENTITY requires client_id and client_secret")
108+
109+
# Construct ADK AuthConfig
110+
scopes = self._auth_config.scopes or [
111+
"https://www.googleapis.com/auth/cloud-platform"
112+
]
113+
scope_dict = {s: "" for s in scopes}
114+
115+
auth_config_adk = AuthConfig(
116+
auth_scheme=OAuth2(
117+
flows=OAuthFlows(
118+
authorizationCode=OAuthFlowAuthorizationCode(
119+
authorizationUrl="https://accounts.google.com/o/oauth2/auth",
120+
tokenUrl="https://oauth2.googleapis.com/token",
121+
scopes=scope_dict,
122+
)
123+
)
124+
),
125+
raw_auth_credential=AuthCredential(
126+
auth_type=AuthCredentialTypes.OAUTH2,
127+
oauth2=OAuth2Auth(
128+
client_id=self._auth_config.client_id,
129+
client_secret=self._auth_config.client_secret,
130+
),
131+
),
132+
)
133+
134+
# Check if we already have credentials from a previous exchange
135+
try:
136+
# get_auth_response returns AuthCredential if found
137+
ctx_any = cast(Any, tool_context)
138+
creds = ctx_any.get_auth_response(auth_config_adk)
139+
if creds and creds.oauth2 and creds.oauth2.access_token:
140+
reset_token = USER_TOKEN_CONTEXT_VAR.set(creds.oauth2.access_token)
141+
else:
142+
# Request credentials and pause execution
143+
ctx_any.request_credential(auth_config_adk)
144+
return None
145+
except Exception as e:
146+
ctx.error = e
147+
if "credential" in str(e).lower() or isinstance(e, ValueError):
148+
raise e
149+
# Fallback to request logic
150+
ctx_any = cast(Any, tool_context)
151+
ctx_any.request_credential(auth_config_adk)
152+
return None
153+
154+
try:
155+
# Execute the core tool
156+
result = await self._core_tool(**ctx.arguments)
157+
158+
ctx.result = result
159+
return result
160+
161+
except Exception as e:
162+
ctx.error = e
163+
raise e
164+
finally:
165+
if reset_token:
166+
USER_TOKEN_CONTEXT_VAR.reset(reset_token)
167+
if self._post_hook:
168+
await self._post_hook(ctx)
169+
170+
def bind_params(self, bounded_params: Dict[str, Any]) -> "ToolboxTool":
171+
"""Allows runtime binding of parameters, delegating to core tool."""
172+
new_core_tool = self._core_tool.bind_params(bounded_params)
173+
# Return a new wrapper
174+
return ToolboxTool(
175+
core_tool=new_core_tool,
176+
pre_hook=self._pre_hook,
177+
post_hook=self._post_hook,
178+
auth_config=self._auth_config,
179+
)

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

Lines changed: 123 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,123 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
from typing import Any, Awaitable, Callable, Dict, List, Mapping, Optional, Union
16+
17+
from google.adk.agents.readonly_context import ReadonlyContext
18+
from google.adk.tools.base_tool import BaseTool
19+
from google.adk.tools.base_toolset import BaseToolset
20+
from typing_extensions import override
21+
22+
from .client import ToolboxClient
23+
from .credentials import CredentialConfig
24+
from .tool import ToolboxContext, ToolboxTool
25+
26+
27+
class ToolboxToolset(BaseToolset):
28+
"""
29+
A Toolset that provides tools from a remote Toolbox server.
30+
"""
31+
32+
def __init__(
33+
self,
34+
server_url: str,
35+
toolset_name: Optional[str] = None,
36+
tool_names: Optional[List[str]] = None,
37+
credentials: Optional[CredentialConfig] = None,
38+
additional_headers: Optional[
39+
Dict[str, Union[str, Callable[[], str], Callable[[], Awaitable[str]]]]
40+
] = None,
41+
bound_params: Optional[Mapping[str, Union[Callable[[], Any], Any]]] = None,
42+
auth_token_getters: Optional[
43+
Mapping[str, Union[Callable[[], str], Callable[[], Awaitable[str]]]]
44+
] = None,
45+
pre_hook: Optional[Callable[[ToolboxContext], Awaitable[None]]] = None,
46+
post_hook: Optional[Callable[[ToolboxContext], Awaitable[None]]] = None,
47+
**kwargs: Any,
48+
):
49+
"""
50+
Args:
51+
server_url: The URL of the Toolbox server.
52+
toolset_name: The name of the remote toolset to load.
53+
tool_names: Specific tool names to load (alternative to toolset_name).
54+
credentials: Authentication configuration.
55+
additional_headers: Extra headers (static or dynamic).
56+
bound_params: Parameters to bind globally to loaded tools.
57+
auth_token_getters: Mapping of auth service names to token getters.
58+
pre_hook: Hook to run before every tool execution.
59+
post_hook: Hook to run after every tool execution.
60+
"""
61+
super().__init__()
62+
self._client = ToolboxClient(
63+
server_url=server_url,
64+
credentials=credentials,
65+
additional_headers=additional_headers,
66+
**kwargs,
67+
)
68+
self._toolset_name = toolset_name
69+
self._tool_names = tool_names
70+
self._bound_params = bound_params
71+
self._auth_token_getters = auth_token_getters
72+
self._pre_hook = pre_hook
73+
self._post_hook = post_hook
74+
75+
@override
76+
async def get_tools(
77+
self, readonly_context: Optional[ReadonlyContext] = None
78+
) -> List[BaseTool]:
79+
"""Loads tools from the toolbox server and wraps them."""
80+
# Note: We don't close the client after get_tools because tools might need it.
81+
82+
tools = []
83+
# 1. Load specific toolset if requested
84+
if self._toolset_name:
85+
core_tools = await self._client.load_toolset(
86+
self._toolset_name,
87+
bound_params=self._bound_params or {},
88+
auth_token_getters=self._auth_token_getters or {},
89+
)
90+
tools.extend(core_tools)
91+
92+
# 2. Load specific tools if requested
93+
if self._tool_names:
94+
for name in self._tool_names:
95+
core_tool = await self._client.load_tool(
96+
name,
97+
bound_params=self._bound_params or {},
98+
auth_token_getters=self._auth_token_getters or {},
99+
)
100+
tools.append(core_tool)
101+
102+
# 3. If NO tools/toolsets were specified, default to loading everything (default toolset)
103+
if not self._toolset_name and not self._tool_names:
104+
core_tools = await self._client.load_toolset(
105+
None,
106+
bound_params=self._bound_params or {},
107+
auth_token_getters=self._auth_token_getters or {},
108+
)
109+
tools.extend(core_tools)
110+
111+
# Wrap all core tools in ToolboxTool
112+
return [
113+
ToolboxTool(
114+
core_tool=t,
115+
pre_hook=self._pre_hook,
116+
post_hook=self._post_hook,
117+
auth_config=self._client.credential_config,
118+
)
119+
for t in tools
120+
]
121+
122+
async def close(self):
123+
await self._client.close()

0 commit comments

Comments
 (0)