From 4a0751941bf965c1d312cffea07527903687005c Mon Sep 17 00:00:00 2001 From: "Xiang (Sean) Zhou" Date: Thu, 15 May 2025 00:29:04 -0700 Subject: [PATCH] fix: fix application integration toolset PiperOrigin-RevId: 759018346 --- .../application_integration_toolset.py | 13 ++++++++++--- .../test_application_integration_toolset.py | 5 +++-- 2 files changed, 13 insertions(+), 5 deletions(-) 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 index f8b48d1..788999b 100644 --- a/src/google/adk/tools/application_integration_tool/application_integration_toolset.py +++ b/src/google/adk/tools/application_integration_tool/application_integration_toolset.py @@ -19,10 +19,12 @@ 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 ..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 @@ -34,7 +36,7 @@ from .integration_connector_tool import IntegrationConnectorTool # TODO(cheliu): Apply a common toolset interface -class ApplicationIntegrationToolset: +class ApplicationIntegrationToolset(BaseToolset): """ApplicationIntegrationToolset generates tools from a given Application Integration or Integration Connector resource. @@ -223,9 +225,14 @@ class ApplicationIntegrationToolset: ) @override - async def get_tools(self) -> List[RestApiTool]: + async def get_tools( + self, + readonly_context: Optional[ReadonlyContext] = None, + ) -> List[RestApiTool]: return ( - [self._tool] if self._tool else await self._openapi_toolset.get_tools() + [self._tool] + if self._tool + else await self._openapi_toolset.get_tools(readonly_context) ) @override diff --git a/tests/unittests/tools/application_integration_tool/test_application_integration_toolset.py b/tests/unittests/tools/application_integration_tool/test_application_integration_toolset.py index 98fd1a7..3dc2393 100644 --- a/tests/unittests/tools/application_integration_tool/test_application_integration_toolset.py +++ b/tests/unittests/tools/application_integration_tool/test_application_integration_toolset.py @@ -15,6 +15,7 @@ import json from unittest import mock from fastapi.openapi.models import Operation +from google.adk.agents.readonly_context import ReadonlyContext from google.adk.auth.auth_credential import AuthCredential from google.adk.tools.application_integration_tool.application_integration_toolset import ApplicationIntegrationToolset from google.adk.tools.application_integration_tool.integration_connector_tool import IntegrationConnectorTool @@ -49,7 +50,7 @@ def mock_openapi_toolset(): mock_rest_api_tool.name = "Test Tool" # Create an async mock for the get_tools method - async def mock_get_tools(): + async def mock_get_tools(context: ReadonlyContext = None): return [mock_rest_api_tool] # Assign the async mock function to get_tools @@ -71,7 +72,7 @@ def mock_openapi_toolset_with_multiple_tools_and_no_tools(): mock_rest_api_tool_2.name = "Test Tool 2" # Create an async mock for the get_tools method - async def mock_get_tools(): + async def mock_get_tools(context: ReadonlyContext = None): return [mock_rest_api_tool, mock_rest_api_tool_2] mock_toolset_instance.get_tools = mock_get_tools