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

google/adk/tools/apihub_tool/apihub_toolset.py

@@ -13,19 +13,25 @@
 # limitations under the License.
 
 
-from typing import Dict, List, Optional
+from typing import List
+from typing import Optional
+from typing import Union
 
+from typing_extensions import override
 import yaml
 
+from ...agents.readonly_context import ReadonlyContext
 from ...auth.auth_credential import AuthCredential
 from ...auth.auth_schemes import AuthScheme
+from ..base_toolset import BaseToolset
+from ..base_toolset import ToolPredicate
 from ..openapi_tool.common.common import to_snake_case
 from ..openapi_tool.openapi_spec_parser.openapi_toolset import OpenAPIToolset
 from ..openapi_tool.openapi_spec_parser.rest_api_tool import RestApiTool
 from .clients.apihub_client import APIHubClient
 
 
-class APIHubToolset:
+class APIHubToolset(BaseToolset):
   """APIHubTool generates tools from a given API Hub resource.
 
   Examples:
@@ -34,16 +40,13 @@ class APIHubToolset:
   apihub_toolset = APIHubToolset(
       apihub_resource_name="projects/test-project/locations/us-central1/apis/test-api",
       service_account_json="...",
+      tool_filter=lambda tool, ctx=None: tool.name in ('my_tool',
+      'my_other_tool')
   )
 
   # Get all available tools
-  agent = LlmAgent(tools=apihub_toolset.get_tools())
+  agent = LlmAgent(tools=apihub_toolset)
 
-  # Get a specific tool
-  agent = LlmAgent(tools=[
-      ...
-      apihub_toolset.get_tool('my_tool'),
-  ])
   ```
 
   **apihub_resource_name** is the resource name from API Hub. It must include
@@ -70,6 +73,7 @@ class APIHubToolset:
       auth_credential: Optional[AuthCredential] = None,
       # Optionally, you can provide a custom API Hub client
       apihub_client: Optional[APIHubClient] = None,
+      tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
   ):
     """Initializes the APIHubTool with the given parameters.
 
@@ -81,12 +85,17 @@ class APIHubToolset:
     )
 
     # Get all available tools
-    agent = LlmAgent(tools=apihub_toolset.get_tools())
+    agent = LlmAgent(tools=[apihub_toolset])
 
