feat: Support custom tool_name_prefix in auto-generated GoogleApiToolset

PiperOrigin-RevId: 795508179
This commit is contained in:
Xiang (Sean) Zhou
2025-08-15 09:44:16 -07:00
committed by Copybara-Service
parent 1328e6ef62
commit a2832d5ac7
4 changed files with 161 additions and 17 deletions
@@ -46,7 +46,8 @@ calendar_toolset = CalendarToolset(
# google calendar tool by adding `calendar_events_list` in the filter list
client_id=oauth_client_id,
client_secret=oauth_client_secret,
tool_filter=["calendar_events_get"],
tool_filter=["calendar_events_get", "calendar_events_update"],
tool_name_prefix="google",
)
@@ -125,7 +126,7 @@ root_agent = Agent(
Scenario2:
User want to know the details of one of the listed calendar events.
Use get_calendar_event to get the details of a calendar event.
Use google_calendar_events_get to get the details of a calendar event.
Current user:
@@ -36,6 +36,15 @@ class GoogleApiToolset(BaseToolset):
Usually one toolsets will contains tools only related to one Google API, e.g.
Google Bigquery API toolset will contains tools only related to Google
Bigquery API, like list dataset tool, list table tool etc.
Args:
api_name: The name of the Google API (e.g., "calendar", "gmail").
api_version: The version of the API (e.g., "v3", "v1").
client_id: OAuth2 client ID for authentication.
client_secret: OAuth2 client secret for authentication.
tool_filter: Optional filter to include only specific tools or use a predicate function.
service_account: Optional service account for authentication.
tool_name_prefix: Optional prefix to add to all tool names in this toolset.
"""
def __init__(
@@ -46,8 +55,9 @@ class GoogleApiToolset(BaseToolset):
client_secret: Optional[str] = None,
tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
service_account: Optional[ServiceAccount] = None,
tool_name_prefix: Optional[str] = None,
):
super().__init__(tool_filter=tool_filter)
super().__init__(tool_filter=tool_filter, tool_name_prefix=tool_name_prefix)
self.api_name = api_name
self.api_version = api_version
self._client_id = client_id
@@ -27,7 +27,15 @@ logger = logging.getLogger("google_adk." + __name__)
class BigQueryToolset(GoogleApiToolset):
"""Auto-generated Bigquery toolset based on Google BigQuery API v2 spec exposed by Google API discovery API"""
"""Auto-generated BigQuery toolset based on Google BigQuery API v2 spec exposed by Google API discovery API.
Args:
client_id: OAuth2 client ID for authentication.
client_secret: OAuth2 client secret for authentication.
tool_filter: Optional filter to include only specific tools or use a predicate function.
service_account: Optional service account for authentication.
tool_name_prefix: Optional prefix to add to all tool names in this toolset.
"""
def __init__(
self,
@@ -35,14 +43,29 @@ class BigQueryToolset(GoogleApiToolset):
client_secret: Optional[str] = None,
tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
service_account: Optional[ServiceAccount] = None,
tool_name_prefix: Optional[str] = None,
):
super().__init__(
"bigquery", "v2", client_id, client_secret, tool_filter, service_account
"bigquery",
"v2",
client_id,
client_secret,
tool_filter,
service_account,
tool_name_prefix,
)
class CalendarToolset(GoogleApiToolset):
"""Auto-generated Calendar toolset based on Google Calendar API v3 spec exposed by Google API discovery API"""
"""Auto-generated Calendar toolset based on Google Calendar API v3 spec exposed by Google API discovery API.
Args:
client_id: OAuth2 client ID for authentication.
client_secret: OAuth2 client secret for authentication.
tool_filter: Optional filter to include only specific tools or use a predicate function.
service_account: Optional service account for authentication.
tool_name_prefix: Optional prefix to add to all tool names in this toolset.
"""
def __init__(
self,
@@ -50,14 +73,29 @@ class CalendarToolset(GoogleApiToolset):
client_secret: Optional[str] = None,
tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
service_account: Optional[ServiceAccount] = None,
tool_name_prefix: Optional[str] = None,
):
super().__init__(
"calendar", "v3", client_id, client_secret, tool_filter, service_account
"calendar",
"v3",
client_id,
client_secret,
tool_filter,
service_account,
tool_name_prefix,
)
class GmailToolset(GoogleApiToolset):
"""Auto-generated Gmail toolset based on Google Gmail API v1 spec exposed by Google API discovery API"""
"""Auto-generated Gmail toolset based on Google Gmail API v1 spec exposed by Google API discovery API.
Args:
client_id: OAuth2 client ID for authentication.
client_secret: OAuth2 client secret for authentication.
tool_filter: Optional filter to include only specific tools or use a predicate function.
service_account: Optional service account for authentication.
tool_name_prefix: Optional prefix to add to all tool names in this toolset.
"""
def __init__(
self,
@@ -65,14 +103,29 @@ class GmailToolset(GoogleApiToolset):
client_secret: Optional[str] = None,
tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
service_account: Optional[ServiceAccount] = None,
tool_name_prefix: Optional[str] = None,
):
super().__init__(
"gmail", "v1", client_id, client_secret, tool_filter, service_account
"gmail",
"v1",
client_id,
client_secret,
tool_filter,
service_account,
tool_name_prefix,
)
class YoutubeToolset(GoogleApiToolset):
"""Auto-generated Youtube toolset based on Youtube API v3 spec exposed by Google API discovery API"""
"""Auto-generated YouTube toolset based on YouTube API v3 spec exposed by Google API discovery API.
Args:
client_id: OAuth2 client ID for authentication.
client_secret: OAuth2 client secret for authentication.
tool_filter: Optional filter to include only specific tools or use a predicate function.
service_account: Optional service account for authentication.
tool_name_prefix: Optional prefix to add to all tool names in this toolset.
"""
def __init__(
self,
@@ -80,14 +133,29 @@ class YoutubeToolset(GoogleApiToolset):
client_secret: Optional[str] = None,
tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
service_account: Optional[ServiceAccount] = None,
tool_name_prefix: Optional[str] = None,
):
super().__init__(
"youtube", "v3", client_id, client_secret, tool_filter, service_account
"youtube",
"v3",
client_id,
client_secret,
tool_filter,
service_account,
tool_name_prefix,
)
class SlidesToolset(GoogleApiToolset):
"""Auto-generated Slides toolset based on Google Slides API v1 spec exposed by Google API discovery API"""
"""Auto-generated Slides toolset based on Google Slides API v1 spec exposed by Google API discovery API.
Args:
client_id: OAuth2 client ID for authentication.
client_secret: OAuth2 client secret for authentication.
tool_filter: Optional filter to include only specific tools or use a predicate function.
service_account: Optional service account for authentication.
tool_name_prefix: Optional prefix to add to all tool names in this toolset.
"""
def __init__(
self,
@@ -95,14 +163,29 @@ class SlidesToolset(GoogleApiToolset):
client_secret: Optional[str] = None,
tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
service_account: Optional[ServiceAccount] = None,
tool_name_prefix: Optional[str] = None,
):
super().__init__(
"slides", "v1", client_id, client_secret, tool_filter, service_account
"slides",
"v1",
client_id,
client_secret,
tool_filter,
service_account,
tool_name_prefix,
)
class SheetsToolset(GoogleApiToolset):
"""Auto-generated Sheets toolset based on Google Sheets API v4 spec exposed by Google API discovery API"""
"""Auto-generated Sheets toolset based on Google Sheets API v4 spec exposed by Google API discovery API.
Args:
client_id: OAuth2 client ID for authentication.
client_secret: OAuth2 client secret for authentication.
tool_filter: Optional filter to include only specific tools or use a predicate function.
service_account: Optional service account for authentication.
tool_name_prefix: Optional prefix to add to all tool names in this toolset.
"""
def __init__(
self,
@@ -110,12 +193,29 @@ class SheetsToolset(GoogleApiToolset):
client_secret: Optional[str] = None,
tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
service_account: Optional[ServiceAccount] = None,
tool_name_prefix: Optional[str] = None,
):
super().__init__("sheets", "v4", client_id, client_secret, tool_filter)
super().__init__(
"sheets",
"v4",
client_id,
client_secret,
tool_filter,
service_account,
tool_name_prefix,
)
class DocsToolset(GoogleApiToolset):
"""Auto-generated Docs toolset based on Google Docs API v1 spec exposed by Google API discovery API"""
"""Auto-generated Docs toolset based on Google Docs API v1 spec exposed by Google API discovery API.
Args:
client_id: OAuth2 client ID for authentication.
client_secret: OAuth2 client secret for authentication.
tool_filter: Optional filter to include only specific tools or use a predicate function.
service_account: Optional service account for authentication.
tool_name_prefix: Optional prefix to add to all tool names in this toolset.
"""
def __init__(
self,
@@ -123,7 +223,14 @@ class DocsToolset(GoogleApiToolset):
client_secret: Optional[str] = None,
tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
service_account: Optional[ServiceAccount] = None,
tool_name_prefix: Optional[str] = None,
):
super().__init__(
"docs", "v1", client_id, client_secret, tool_filter, service_account
"docs",
"v1",
client_id,
client_secret,
tool_filter,
service_account,
tool_name_prefix,
)
@@ -425,3 +425,29 @@ class TestGoogleApiToolset:
tool_set.set_tool_filter(new_filter_predicate)
assert tool_set.tool_filter == new_filter_predicate
@mock.patch(
"google.adk.tools.google_api_tool.google_api_toolset.OpenAPIToolset"
)
@mock.patch(
"google.adk.tools.google_api_tool.google_api_toolset.GoogleApiToOpenApiConverter"
)
def test_init_with_tool_name_prefix(
self,
mock_converter_class,
mock_openapi_toolset_class,
mock_converter_instance,
mock_openapi_toolset_instance,
):
"""Test GoogleApiToolset initialization with tool_name_prefix."""
mock_converter_class.return_value = mock_converter_instance
mock_openapi_toolset_class.return_value = mock_openapi_toolset_instance
tool_name_prefix = "test_prefix"
tool_set = GoogleApiToolset(
api_name=TEST_API_NAME,
api_version=TEST_API_VERSION,
tool_name_prefix=tool_name_prefix,
)
assert tool_set.tool_name_prefix == tool_name_prefix