fix: fix application integration toolset

PiperOrigin-RevId: 759018346
This commit is contained in:
Xiang (Sean) Zhou 2025-05-15 00:29:04 -07:00 committed by Copybara-Service
parent 9c2358e2a5
commit 4a0751941b
2 changed files with 13 additions and 5 deletions

View File

@ -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

View File

@ -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