google-adk 0.5.0__py3-none-any.whl → 1.0.0__py3-none-any.whl

google/adk/tools/application_integration_tool/integration_connector_tool.py

@@ -12,21 +12,24 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
 import logging
 from typing import Any
 from typing import Dict
 from typing import Optional
+from typing import Union
 
-from google.adk.tools.openapi_tool.openapi_spec_parser.rest_api_tool import RestApiTool
-from google.adk.tools.openapi_tool.openapi_spec_parser.rest_api_tool import to_gemini_schema
 from google.genai.types import FunctionDeclaration
 from typing_extensions import override
 
+from ...auth.auth_credential import AuthCredential
+from ...auth.auth_schemes import AuthScheme
 from .. import BaseTool
+from ..openapi_tool.openapi_spec_parser.rest_api_tool import RestApiTool
+from ..openapi_tool.openapi_spec_parser.rest_api_tool import to_gemini_schema
+from ..openapi_tool.openapi_spec_parser.tool_auth_handler import ToolAuthHandler
 from ..tool_context import ToolContext
 
-logger = logging.getLogger(__name__)
+logger = logging.getLogger('google_adk.' + __name__)
 
 
 class IntegrationConnectorTool(BaseTool):
@@ -56,6 +59,7 @@ class IntegrationConnectorTool(BaseTool):
       'entity',
       'operation',
       'action',
