feat: Use json schema for IntegrationConnectorTool declaration when feature enabled

Co-authored-by: Xuan Yang <xygoogle@google.com>
PiperOrigin-RevId: 856508415
This commit is contained in:
Xuan Yang
2026-01-14 23:09:44 -08:00
committed by Copybara-Service
parent ec6abf4010
commit 2ed686527a
2 changed files with 35 additions and 4 deletions
@@ -25,6 +25,8 @@ from typing_extensions import override
from ...auth.auth_credential import AuthCredential from ...auth.auth_credential import AuthCredential
from ...auth.auth_schemes import AuthScheme from ...auth.auth_schemes import AuthScheme
from ...features import FeatureName
from ...features import is_feature_enabled
from .._gemini_schema_util import _to_gemini_schema from .._gemini_schema_util import _to_gemini_schema
from ..base_tool import BaseTool from ..base_tool import BaseTool
from ..openapi_tool.openapi_spec_parser.rest_api_tool import RestApiTool from ..openapi_tool.openapi_spec_parser.rest_api_tool import RestApiTool
@@ -125,10 +127,17 @@ class IntegrationConnectorTool(BaseTool):
if field in schema_dict['required']: if field in schema_dict['required']:
schema_dict['required'].remove(field) schema_dict['required'].remove(field)
parameters = _to_gemini_schema(schema_dict) if is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL):
function_decl = FunctionDeclaration( function_decl = FunctionDeclaration(
name=self.name, description=self.description, parameters=parameters name=self.name,
) description=self.description,
parameters_json_schema=schema_dict,
)
else:
parameters = _to_gemini_schema(schema_dict)
function_decl = FunctionDeclaration(
name=self.name, description=self.description, parameters=parameters
)
return function_decl return function_decl
def _prepare_dynamic_euc(self, auth_credential: AuthCredential) -> str: def _prepare_dynamic_euc(self, auth_credential: AuthCredential) -> str:
@@ -18,6 +18,8 @@ from google.adk.auth.auth_credential import AuthCredential
from google.adk.auth.auth_credential import AuthCredentialTypes from google.adk.auth.auth_credential import AuthCredentialTypes
from google.adk.auth.auth_credential import HttpAuth from google.adk.auth.auth_credential import HttpAuth
from google.adk.auth.auth_credential import HttpCredentials from google.adk.auth.auth_credential import HttpCredentials
from google.adk.features import FeatureName
from google.adk.features._feature_registry import temporary_feature_override
from google.adk.tools.application_integration_tool.integration_connector_tool import IntegrationConnectorTool from google.adk.tools.application_integration_tool.integration_connector_tool import IntegrationConnectorTool
from google.adk.tools.openapi_tool.openapi_spec_parser.rest_api_tool import RestApiTool from google.adk.tools.openapi_tool.openapi_spec_parser.rest_api_tool import RestApiTool
from google.adk.tools.openapi_tool.openapi_spec_parser.tool_auth_handler import AuthPreparationResult from google.adk.tools.openapi_tool.openapi_spec_parser.tool_auth_handler import AuthPreparationResult
@@ -254,3 +256,23 @@ async def test_run_with_auth_async(
args=expected_call_args, tool_context={} args=expected_call_args, tool_context={}
) )
assert result == {"status": "success", "data": "mock_data"} assert result == {"status": "success", "data": "mock_data"}
def test_get_declaration_with_json_schema_feature_enabled(integration_tool):
"""Tests the generation of the function declaration with JSON schema feature enabled."""
with temporary_feature_override(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL, True):
declaration = integration_tool._get_declaration()
assert isinstance(declaration, FunctionDeclaration)
assert declaration.name == "test_integration_tool"
assert declaration.description == "Test integration tool description."
assert declaration.parameters is None
assert declaration.parameters_json_schema == {
"type": "object",
"properties": {
"user_id": {"type": "string", "description": "User ID"},
"page_size": {"type": "integer"},
"filter": {"type": "string"},
},
"required": ["user_id"],
}