google-adk 0.0.1__py3-none-any.whl → 0.0.2__py3-none-any.whl

google/adk/tools/application_integration_tool/application_integration_toolset.py

@@ -0,0 +1,230 @@
+# 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 typing import Dict
+from typing import List
+from typing import Optional
+
+from fastapi.openapi.models import HTTPBearer
+from google.adk.tools.application_integration_tool.clients.connections_client import ConnectionsClient
+from google.adk.tools.application_integration_tool.clients.integration_client import IntegrationClient
+from google.adk.tools.openapi_tool.auth.auth_helpers import service_account_scheme_credential
+from google.adk.tools.openapi_tool.openapi_spec_parser.openapi_toolset import OpenAPIToolset
+from google.adk.tools.openapi_tool.openapi_spec_parser.rest_api_tool import RestApiTool
+
+from ...auth.auth_credential import AuthCredential
+from ...auth.auth_credential import AuthCredentialTypes
+from ...auth.auth_credential import ServiceAccount
+from ...auth.auth_credential import ServiceAccountCredential
+
+
+# TODO(cheliu): Apply a common toolset interface
+class ApplicationIntegrationToolset:
+  """ApplicationIntegrationToolset generates tools from a given Application
+
+  Integration or Integration Connector resource.
+  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(),
+  ])
+  ```
+  """
+
+  def __init__(
+      self,
+      project: str,
+      location: str,
+      integration: Optional[str] = None,
+      trigger: Optional[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] = "",
+      # 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,
+  ):
+    """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:
+        project: The GCP project ID.
+        location: The GCP location.
+        integration: The integration name.
+        trigger: The trigger name.
+        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_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.
+
+    Raises:
+        ValueError: If neither integration and trigger nor connection and
+            (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] = {}
+
+    integration_client = IntegrationClient(
+        project,
+        location,
+        integration,
+        trigger,
+        connection,
+        entity_operations,
+        actions,
+        service_account_json,
+    )
+    if integration and trigger:
+      spec = integration_client.get_openapi_spec_for_integration()
+    elif connection and (entity_operations or actions):
+      connections_client = ConnectionsClient(
+          project, location, connection, service_account_json
+      )
+      connection_details = connections_client.get_connection_details()
+      tool_instructions += (
+          "ALWAYS use serviceName = "
+          + connection_details["serviceName"]
+          + ", host = "
+          + connection_details["host"]
+          + " and the connection name = "
+          + f"projects/{project}/locations/{location}/connections/{connection} when"
+          " using this tool"
+          + ". DONOT ask the user for these values as you already have those."
+      )
+      spec = integration_client.get_openapi_spec_for_connection(
+          tool_name,
+          tool_instructions,
+      )
+    else:
+      raise ValueError(
+          "Either (integration and trigger) or (connection and"
+          " (entity_operations or actions)) should be provided."
+      )
+    self._parse_spec_to_tools(spec)
+
+  def _parse_spec_to_tools(self, spec_dict):
+    """Parses the spec dict to a list of RestApiTool."""
+    if self.service_account_json:
+      sa_credential = ServiceAccountCredential.model_validate_json(
+          self.service_account_json
+      )
+      service_account = ServiceAccount(
+          service_account_credential=sa_credential,
+          scopes=["https://www.googleapis.com/auth/cloud-platform"],
+      )
+      auth_scheme, auth_credential = service_account_scheme_credential(
+          config=service_account
+      )
+    else:
+      auth_credential = AuthCredential(
+          auth_type=AuthCredentialTypes.SERVICE_ACCOUNT,
+          service_account=ServiceAccount(
+              use_default_credential=True,
+              scopes=["https://www.googleapis.com/auth/cloud-platform"],
+          ),
+      )
+      auth_scheme = HTTPBearer(bearerFormat="JWT")
+    tools = 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
+
+  def get_tools(self) -> List[RestApiTool]:
+    return list(self.generated_tools.values())