feat: adds the ability to set up google api tools with a custom discovery doc url

Author: tonycoco

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

@@ -24,7 +24,6 @@
 from ...agents.readonly_context import ReadonlyContext
 from ...auth.auth_credential import ServiceAccount
 from ...auth.auth_schemes import OpenIdConnectWithConfig
-from ...tools.base_tool import BaseTool
 from ...tools.base_toolset import BaseToolset
 from ...tools.base_toolset import ToolPredicate
 from ..openapi_tool import OpenAPIToolset
@@ -64,6 +63,7 @@ def __init__(
       *,
       additional_headers: Optional[Dict[str, str]] = None,
       additional_scopes: Optional[List[str]] = None,
+      discovery_url: Optional[str] = None,
   ):
     super().__init__(tool_filter=tool_filter, tool_name_prefix=tool_name_prefix)
     self.api_name = api_name
@@ -73,12 +73,13 @@ def __init__(
     self._service_account = service_account
     self._additional_headers = additional_headers
     self._additional_scopes = additional_scopes
+    self._discovery_url = discovery_url
     self._openapi_toolset = self._load_toolset_with_oidc_auth()
 
   @override
   async def get_tools(
       self, readonly_context: Optional[ReadonlyContext] = None
-  ) -> list[BaseTool]:
+  ) -> List[GoogleApiTool]:
     """Get all tools in the toolset."""
     return [
         GoogleApiTool(
@@ -97,7 +98,7 @@ def set_tool_filter(self, tool_filter: Union[ToolPredicate, List[str]]):
 
   def _load_toolset_with_oidc_auth(self) -> OpenAPIToolset:
     spec_dict = GoogleApiToOpenApiConverter(
-        self.api_name, self.api_version
+        self.api_name, self.api_version, discovery_url=self._discovery_url
     ).convert()
     discovery_scopes = list(
         spec_dict['components']['securitySchemes']['oauth2']['flows'][

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

@@ -32,15 +32,19 @@
 class GoogleApiToOpenApiConverter:
   """Converts Google API Discovery documents to OpenAPI v3 format."""
 
-  def __init__(self, api_name: str, api_version: str):
+  def __init__(
+      self, api_name: str, api_version: str, discovery_url: str | None = None
+  ):
     """Initialize the converter with the API name and version.
 
     Args:
         api_name: The name of the Google API (e.g., "calendar")
         api_version: The version of the API (e.g., "v3")
+        discovery_url: Optional custom discovery document URL.
     """
     self._api_name = api_name
     self._api_version = api_version
+    self._discovery_url = discovery_url
     self._google_api_resource = None
     self._google_api_spec = None
     self._openapi_spec = {
@@ -60,7 +64,14 @@ def fetch_google_api_spec(self) -> None:
           self._api_version,
       )
       # Build a resource object for the specified API
-      self._google_api_resource = build(self._api_name, self._api_version)
+      if self._discovery_url:
+        self._google_api_resource = build(
+            self._api_name,
+            self._api_version,
+            discoveryServiceUrl=self._discovery_url,
+        )
+      else:
+        self._google_api_resource = build(self._api_name, self._api_version)
 
       # Access the underlying API discovery document
       self._google_api_spec = self._google_api_resource._rootDesc

tests/unittests/tools/google_api_tool/test_google_api_toolset.py

@@ -146,7 +146,7 @@ def test_init(
     assert tool_set._additional_headers == additional_headers
 
     mock_converter_class.assert_called_once_with(
-        TEST_API_NAME, TEST_API_VERSION
+        TEST_API_NAME, TEST_API_VERSION, discovery_url=None
     )
     mock_converter_instance.convert.assert_called_once()
     spec_dict = mock_converter_instance.convert.return_value
@@ -187,6 +187,34 @@ def test_init_with_additional_scopes(
     assert isinstance(kwargs["auth_scheme"], OpenIdConnectWithConfig)
     assert kwargs["auth_scheme"].scopes == [DEFAULT_SCOPE] + extra_scopes
 
+  @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_discovery_url(
+      self,
+      mock_converter_class,
+      mock_openapi_toolset_class,
+      mock_converter_instance,
+      mock_openapi_toolset_instance,
+  ):
+    """Test GoogleApiToolset initialization with custom discovery URL."""
+    mock_converter_class.return_value = mock_converter_instance
+    mock_openapi_toolset_class.return_value = mock_openapi_toolset_instance
+
+    discovery_url = "https://example.com/discovery"
+    tool_set = GoogleApiToolset(
+        api_name=TEST_API_NAME,
+        api_version=TEST_API_VERSION,
+        discovery_url=discovery_url,
+    )
+
+    mock_converter_class.assert_called_once_with(
+        TEST_API_NAME, TEST_API_VERSION, discovery_url=discovery_url
+    )
+
   @mock.patch(
       "google.adk.tools.google_api_tool.google_api_toolset.GoogleApiTool"
   )

tests/unittests/tools/google_api_tool/test_googleapi_to_openapi_converter.py

@@ -261,6 +261,27 @@ def test_fetch_google_api_spec(
     # Verify the results
     assert converter_with_patched_build._google_api_spec == calendar_api_spec
 
+  def test_fetch_google_api_spec_with_discovery_url(
+      self, monkeypatch, mock_api_resource, calendar_api_spec
+  ):
+    """Test fetching Google API specification with custom discovery URL."""
+    mock_build = MagicMock(return_value=mock_api_resource)
+    monkeypatch.setattr(
+        "google.adk.tools.google_api_tool.googleapi_to_openapi_converter.build",
+        mock_build,
+    )
+
+    discovery_url = "https://example.com/discovery"
+    converter = GoogleApiToOpenApiConverter(
+        "calendar", "v3", discovery_url=discovery_url
+    )
+    converter.fetch_google_api_spec()
+
+    assert converter._google_api_spec == calendar_api_spec
+    mock_build.assert_called_once_with(
+        "calendar", "v3", discoveryServiceUrl=discovery_url
+    )
+
   def test_fetch_google_api_spec_error(self, monkeypatch, converter):
     """Test error handling when fetching Google API specification."""
     # Create a mock that raises an error