+    apihub_toolset = APIHubToolset(
+        apihub_resource_name="projects/test-project/locations/us-central1/apis/test-api",
+        service_account_json="...",
+        tool_filter = ['my_tool']
+    )
     # Get a specific tool
     agent = LlmAgent(tools=[
-        ...
-        apihub_toolset.get_tool('my_tool'),
+        ...,
+        apihub_toolset,
     ])
     ```
 
@@ -118,82 +127,49 @@ class APIHubToolset:
         lazy_load_spec: If True, the spec will be loaded lazily when needed.
           Otherwise, the spec will be loaded immediately and the tools will be
           generated during initialization.
+        tool_filter: The filter used to filter the tools in the toolset. It can
+          be either a tool predicate or a list of tool names of the tools to
+          expose.
     """
     self.name = name
     self.description = description
-    self.apihub_resource_name = apihub_resource_name
-    self.lazy_load_spec = lazy_load_spec
-    self.apihub_client = apihub_client or APIHubClient(
+    self._apihub_resource_name = apihub_resource_name
+    self._lazy_load_spec = lazy_load_spec
+    self._apihub_client = apihub_client or APIHubClient(
         access_token=access_token,
         service_account_json=service_account_json,
     )
 
-    self.generated_tools: Dict[str, RestApiTool] = {}
-    self.auth_scheme = auth_scheme
-    self.auth_credential = auth_credential
-
-    if not self.lazy_load_spec:
-      self._prepare_tools()
-
-  def get_tool(self, name: str) -> Optional[RestApiTool]:
-    """Retrieves a specific tool by its name.
-
-    Example:
-    ```
-    apihub_tool = apihub_toolset.get_tool('my_tool')
-    ```
-
-    Args:
-        name: The name of the tool to retrieve.
-
-    Returns:
-        The tool with the given name, or None if no such tool exists.
-    """
-    if not self._are_tools_ready():
-      self._prepare_tools()
+    self._openapi_toolset = None
+    self._auth_scheme = auth_scheme
+    self._auth_credential = auth_credential
+    self.tool_filter = tool_filter
 
-    return self.generated_tools[name] if name in self.generated_tools else None
+    if not self._lazy_load_spec:
+      self._prepare_toolset()
 
-  def get_tools(self) -> List[RestApiTool]:
+  @override
+  async def get_tools(
+      self, readonly_context: Optional[ReadonlyContext] = None
+  ) -> List[RestApiTool]:
     """Retrieves all available tools.
 
     Returns:
         A list of all available RestApiTool objects.
     """
-    if not self._are_tools_ready():
-      self._prepare_tools()
-
-    return list(self.generated_tools.values())
-
-  def _are_tools_ready(self) -> bool:
-    return not self.lazy_load_spec or self.generated_tools
-
-  def _prepare_tools(self) -> str:
-    """Fetches the spec from API Hub and generates the tools.
+    if not self._openapi_toolset:
+      self._prepare_toolset()
+    if not self._openapi_toolset:
+      return []
+    return await self._openapi_toolset.get_tools(readonly_context)
 
-    Returns:
-        True if the tools are ready, False otherwise.
-    """
+  def _prepare_toolset(self) -> None:
+    """Fetches the spec from API Hub and generates the toolset."""
     # For each API, get the first version and the first spec of that version.
-    spec = self.apihub_client.get_spec_content(self.apihub_resource_name)
-    self.generated_tools: Dict[str, RestApiTool] = {}
-
-    tools = self._parse_spec_to_tools(spec)
-    for tool in tools:
-      self.generated_tools[tool.name] = tool
-
-  def _parse_spec_to_tools(self, spec_str: str) -> List[RestApiTool]:
-    """Parses the spec string to a list of RestApiTool.
-
-    Args:
-        spec_str: The spec string to parse.
-
-    Returns:
-        A list of RestApiTool objects.
-    """
+    spec_str = self._apihub_client.get_spec_content(self._apihub_resource_name)
     spec_dict = yaml.safe_load(spec_str)
     if not spec_dict:
-      return []
+      return
 
     self.name = self.name or to_snake_case(
         spec_dict.get('info', {}).get('title', 'unnamed')
@@ -201,9 +177,14 @@ class APIHubToolset:
     self.description = self.description or spec_dict.get('info', {}).get(
         'description', ''
     )
-    tools = OpenAPIToolset(
+    self._openapi_toolset = OpenAPIToolset(
         spec_dict=spec_dict,
-        auth_credential=self.auth_credential,
-        auth_scheme=self.auth_scheme,
-    ).get_tools()
-    return tools
+        auth_credential=self._auth_credential,
+        auth_scheme=self._auth_scheme,
+        tool_filter=self.tool_filter,
+    )
+
+  @override
+  async def close(self):
+    if self._openapi_toolset:
+      await self._openapi_toolset.close()

google/adk/tools/application_integration_tool/application_integration_toolset.py

@@ -12,14 +12,22 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Dict, List, Optional
+import logging
+from typing import List
+from typing import Optional
+from typing import Union
 
 from fastapi.openapi.models import HTTPBearer
+from typing_extensions import override
 
+from ...agents.readonly_context import ReadonlyContext
 from ...auth.auth_credential import AuthCredential
 from ...auth.auth_credential import AuthCredentialTypes
 from ...auth.auth_credential import ServiceAccount
 from ...auth.auth_credential import ServiceAccountCredential
+from ...auth.auth_schemes import AuthScheme
+from ..base_toolset import BaseToolset
+from ..base_toolset import ToolPredicate
 from ..openapi_tool.auth.auth_helpers import service_account_scheme_credential
 from ..openapi_tool.openapi_spec_parser.openapi_spec_parser import OpenApiSpecParser
 from ..openapi_tool.openapi_spec_parser.openapi_toolset import OpenAPIToolset
@@ -28,9 +36,11 @@ from .clients.connections_client import ConnectionsClient
 from .clients.integration_client import IntegrationClient
 from .integration_connector_tool import IntegrationConnectorTool
 
+logger = logging.getLogger("google_adk." + __name__)
+
 
 # TODO(cheliu): Apply a common toolset interface
-class ApplicationIntegrationToolset:
+class ApplicationIntegrationToolset(BaseToolset):
   """ApplicationIntegrationToolset generates tools from a given Application
 
   Integration or Integration Connector resource.
@@ -42,7 +52,7 @@ class ApplicationIntegrationToolset:
       project="test-project",
       location="us-central1"
       integration="test-integration",
-      trigger="api_trigger/test_trigger",
+      triggers=["api_trigger/test_trigger"],
       service_account_credentials={...},
   )
 
@@ -63,10 +73,10 @@ class ApplicationIntegrationToolset:
       service_account_credentials={...},
   )
 
-  # Get all available tools
+  # Feed the toolset to agent
   agent = LlmAgent(tools=[
-      ...
-      *application_integration_toolset.get_tools(),
+      ...,
+      application_integration_toolset,
   ])
   ```
   """
@@ -76,100 +86,74 @@ class ApplicationIntegrationToolset:
       project: str,
       location: str,
       integration: Optional[str] = None,
-      trigger: Optional[str] = None,
+      triggers: Optional[List[str]] = None,
       connection: Optional[str] = None,
       entity_operations: Optional[str] = None,
       actions: Optional[str] = None,
       # Optional parameter for the toolset. This is prepended to the generated
       # tool/python function name.
-      tool_name: Optional[str] = "",
+      tool_name_prefix: Optional[str] = "",
       # Optional parameter for the toolset. This is appended to the generated
       # tool/python function description.
       tool_instructions: Optional[str] = "",
       service_account_json: Optional[str] = None,
+      auth_scheme: Optional[AuthScheme] = None,
+      auth_credential: Optional[AuthCredential] = None,
+      tool_filter: Optional[Union[ToolPredicate, List[str]]] = None,
   ):
-    """Initializes the ApplicationIntegrationToolset.
-
-    Example Usage:
-    ```
-    # Get all available tools for an integration with api trigger
-    application_integration_toolset = ApplicationIntegrationToolset(
-
-        project="test-project",
-        location="us-central1"
-        integration="test-integration",
-        trigger="api_trigger/test_trigger",
-        service_account_credentials={...},
-    )
-
-    # Get all available tools for a connection using entity operations and
-    # actions
-    # Note: Find the list of supported entity operations and actions for a
-    connection
-    # using integration connector apis:
-    #
-    https://cloud.google.com/integration-connectors/docs/reference/rest/v1/projects.locations.connections.connectionSchemaMetadata
-    application_integration_toolset = ApplicationIntegrationToolset(
-        project="test-project",
-        location="us-central1"
-        connection="test-connection",
-        entity_operations=["EntityId1": ["LIST","CREATE"], "EntityId2": []],
-        #empty list for actions means all operations on the entity are supported
-        actions=["action1"],
-        service_account_credentials={...},
-    )
-
-    # Get all available tools
-    agent = LlmAgent(tools=[
-        ...
-        *application_integration_toolset.get_tools(),
-    ])
-    ```
+    """Args:
 
     Args:
         project: The GCP project ID.
         location: The GCP location.
         integration: The integration name.
-        trigger: The trigger name.
+        triggers: The list of trigger names in the integration.
         connection: The connection name.
         entity_operations: The entity operations supported by the connection.
         actions: The actions supported by the connection.
-        tool_name: The name of the tool.
+        tool_name_prefix: The name prefix of the generated tools.
         tool_instructions: The instructions for the tool.
         service_account_json: The service account configuration as a dictionary.
           Required if not using default service credential. Used for fetching
           the Application Integration or Integration Connector resource.
+        tool_filter: The filter used to filter the tools in the toolset. It can
+          be either a tool predicate or a list of tool names of the tools to
+          expose.
 
     Raises:
-        ValueError: If neither integration and trigger nor connection and
-            (entity_operations or actions) is provided.
+        ValueError: If none of the following conditions are met:
+            - `integration` is provided.
+            - `connection` is provided and at least one of `entity_operations`
+              or `actions` is provided.
         Exception: If there is an error during the initialization of the
             integration or connection client.
     """
     self.project = project
     self.location = location
