adapt application integration toolset to new toolset interface

seanzhougoogle · copybara-github · commit 57d1315533d9 · 2025-05-12T16:36:43.000-07:00
PiperOrigin-RevId: 757960706
diff --git a/src/google/adk/tools/application_integration_tool/application_integration_toolset.py b/src/google/adk/tools/application_integration_tool/application_integration_toolset.py
@@ -12,14 +12,18 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Dict, List, Optional
+from typing import List
+from typing import Optional
+from typing import override
+from typing import Union
 
 from fastapi.openapi.models import HTTPBearer
 
 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 ..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
@@ -63,10 +67,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,
   ])
   ```
   """
@@ -87,46 +91,10 @@ def __init__(
       # tool/python function description.
       tool_instructions: Optional[str] = "",
       service_account_json: Optional[str] = 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={...},
-    )
+    """Args:
 
-    # 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.
@@ -139,6 +107,9 @@ def __init__(
         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
@@ -156,7 +127,7 @@ def __init__(
     self.tool_name = tool_name
     self.tool_instructions = tool_instructions
     self.service_account_json = service_account_json
-    self.generated_tools: Dict[str, RestApiTool] = {}
+    self.tool_filter = tool_filter
 
     integration_client = IntegrationClient(
         project,
@@ -185,10 +156,12 @@ def __init__(
           "Either (integration and trigger) or (connection and"
           " (entity_operations or actions)) should be provided."
       )
-    self._parse_spec_to_tools(spec, connection_details)
+    self.openapi_toolset = None
+    self.tool = None
+    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."""
+  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
@@ -211,13 +184,12 @@ def _parse_spec_to_tools(self, spec_dict, connection_details):
       auth_scheme = HTTPBearer(bearerFormat="JWT")
 
     if self.integration and self.trigger:
-      tools = OpenAPIToolset(
+      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,7 +207,7 @@ def _parse_spec_to_tools(self, spec_dict, connection_details):
         rest_api_tool.configure_auth_scheme(auth_scheme)
       if auth_credential:
         rest_api_tool.configure_auth_credential(auth_credential)
-      tool = IntegrationConnectorTool(
+      self.tool = IntegrationConnectorTool(
           name=rest_api_tool.name,
           description=rest_api_tool.description,
           connection_name=connection_details["name"],
@@ -246,7 +218,12 @@ def _parse_spec_to_tools(self, spec_dict, connection_details):
           operation=operation,
           rest_api_tool=rest_api_tool,
       )
-      self.generated_tools[tool.name] = tool
 
-  def get_tools(self) -> List[RestApiTool]:
-    return list(self.generated_tools.values())
+  @override
+  async def get_tools(self) -> List[RestApiTool]:
+    return [self.tool] if self.tool else await self.openapi_toolset.get_tools()
+
+  @override
+  async def close(self) -> None:
+    if self.openapi_toolset:
+      await self.openapi_toolset.close()