+      'dynamic_auth_config',
   ]
 
   OPTIONAL_FIELDS = [
@@ -75,6 +79,8 @@ class IntegrationConnectorTool(BaseTool):
       operation: str,
       action: str,
       rest_api_tool: RestApiTool,
+      auth_scheme: Optional[Union[AuthScheme, str]] = None,
+      auth_credential: Optional[Union[AuthCredential, str]] = None,
   ):
     """Initializes the ApplicationIntegrationTool.
 
@@ -101,18 +107,20 @@ class IntegrationConnectorTool(BaseTool):
         name=name,
         description=description,
     )
-    self.connection_name = connection_name
-    self.connection_host = connection_host
-    self.connection_service_name = connection_service_name
-    self.entity = entity
-    self.operation = operation
-    self.action = action
-    self.rest_api_tool = rest_api_tool
+    self._connection_name = connection_name
+    self._connection_host = connection_host
+    self._connection_service_name = connection_service_name
+    self._entity = entity
+    self._operation = operation
+    self._action = action
+    self._rest_api_tool = rest_api_tool
+    self._auth_scheme = auth_scheme
+    self._auth_credential = auth_credential
 
   @override
   def _get_declaration(self) -> FunctionDeclaration:
     """Returns the function declaration in the Gemini Schema format."""
-    schema_dict = self.rest_api_tool._operation_parser.get_json_schema()
+    schema_dict = self._rest_api_tool._operation_parser.get_json_schema()
     for field in self.EXCLUDE_FIELDS:
       if field in schema_dict['properties']:
         del schema_dict['properties'][field]
@@ -126,34 +134,69 @@ class IntegrationConnectorTool(BaseTool):
     )
     return function_decl
 
+  def _prepare_dynamic_euc(self, auth_credential: AuthCredential) -> str:
+    if (
+        auth_credential
+        and auth_credential.http
+        and auth_credential.http.credentials
+        and auth_credential.http.credentials.token
+    ):
+      return auth_credential.http.credentials.token
+    return None
+
   @override
   async def run_async(
       self, *, args: dict[str, Any], tool_context: Optional[ToolContext]
   ) -> Dict[str, Any]:
-    args['connection_name'] = self.connection_name
-    args['service_name'] = self.connection_service_name
-    args['host'] = self.connection_host
-    args['entity'] = self.entity
-    args['operation'] = self.operation
-    args['action'] = self.action
+
+    tool_auth_handler = ToolAuthHandler.from_tool_context(
+        tool_context, self._auth_scheme, self._auth_credential
+    )
+    auth_result = tool_auth_handler.prepare_auth_credentials()
+
+    if auth_result.state == 'pending':
+      return {
+          'pending': True,
+          'message': 'Needs your authorization to access your data.',
+      }
+
+    # Attach parameters from auth into main parameters list
+    if auth_result.auth_credential:
+      # Attach parameters from auth into main parameters list
+      auth_credential_token = self._prepare_dynamic_euc(
+          auth_result.auth_credential
+      )
+      if auth_credential_token:
+        args['dynamic_auth_config'] = {
+            'oauth2_auth_code_flow.access_token': auth_credential_token
+        }
+      else:
+        args['dynamic_auth_config'] = {'oauth2_auth_code_flow.access_token': {}}
+
+    args['connection_name'] = self._connection_name
+    args['service_name'] = self._connection_service_name
+    args['host'] = self._connection_host
+    args['entity'] = self._entity
+    args['operation'] = self._operation
+    args['action'] = self._action
     logger.info('Running tool: %s with args: %s', self.name, args)
-    return self.rest_api_tool.call(args=args, tool_context=tool_context)
+    return self._rest_api_tool.call(args=args, tool_context=tool_context)
 
   def __str__(self):
     return (
         f'ApplicationIntegrationTool(name="{self.name}",'
         f' description="{self.description}",'
-        f' connection_name="{self.connection_name}", entity="{self.entity}",'
-        f' operation="{self.operation}", action="{self.action}")'
+        f' connection_name="{self._connection_name}", entity="{self._entity}",'
+        f' operation="{self._operation}", action="{self._action}")'
     )
 
   def __repr__(self):
     return (
         f'ApplicationIntegrationTool(name="{self.name}",'
         f' description="{self.description}",'
-        f' connection_name="{self.connection_name}",'
-        f' connection_host="{self.connection_host}",'
-        f' connection_service_name="{self.connection_service_name}",'
-        f' entity="{self.entity}", operation="{self.operation}",'
-        f' action="{self.action}", rest_api_tool={repr(self.rest_api_tool)})'
+        f' connection_name="{self._connection_name}",'
+        f' connection_host="{self._connection_host}",'
+        f' connection_service_name="{self._connection_service_name}",'
+        f' entity="{self._entity}", operation="{self._operation}",'
+        f' action="{self._action}", rest_api_tool={repr(self._rest_api_tool)})'
     )

google/adk/tools/base_toolset.py

@@ -0,0 +1,58 @@
+from abc import ABC
+from abc import abstractmethod
+from typing import Optional, runtime_checkable
+from typing import Protocol
+
+from ..agents.readonly_context import ReadonlyContext
+from .base_tool import BaseTool
+
+
+@runtime_checkable
+class ToolPredicate(Protocol):
+  """Base class for a predicate that defines the interface to decide whether a
+
+  tool should be exposed to LLM. Toolset implementer could consider whether to
+  accept such instance in the toolset's constructor and apply the predicate in
+  get_tools method.
+  """
+
+  def __call__(
+      self, tool: BaseTool, readonly_context: Optional[ReadonlyContext] = None
+  ) -> bool:
+    """Decide whether the passed-in tool should be exposed to LLM based on the
+
+    current context. True if the tool is usable by the LLM.
+
+    It's used to filter tools in the toolset.
+    """
+
+
+class BaseToolset(ABC):
+  """Base class for toolset.
+
+  A toolset is a collection of tools that can be used by an agent.
+  """
+
+  @abstractmethod
+  async def get_tools(
+      self, readonly_context: Optional[ReadonlyContext] = None
+  ) -> list[BaseTool]:
+    """Return all tools in the toolset based on the provided context.
+
+    Args:
+      readony_context (ReadonlyContext, optional): Context used to filter tools
+        available to the agent. If None, all tools in the toolset are returned.
+
+    Returns:
+      list[BaseTool]: A list of tools available under the specified context.
+    """
+
+  @abstractmethod
+  async def close(self) -> None:
+    """Performs cleanup and releases resources held by the toolset.
+
+    NOTE: This method is invoked, for example, at the end of an agent server's
+    lifecycle or when the toolset is no longer needed. Implementations
+    should ensure that any open connections, files, or other managed
+    resources are properly released to prevent leaks.
+    """

google/adk/tools/enterprise_search_tool.py

@@ -0,0 +1,65 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from google.genai import types
+from typing_extensions import override
+
+from .base_tool import BaseTool
+from .tool_context import ToolContext
+
+if TYPE_CHECKING:
+  from ..models import LlmRequest
+
+
+class EnterpriseWebSearchTool(BaseTool):
+  """A Gemini 2+ built-in tool using web grounding for Enterprise compliance.
+
+  See the documentation for more details:
+  https://cloud.google.com/vertex-ai/generative-ai/docs/grounding/web-grounding-enterprise.
+  """
+
+  def __init__(self):
+    """Initializes the Vertex AI Search tool."""
+    # Name and description are not used because this is a model built-in tool.
+    super().__init__(
+        name='enterprise_web_search', description='enterprise_web_search'
+    )
+
+  @override
+  async def process_llm_request(
+      self,
+      *,
+      tool_context: ToolContext,
+      llm_request: LlmRequest,
+  ) -> None:
+    if llm_request.model and llm_request.model.startswith('gemini-'):
+      if llm_request.model.startswith('gemini-1') and llm_request.config.tools:
+        raise ValueError(
+            'Enterprise web search tool can not be used with other tools in'
+            ' Gemini 1.x.'
+        )
+      llm_request.config = llm_request.config or types.GenerateContentConfig()
+      llm_request.config.tools = llm_request.config.tools or []
+      llm_request.config.tools.append(
+          types.Tool(enterprise_web_search=types.EnterpriseWebSearch())
+      )
+    else:
+      raise ValueError(
+          'Enterprise web search tool is not supported for model'
+          f' {llm_request.model}'
+      )

google/adk/tools/function_parameter_parse_util.py

@@ -35,7 +35,7 @@ _py_builtin_type_to_schema_type = {
     dict: types.Type.OBJECT,
 }
 
-logger = logging.getLogger(__name__)
+logger = logging.getLogger('google_adk.' + __name__)
 
 
 def _is_builtin_primitive_or_compound(
@@ -292,7 +292,7 @@ def _parse_schema_from_parameter(
   raise ValueError(
       f'Failed to parse the parameter {param} of function {func_name} for'
       ' automatic function calling. Automatic function calling works best with'
-      ' simpler function signature schema,consider manually parse your'
+      ' simpler function signature schema, consider manually parsing your'
       f' function declaration for function {func_name}.'
   )
 

google/adk/tools/google_api_tool/__init__.py

@@ -12,76 +12,24 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 __all__ = [
-    'bigquery_tool_set',
-    'calendar_tool_set',
-    'gmail_tool_set',
-    'youtube_tool_set',
-    'slides_tool_set',
-    'sheets_tool_set',
-    'docs_tool_set',
+    'BigQueryToolset',
+    'CalendarToolset',
+    'GmailToolset',
+    'YoutubeToolset',
+    'SlidesToolset',
+    'SheetsToolset',
+    'DocsToolset',
+    'GoogleApiToolset',
+    'GoogleApiTool',
 ]
 
-# Nothing is imported here automatically
-# Each tool set will only be imported when accessed
 
-_bigquery_tool_set = None
-_calendar_tool_set = None
-_gmail_tool_set = None
-_youtube_tool_set = None
-_slides_tool_set = None
-_sheets_tool_set = None
-_docs_tool_set = None
-
-
-def __getattr__(name):
-  global _bigquery_tool_set, _calendar_tool_set, _gmail_tool_set, _youtube_tool_set, _slides_tool_set, _sheets_tool_set, _docs_tool_set
-
-  match name:
-    case 'bigquery_tool_set':
-      if _bigquery_tool_set is None:
-        from .google_api_tool_sets import bigquery_tool_set as bigquery
-
-        _bigquery_tool_set = bigquery
-      return _bigquery_tool_set
-
-    case 'calendar_tool_set':
-      if _calendar_tool_set is None:
-        from .google_api_tool_sets import calendar_tool_set as calendar
-
-        _calendar_tool_set = calendar
-      return _calendar_tool_set
-
-    case 'gmail_tool_set':
-      if _gmail_tool_set is None:
-        from .google_api_tool_sets import gmail_tool_set as gmail
-
-        _gmail_tool_set = gmail
-      return _gmail_tool_set
-
-    case 'youtube_tool_set':
-      if _youtube_tool_set is None:
-        from .google_api_tool_sets import youtube_tool_set as youtube
-
-        _youtube_tool_set = youtube
-      return _youtube_tool_set
-
-    case 'slides_tool_set':
-      if _slides_tool_set is None:
-        from .google_api_tool_sets import slides_tool_set as slides
-
-        _slides_tool_set = slides
-      return _slides_tool_set
-
-    case 'sheets_tool_set':
-      if _sheets_tool_set is None:
-        from .google_api_tool_sets import sheets_tool_set as sheets
-
-        _sheets_tool_set = sheets
-      return _sheets_tool_set
-
-    case 'docs_tool_set':
-      if _docs_tool_set is None:
-        from .google_api_tool_sets import docs_tool_set as docs
-
-        _docs_tool_set = docs
-      return _docs_tool_set
+from .google_api_tool import GoogleApiTool
+from .google_api_toolset import GoogleApiToolset
+from .google_api_toolsets import BigQueryToolset
+from .google_api_toolsets import CalendarToolset
+from .google_api_toolsets import DocsToolset
+from .google_api_toolsets import GmailToolset
+from .google_api_toolsets import SheetsToolset
+from .google_api_toolsets import SlidesToolset
+from .google_api_toolsets import YoutubeToolset

google/adk/tools/google_api_tool/google_api_tool.py

@@ -29,28 +29,34 @@ from ..tool_context import ToolContext
 
 class GoogleApiTool(BaseTool):
 
-  def __init__(self, rest_api_tool: RestApiTool):
+  def __init__(
+      self,
+      rest_api_tool: RestApiTool,
+      client_id: Optional[str] = None,
+      client_secret: Optional[str] = None,
+  ):
     super().__init__(
         name=rest_api_tool.name,
         description=rest_api_tool.description,
         is_long_running=rest_api_tool.is_long_running,
     )
-    self.rest_api_tool = rest_api_tool
+    self._rest_api_tool = rest_api_tool
+    self.configure_auth(client_id, client_secret)
 
   @override
   def _get_declaration(self) -> FunctionDeclaration:
-    return self.rest_api_tool._get_declaration()
+    return self._rest_api_tool._get_declaration()
 
   @override
   async def run_async(
       self, *, args: dict[str, Any], tool_context: Optional[ToolContext]
   ) -> Dict[str, Any]:
-    return await self.rest_api_tool.run_async(
+    return await self._rest_api_tool.run_async(
         args=args, tool_context=tool_context
     )
 
   def configure_auth(self, client_id: str, client_secret: str):
-    self.rest_api_tool.auth_credential = AuthCredential(
+    self._rest_api_tool.auth_credential = AuthCredential(
         auth_type=AuthCredentialTypes.OPEN_ID_CONNECT,
         oauth2=OAuth2Auth(
             client_id=client_id,

google/adk/tools/google_api_tool/google_api_toolset.py

@@ -0,0 +1,126 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import inspect
+import os
+from typing import Any
+from typing import List
+from typing import Optional
+from typing import Type
+from typing import Union
+
+from typing_extensions import override
+
+from ...agents.readonly_context import ReadonlyContext
+from ...auth import OpenIdConnectWithConfig
+from ...tools.base_toolset import BaseToolset
+from ...tools.base_toolset import ToolPredicate
+from ..openapi_tool import OpenAPIToolset
+from .google_api_tool import GoogleApiTool
+from .googleapi_to_openapi_converter import GoogleApiToOpenApiConverter
+
+
+class GoogleApiToolset(BaseToolset):
+  """Google API Toolset contains tools for interacting with Google APIs.
+
+  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.
+  """
+
+  def __init__(
+      self,
+      api_name: str,
+      api_version: str,
+      client_id: Optional[str] = None,
+      client_secret: Optional[str] = None,
+      tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
+  ):
+    self.api_name = api_name
+    self.api_version = api_version
+    self._client_id = client_id
+    self._client_secret = client_secret
+    self._openapi_toolset = self._load_toolset_with_oidc_auth()
+    self.tool_filter = tool_filter
+
+  def _is_tool_selected(
+      self, tool: GoogleApiTool, readonly_context: ReadonlyContext
+  ) -> bool:
+    if not self.tool_filter:
+      return True
+
+    if isinstance(self.tool_filter, ToolPredicate):
+      return self.tool_filter(tool, readonly_context)
+
+    if isinstance(self.tool_filter, list):
+      return tool.name in self.tool_filter
+
+    return False
+
+  @override
+  async def get_tools(
+      self, readonly_context: Optional[ReadonlyContext] = None
+  ) -> List[GoogleApiTool]:
+    """Get all tools in the toolset."""
+    tools = []
+
+    return [
+        GoogleApiTool(tool, self._client_id, self._client_secret)
+        for tool in await self._openapi_toolset.get_tools(readonly_context)
+        if self._is_tool_selected(tool, readonly_context)
+    ]
+
+  def set_tool_filter(self, tool_filter: Union[ToolPredicate, List[str]]):
+    self.tool_filter = tool_filter
+
+  def _load_toolset_with_oidc_auth(self) -> OpenAPIToolset:
+    spec_dict = GoogleApiToOpenApiConverter(
+        self.api_name, self.api_version
+    ).convert()
+    scope = list(
+        spec_dict['components']['securitySchemes']['oauth2']['flows'][
+            'authorizationCode'
+        ]['scopes'].keys()
+    )[0]
+    return OpenAPIToolset(
+        spec_dict=spec_dict,
+        spec_str_type='yaml',
+        auth_scheme=OpenIdConnectWithConfig(
+            authorization_endpoint=(
+                'https://accounts.google.com/o/oauth2/v2/auth'
+            ),
+            token_endpoint='https://oauth2.googleapis.com/token',
+            userinfo_endpoint=(
+                'https://openidconnect.googleapis.com/v1/userinfo'
+            ),
+            revocation_endpoint='https://oauth2.googleapis.com/revoke',
+            token_endpoint_auth_methods_supported=[
+                'client_secret_post',
+                'client_secret_basic',
+            ],
+            grant_types_supported=['authorization_code'],
+            scopes=[scope],
+        ),
+    )
+
+  def configure_auth(self, client_id: str, client_secret: str):
+    self._client_id = client_id
+    self._client_secret = client_secret
+
+  @override
+  async def close(self):
+    if self._openapi_toolset:
+      await self._openapi_toolset.close()

google/adk/tools/google_api_tool/google_api_toolsets.py

@@ -0,0 +1,102 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import logging
+from typing import List
+from typing import Optional
+from typing import Union
+
+from google.adk.tools.base_toolset import ToolPredicate
+
+from .google_api_toolset import GoogleApiToolset
+
+logger = logging.getLogger("google_adk." + __name__)
+
+
+class BigQueryToolset(GoogleApiToolset):
+
+  def __init__(
+      self,
+      client_id: str = None,
+      client_secret: str = None,
+      tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
+  ):
+    super().__init__("bigquery", "v2", client_id, client_secret, tool_filter)
+
+
+class CalendarToolset(GoogleApiToolset):
+
+  def __init__(
+      self,
+      client_id: str = None,
+      client_secret: str = None,
+      tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
+  ):
+    super().__init__("calendar", "v3", client_id, client_secret, tool_filter)
+
+
+class GmailToolset(GoogleApiToolset):
+
+  def __init__(
+      self,
+      client_id: str = None,
+      client_secret: str = None,
+      tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
+  ):
+    super().__init__("gmail", "v1", client_id, client_secret, tool_filter)
+
+
+class YoutubeToolset(GoogleApiToolset):
+
+  def __init__(
+      self,
+      client_id: str = None,
+      client_secret: str = None,
+      tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
+  ):
+    super().__init__("youtube", "v3", client_id, client_secret, tool_filter)
+
+
+class SlidesToolset(GoogleApiToolset):
+
+  def __init__(
+      self,
+      client_id: str = None,
+      client_secret: str = None,
+      tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
+  ):
+    super().__init__("slides", "v1", client_id, client_secret, tool_filter)
+
+
+class SheetsToolset(GoogleApiToolset):
+
+  def __init__(
+      self,
+      client_id: str = None,
+      client_secret: str = None,
+      tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
+  ):
+    super().__init__("sheets", "v4", client_id, client_secret, tool_filter)
+
+
+class DocsToolset(GoogleApiToolset):
+
+  def __init__(
+      self,
+      client_id: str = None,
+      client_secret: str = None,
+      tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
+  ):
+    super().__init__("docs", "v1", client_id, client_secret, tool_filter)