-    self.integration = integration
-    self.trigger = trigger
-    self.connection = connection
-    self.entity_operations = entity_operations
-    self.actions = actions
-    self.tool_name = tool_name
-    self.tool_instructions = tool_instructions
-    self.service_account_json = service_account_json
-    self.generated_tools: Dict[str, RestApiTool] = {}
+    self._integration = integration
+    self._triggers = triggers
+    self._connection = connection
+    self._entity_operations = entity_operations
+    self._actions = actions
+    self._tool_name_prefix = tool_name_prefix
+    self._tool_instructions = tool_instructions
+    self._service_account_json = service_account_json
+    self._auth_scheme = auth_scheme
+    self._auth_credential = auth_credential
+    self.tool_filter = tool_filter
 
     integration_client = IntegrationClient(
         project,
         location,
         integration,
-        trigger,
+        triggers,
         connection,
         entity_operations,
         actions,
         service_account_json,
     )
     connection_details = {}
-    if integration and trigger:
+    if integration:
       spec = integration_client.get_openapi_spec_for_integration()
     elif connection and (entity_operations or actions):
       connections_client = ConnectionsClient(
@@ -177,21 +161,23 @@ class ApplicationIntegrationToolset:
       )
       connection_details = connections_client.get_connection_details()
       spec = integration_client.get_openapi_spec_for_connection(
-          tool_name,
+          tool_name_prefix,
           tool_instructions,
       )
     else:
       raise ValueError(
-          "Either (integration and trigger) or (connection and"
+          "Invalid request, Either integration or (connection and"
           " (entity_operations or actions)) should be provided."
       )
