adapt application integration toolset to new toolset interface

PiperOrigin-RevId: 757960706
This commit is contained in:
Xiang (Sean) Zhou 2025-05-12 16:36:09 -07:00 committed by Copybara-Service
parent d19927b1bc
commit 57d1315533

View File

@ -12,7 +12,10 @@
# 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
@ -20,6 +23,7 @@ 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 @@ class ApplicationIntegrationToolset:
# 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.
"""Args:
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.
@ -139,6 +107,9 @@ class ApplicationIntegrationToolset:
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 @@ class ApplicationIntegrationToolset:
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 @@ class ApplicationIntegrationToolset:
"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 @@ class ApplicationIntegrationToolset:
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 @@ class ApplicationIntegrationToolset:
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 @@ class ApplicationIntegrationToolset:
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()