Skip to content

Commit a2832d5

Browse files
seanzhougooglecopybara-github
authored andcommitted
feat: Support custom tool_name_prefix in auto-generated GoogleApiToolset
PiperOrigin-RevId: 795508179
1 parent 1328e6e commit a2832d5

File tree

4 files changed

+161
-17
lines changed

4 files changed

+161
-17
lines changed

contributing/samples/oauth_calendar_agent/agent.py

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,8 @@
4646
# google calendar tool by adding `calendar_events_list` in the filter list
4747
client_id=oauth_client_id,
4848
client_secret=oauth_client_secret,
49-
tool_filter=["calendar_events_get"],
49+
tool_filter=["calendar_events_get", "calendar_events_update"],
50+
tool_name_prefix="google",
5051
)
5152

5253

@@ -125,7 +126,7 @@ def update_time(callback_context: CallbackContext):
125126
126127
Scenario2:
127128
User want to know the details of one of the listed calendar events.
128-
Use get_calendar_event to get the details of a calendar event.
129+
Use google_calendar_events_get to get the details of a calendar event.
129130
130131
131132
Current user:

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

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,15 @@ class GoogleApiToolset(BaseToolset):
3636
Usually one toolsets will contains tools only related to one Google API, e.g.
3737
Google Bigquery API toolset will contains tools only related to Google
3838
Bigquery API, like list dataset tool, list table tool etc.
39+
40+
Args:
41+
api_name: The name of the Google API (e.g., "calendar", "gmail").
42+
api_version: The version of the API (e.g., "v3", "v1").
43+
client_id: OAuth2 client ID for authentication.
44+
client_secret: OAuth2 client secret for authentication.
45+
tool_filter: Optional filter to include only specific tools or use a predicate function.
46+
service_account: Optional service account for authentication.
47+
tool_name_prefix: Optional prefix to add to all tool names in this toolset.
3948
"""
4049

4150
def __init__(
@@ -46,8 +55,9 @@ def __init__(
4655
client_secret: Optional[str] = None,
4756
tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
4857
service_account: Optional[ServiceAccount] = None,
58+
tool_name_prefix: Optional[str] = None,
4959
):
50-
super().__init__(tool_filter=tool_filter)
60+
super().__init__(tool_filter=tool_filter, tool_name_prefix=tool_name_prefix)
5161
self.api_name = api_name
5262
self.api_version = api_version
5363
self._client_id = client_id

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

Lines changed: 121 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -27,103 +27,210 @@
2727

2828

2929
class BigQueryToolset(GoogleApiToolset):
30-
"""Auto-generated Bigquery toolset based on Google BigQuery API v2 spec exposed by Google API discovery API"""
30+
"""Auto-generated BigQuery toolset based on Google BigQuery API v2 spec exposed by Google API discovery API.
31+
32+
Args:
33+
client_id: OAuth2 client ID for authentication.
34+
client_secret: OAuth2 client secret for authentication.
35+
tool_filter: Optional filter to include only specific tools or use a predicate function.
36+
service_account: Optional service account for authentication.
37+
tool_name_prefix: Optional prefix to add to all tool names in this toolset.
38+
"""
3139

3240
def __init__(
3341
self,
3442
client_id: Optional[str] = None,
3543
client_secret: Optional[str] = None,
3644
tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
3745
service_account: Optional[ServiceAccount] = None,
46+
tool_name_prefix: Optional[str] = None,
3847
):
3948
super().__init__(
40-
"bigquery", "v2", client_id, client_secret, tool_filter, service_account
49+
"bigquery",
50+
"v2",
51+
client_id,
52+
client_secret,
53+
tool_filter,
54+
service_account,
55+
tool_name_prefix,
4156
)
4257

4358

4459
class CalendarToolset(GoogleApiToolset):
45-
"""Auto-generated Calendar toolset based on Google Calendar API v3 spec exposed by Google API discovery API"""
60+
"""Auto-generated Calendar toolset based on Google Calendar API v3 spec exposed by Google API discovery API.
61+
62+
Args:
63+
client_id: OAuth2 client ID for authentication.
64+
client_secret: OAuth2 client secret for authentication.
65+
tool_filter: Optional filter to include only specific tools or use a predicate function.
66+
service_account: Optional service account for authentication.
67+
tool_name_prefix: Optional prefix to add to all tool names in this toolset.
68+
"""
4669

4770
def __init__(
4871
self,
4972
client_id: Optional[str] = None,
5073
client_secret: Optional[str] = None,
5174
tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
5275
service_account: Optional[ServiceAccount] = None,
76+
tool_name_prefix: Optional[str] = None,
5377
):
5478
super().__init__(
55-
"calendar", "v3", client_id, client_secret, tool_filter, service_account
79+
"calendar",
80+
"v3",
81+
client_id,
82+
client_secret,
83+
tool_filter,
84+
service_account,
85+
tool_name_prefix,
5686
)
5787

5888

5989
class GmailToolset(GoogleApiToolset):
60-
"""Auto-generated Gmail toolset based on Google Gmail API v1 spec exposed by Google API discovery API"""
90+
"""Auto-generated Gmail toolset based on Google Gmail API v1 spec exposed by Google API discovery API.
91+
92+
Args:
93+
client_id: OAuth2 client ID for authentication.
94+
client_secret: OAuth2 client secret for authentication.
95+
tool_filter: Optional filter to include only specific tools or use a predicate function.
96+
service_account: Optional service account for authentication.
97+
tool_name_prefix: Optional prefix to add to all tool names in this toolset.
98+
"""
6199

62100
def __init__(
63101
self,
64102
client_id: Optional[str] = None,
65103
client_secret: Optional[str] = None,
66104
tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
67105
service_account: Optional[ServiceAccount] = None,
106+
tool_name_prefix: Optional[str] = None,
68107
):
69108
super().__init__(
70-
"gmail", "v1", client_id, client_secret, tool_filter, service_account
109+
"gmail",
110+
"v1",
111+
client_id,
112+
client_secret,
113+
tool_filter,
114+
service_account,
115+
tool_name_prefix,
71116
)
72117

73118

74119
class YoutubeToolset(GoogleApiToolset):
75-
"""Auto-generated Youtube toolset based on Youtube API v3 spec exposed by Google API discovery API"""
120+
"""Auto-generated YouTube toolset based on YouTube API v3 spec exposed by Google API discovery API.
121+
122+
Args:
123+
client_id: OAuth2 client ID for authentication.
124+
client_secret: OAuth2 client secret for authentication.
125+
tool_filter: Optional filter to include only specific tools or use a predicate function.
126+
service_account: Optional service account for authentication.
127+
tool_name_prefix: Optional prefix to add to all tool names in this toolset.
128+
"""
76129

77130
def __init__(
78131
self,
79132
client_id: Optional[str] = None,
80133
client_secret: Optional[str] = None,
81134
tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
82135
service_account: Optional[ServiceAccount] = None,
136+
tool_name_prefix: Optional[str] = None,
83137
):
84138
super().__init__(
85-
"youtube", "v3", client_id, client_secret, tool_filter, service_account
139+
"youtube",
140+
"v3",
141+
client_id,
142+
client_secret,
143+
tool_filter,
144+
service_account,
145+
tool_name_prefix,
86146
)
87147

88148

89149
class SlidesToolset(GoogleApiToolset):
90-
"""Auto-generated Slides toolset based on Google Slides API v1 spec exposed by Google API discovery API"""
150+
"""Auto-generated Slides toolset based on Google Slides API v1 spec exposed by Google API discovery API.
151+
152+
Args:
153+
client_id: OAuth2 client ID for authentication.
154+
client_secret: OAuth2 client secret for authentication.
155+
tool_filter: Optional filter to include only specific tools or use a predicate function.
156+
service_account: Optional service account for authentication.
157+
tool_name_prefix: Optional prefix to add to all tool names in this toolset.
158+
"""
91159

92160
def __init__(
93161
self,
94162
client_id: Optional[str] = None,
95163
client_secret: Optional[str] = None,
96164
tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
97165
service_account: Optional[ServiceAccount] = None,
166+
tool_name_prefix: Optional[str] = None,
98167
):
99168
super().__init__(
100-
"slides", "v1", client_id, client_secret, tool_filter, service_account
169+
"slides",
170+
"v1",
171+
client_id,
172+
client_secret,
173+
tool_filter,
174+
service_account,
175+
tool_name_prefix,
101176
)
102177

103178

104179
class SheetsToolset(GoogleApiToolset):
105-
"""Auto-generated Sheets toolset based on Google Sheets API v4 spec exposed by Google API discovery API"""
180+
"""Auto-generated Sheets toolset based on Google Sheets API v4 spec exposed by Google API discovery API.
181+
182+
Args:
183+
client_id: OAuth2 client ID for authentication.
184+
client_secret: OAuth2 client secret for authentication.
185+
tool_filter: Optional filter to include only specific tools or use a predicate function.
186+
service_account: Optional service account for authentication.
187+
tool_name_prefix: Optional prefix to add to all tool names in this toolset.
188+
"""
106189

107190
def __init__(
108191
self,
109192
client_id: Optional[str] = None,
110193
client_secret: Optional[str] = None,
111194
tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
112195
service_account: Optional[ServiceAccount] = None,
196+
tool_name_prefix: Optional[str] = None,
113197
):
114-
super().__init__("sheets", "v4", client_id, client_secret, tool_filter)
198+
super().__init__(
199+
"sheets",
200+
"v4",
201+
client_id,
202+
client_secret,
203+
tool_filter,
204+
service_account,
205+
tool_name_prefix,
206+
)
115207

116208

117209
class DocsToolset(GoogleApiToolset):
118-
"""Auto-generated Docs toolset based on Google Docs API v1 spec exposed by Google API discovery API"""
210+
"""Auto-generated Docs toolset based on Google Docs API v1 spec exposed by Google API discovery API.
211+
212+
Args:
213+
client_id: OAuth2 client ID for authentication.
214+
client_secret: OAuth2 client secret for authentication.
215+
tool_filter: Optional filter to include only specific tools or use a predicate function.
216+
service_account: Optional service account for authentication.
217+
tool_name_prefix: Optional prefix to add to all tool names in this toolset.
218+
"""
119219

120220
def __init__(
121221
self,
122222
client_id: Optional[str] = None,
123223
client_secret: Optional[str] = None,
124224
tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
125225
service_account: Optional[ServiceAccount] = None,
226+
tool_name_prefix: Optional[str] = None,
126227
):
127228
super().__init__(
128-
"docs", "v1", client_id, client_secret, tool_filter, service_account
229+
"docs",
230+
"v1",
231+
client_id,
232+
client_secret,
233+
tool_filter,
234+
service_account,
235+
tool_name_prefix,
129236
)

tests/unittests/tools/google_api_tool/test_google_api_toolset.py

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -425,3 +425,29 @@ def new_filter_predicate(
425425

426426
tool_set.set_tool_filter(new_filter_predicate)
427427
assert tool_set.tool_filter == new_filter_predicate
428+
429+
@mock.patch(
430+
"google.adk.tools.google_api_tool.google_api_toolset.OpenAPIToolset"
431+
)
432+
@mock.patch(
433+
"google.adk.tools.google_api_tool.google_api_toolset.GoogleApiToOpenApiConverter"
434+
)
435+
def test_init_with_tool_name_prefix(
436+
self,
437+
mock_converter_class,
438+
mock_openapi_toolset_class,
439+
mock_converter_instance,
440+
mock_openapi_toolset_instance,
441+
):
442+
"""Test GoogleApiToolset initialization with tool_name_prefix."""
443+
mock_converter_class.return_value = mock_converter_instance
444+
mock_openapi_toolset_class.return_value = mock_openapi_toolset_instance
445+
446+
tool_name_prefix = "test_prefix"
447+
tool_set = GoogleApiToolset(
448+
api_name=TEST_API_NAME,
449+
api_version=TEST_API_VERSION,
450+
tool_name_prefix=tool_name_prefix,
451+
)
452+
453+
assert tool_set.tool_name_prefix == tool_name_prefix

0 commit comments

Comments
 (0)