-    self._parse_spec_to_tools(spec, connection_details)
+    self._openapi_toolset = None
+    self._tools = []
+    self._parse_spec_to_toolset(spec, connection_details)
 
-  def _parse_spec_to_tools(self, spec_dict, connection_details):
-    """Parses the spec dict to a list of RestApiTool."""
-    if self.service_account_json:
+  def _parse_spec_to_toolset(self, spec_dict, connection_details):
+    """Parses the spec dict to OpenAPI toolset."""
+    if self._service_account_json:
       sa_credential = ServiceAccountCredential.model_validate_json(
-          self.service_account_json
+          self._service_account_json
       )
       service_account = ServiceAccount(
           service_account_credential=sa_credential,
@@ -210,14 +196,13 @@ class ApplicationIntegrationToolset:
       )
       auth_scheme = HTTPBearer(bearerFormat="JWT")
 
-    if self.integration and self.trigger:
-      tools = OpenAPIToolset(
+    if self._integration:
+      self._openapi_toolset = OpenAPIToolset(
           spec_dict=spec_dict,
           auth_credential=auth_credential,
           auth_scheme=auth_scheme,
-      ).get_tools()
-      for tool in tools:
-        self.generated_tools[tool.name] = tool
+          tool_filter=self.tool_filter,
+      )
       return
 
     operations = OpenApiSpecParser().parse(spec_dict)
@@ -235,18 +220,55 @@ class ApplicationIntegrationToolset:
         rest_api_tool.configure_auth_scheme(auth_scheme)
       if auth_credential:
         rest_api_tool.configure_auth_credential(auth_credential)
-      tool = IntegrationConnectorTool(
-          name=rest_api_tool.name,
-          description=rest_api_tool.description,
-          connection_name=connection_details["name"],
-          connection_host=connection_details["host"],
-          connection_service_name=connection_details["serviceName"],
-          entity=entity,
-          action=action,
-          operation=operation,
-          rest_api_tool=rest_api_tool,
+
+      auth_override_enabled = connection_details.get(
+          "authOverrideEnabled", False
       )
-      self.generated_tools[tool.name] = tool
 
-  def get_tools(self) -> List[RestApiTool]:
-    return list(self.generated_tools.values())
+      if (
+          self._auth_scheme
+          and self._auth_credential
+          and not auth_override_enabled
+      ):
+        # Case: Auth provided, but override is OFF. Don't use provided auth.
+        logger.warning(
+            "Authentication schema and credentials are not used because"
+            " authOverrideEnabled is not enabled in the connection."
+        )
+        connector_auth_scheme = None
+        connector_auth_credential = None
+      else:
+        connector_auth_scheme = self._auth_scheme
+        connector_auth_credential = self._auth_credential
+
+      self._tools.append(
+          IntegrationConnectorTool(
+              name=rest_api_tool.name,
+              description=rest_api_tool.description,
+              connection_name=connection_details["name"],
+              connection_host=connection_details["host"],
+              connection_service_name=connection_details["serviceName"],
+              entity=entity,
+              action=action,
+              operation=operation,
+              rest_api_tool=rest_api_tool,
+              auth_scheme=connector_auth_scheme,
+              auth_credential=connector_auth_credential,
+          )
+      )
+
+  @override
+  async def get_tools(
+      self,
+      readonly_context: Optional[ReadonlyContext] = None,
+  ) -> List[RestApiTool]:
+    return (
+        self._tools
+        if self._openapi_toolset is None
+        else await self._openapi_toolset.get_tools(readonly_context)
+    )
+
+  @override
+  async def close(self) -> None:
+    if self._openapi_toolset:
+      await self._openapi_toolset.close()

google/adk/tools/application_integration_tool/clients/connections_client.py

@@ -554,6 +554,9 @@ class ConnectionsClient:
             "serviceName": {"$ref": "#/components/schemas/serviceName"},
             "host": {"$ref": "#/components/schemas/host"},
             "entity": {"$ref": "#/components/schemas/entity"},
+            "dynamicAuthConfig": {
+                "$ref": "#/components/schemas/dynamicAuthConfig"
+            },
         },
     }
 
