Skip to content

Commit f5a12f6

Browse files
committed
feat: adds the ability to extend the scopes for google api tools that require more than the discovery doc requires
1 parent e7ca943 commit f5a12f6

3 files changed

Lines changed: 44 additions & 5 deletions

File tree

src/google/adk/tools/base_toolset.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -192,7 +192,7 @@ def from_config(
192192
raise ValueError(f"from_config() not implemented for toolset: {cls}")
193193

194194
def _is_tool_selected(
195-
self, tool: BaseTool, readonly_context: ReadonlyContext
195+
self, tool: BaseTool, readonly_context: Optional[ReadonlyContext]
196196
) -> bool:
197197
if not self.tool_filter:
198198
return True

src/google/adk/tools/google_api_tool/google_api_toolset.py

Lines changed: 14 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
from ...agents.readonly_context import ReadonlyContext
2525
from ...auth.auth_credential import ServiceAccount
2626
from ...auth.auth_schemes import OpenIdConnectWithConfig
27+
from ...tools.base_tool import BaseTool
2728
from ...tools.base_toolset import BaseToolset
2829
from ...tools.base_toolset import ToolPredicate
2930
from ..openapi_tool import OpenAPIToolset
@@ -48,6 +49,7 @@ class GoogleApiToolset(BaseToolset):
4849
tool_name_prefix: Optional prefix to add to all tool names in this toolset.
4950
additional_headers: Optional dict of HTTP headers to inject into every request
5051
executed by this toolset.
52+
additional_scopes: Optional list of additional scopes to request.
5153
"""
5254

5355
def __init__(
@@ -61,6 +63,7 @@ def __init__(
6163
tool_name_prefix: Optional[str] = None,
6264
*,
6365
additional_headers: Optional[Dict[str, str]] = None,
66+
additional_scopes: Optional[List[str]] = None,
6467
):
6568
super().__init__(tool_filter=tool_filter, tool_name_prefix=tool_name_prefix)
6669
self.api_name = api_name
@@ -69,12 +72,13 @@ def __init__(
6972
self._client_secret = client_secret
7073
self._service_account = service_account
7174
self._additional_headers = additional_headers
75+
self._additional_scopes = additional_scopes
7276
self._openapi_toolset = self._load_toolset_with_oidc_auth()
7377

7478
@override
7579
async def get_tools(
7680
self, readonly_context: Optional[ReadonlyContext] = None
77-
) -> List[GoogleApiTool]:
81+
) -> list[BaseTool]:
7882
"""Get all tools in the toolset."""
7983
return [
8084
GoogleApiTool(
@@ -95,11 +99,17 @@ def _load_toolset_with_oidc_auth(self) -> OpenAPIToolset:
9599
spec_dict = GoogleApiToOpenApiConverter(
96100
self.api_name, self.api_version
97101
).convert()
98-
scope = list(
102+
discovery_scopes = list(
99103
spec_dict['components']['securitySchemes']['oauth2']['flows'][
100104
'authorizationCode'
101105
]['scopes'].keys()
102-
)[0]
106+
)
107+
default_scope = discovery_scopes[0] if discovery_scopes else None
108+
109+
scopes = ([default_scope] if default_scope else []) + (
110+
self._additional_scopes or []
111+
)
112+
103113
return OpenAPIToolset(
104114
spec_dict=spec_dict,
105115
spec_str_type='yaml',
@@ -117,7 +127,7 @@ def _load_toolset_with_oidc_auth(self) -> OpenAPIToolset:
117127
'client_secret_basic',
118128
],
119129
grant_types_supported=['authorization_code'],
120-
scopes=[scope],
130+
scopes=scopes,
121131
),
122132
)
123133

tests/unittests/tools/google_api_tool/test_google_api_toolset.py

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -158,6 +158,35 @@ def test_init(
158158
assert isinstance(kwargs["auth_scheme"], OpenIdConnectWithConfig)
159159
assert kwargs["auth_scheme"].scopes == [DEFAULT_SCOPE]
160160

161+
@mock.patch(
162+
"google.adk.tools.google_api_tool.google_api_toolset.OpenAPIToolset"
163+
)
164+
@mock.patch(
165+
"google.adk.tools.google_api_tool.google_api_toolset.GoogleApiToOpenApiConverter"
166+
)
167+
def test_init_with_additional_scopes(
168+
self,
169+
mock_converter_class,
170+
mock_openapi_toolset_class,
171+
mock_converter_instance,
172+
mock_openapi_toolset_instance,
173+
):
174+
"""Test GoogleApiToolset initialization with additional scopes."""
175+
mock_converter_class.return_value = mock_converter_instance
176+
mock_openapi_toolset_class.return_value = mock_openapi_toolset_instance
177+
178+
extra_scopes = ["https://www.googleapis.com/auth/calendar.readonly"]
179+
tool_set = GoogleApiToolset(
180+
api_name=TEST_API_NAME,
181+
api_version=TEST_API_VERSION,
182+
additional_scopes=extra_scopes,
183+
)
184+
185+
mock_openapi_toolset_class.assert_called_once()
186+
_, kwargs = mock_openapi_toolset_class.call_args
187+
assert isinstance(kwargs["auth_scheme"], OpenIdConnectWithConfig)
188+
assert kwargs["auth_scheme"].scopes == [DEFAULT_SCOPE] + extra_scopes
189+
161190
@mock.patch(
162191
"google.adk.tools.google_api_tool.google_api_toolset.GoogleApiTool"
163192
)

0 commit comments

Comments
 (0)