@@ -580,6 +583,10 @@ class ConnectionsClient:
             "serviceName": {"$ref": "#/components/schemas/serviceName"},
             "host": {"$ref": "#/components/schemas/host"},
             "entity": {"$ref": "#/components/schemas/entity"},
+            "dynamicAuthConfig": {
+                "$ref": "#/components/schemas/dynamicAuthConfig"
+            },
+            "filterClause": {"$ref": "#/components/schemas/filterClause"},
         },
     }
 
@@ -602,6 +609,9 @@ class ConnectionsClient:
             "serviceName": {"$ref": "#/components/schemas/serviceName"},
             "host": {"$ref": "#/components/schemas/host"},
             "entity": {"$ref": "#/components/schemas/entity"},
+            "dynamicAuthConfig": {
+                "$ref": "#/components/schemas/dynamicAuthConfig"
+            },
         },
     }
 
@@ -624,6 +634,10 @@ class ConnectionsClient:
             "serviceName": {"$ref": "#/components/schemas/serviceName"},
             "host": {"$ref": "#/components/schemas/host"},
             "entity": {"$ref": "#/components/schemas/entity"},
+            "dynamicAuthConfig": {
+                "$ref": "#/components/schemas/dynamicAuthConfig"
+            },
+            "filterClause": {"$ref": "#/components/schemas/filterClause"},
         },
     }
 
@@ -647,6 +661,9 @@ class ConnectionsClient:
             "serviceName": {"$ref": "#/components/schemas/serviceName"},
             "host": {"$ref": "#/components/schemas/host"},
             "entity": {"$ref": "#/components/schemas/entity"},
+            "dynamicAuthConfig": {
+                "$ref": "#/components/schemas/dynamicAuthConfig"
+            },
         },
     }
 
@@ -671,6 +688,9 @@ class ConnectionsClient:
             "connectorInputPayload": {
                 "$ref": f"#/components/schemas/connectorInputPayload_{action}"
             },
+            "dynamicAuthConfig": {
+                "$ref": "#/components/schemas/dynamicAuthConfig"
+            },
         },
     }
 

google/adk/tools/application_integration_tool/clients/integration_client.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import json
-from typing import Optional
+from typing import List, Optional
 from google.adk.tools.application_integration_tool.clients.connections_client import ConnectionsClient
 import google.auth
 from google.auth import default as default_service_credential
@@ -35,7 +35,7 @@ class IntegrationClient:
       project: str,
       location: str,
       integration: Optional[str] = None,
-      trigger: Optional[str] = None,
+      triggers: Optional[List[str]] = None,
       connection: Optional[str] = None,
       entity_operations: Optional[dict[str, list[str]]] = None,
       actions: Optional[list[str]] = None,
@@ -47,7 +47,7 @@ class IntegrationClient:
         project: The Google Cloud project ID.
         location: The Google Cloud location (e.g., us-central1).
         integration: The integration name.
-        trigger: The trigger ID for the integration.
+        triggers: The list of trigger IDs for the integration.
         connection: The connection name.
         entity_operations: A dictionary mapping entity names to a list of
           operations (e.g., LIST, CREATE, UPDATE, DELETE, GET).
@@ -59,7 +59,7 @@ class IntegrationClient:
     self.project = project
     self.location = location
     self.integration = integration
-    self.trigger = trigger
+    self.triggers = triggers
     self.connection = connection
     self.entity_operations = (
         entity_operations if entity_operations is not None else {}
@@ -88,7 +88,7 @@ class IntegrationClient:
           "apiTriggerResources": [
               {
                   "integrationResource": self.integration,
-                  "triggerId": [self.trigger],
+                  "triggerId": self.triggers,
               },
           ],
           "fileFormat": "JSON",
@@ -109,7 +109,7 @@ class IntegrationClient:
         raise ValueError(
             "Invalid request. Please check the provided values of"
             f" project({self.project}), location({self.location}),"
-            f" integration({self.integration}) and trigger({self.trigger})."
+            f" integration({self.integration})."
         ) from e
       raise ValueError(f"Request error: {e}") from e
     except Exception as e: