feat: Use json schema for RestApiTool declaration when feature enabled

Co-authored-by: Xuan Yang <xygoogle@google.com>
PiperOrigin-RevId: 855767527
This commit is contained in:
Xuan Yang
2026-01-13 09:50:59 -08:00
committed by Copybara-Service
parent fd2c0f556b
commit a5f0d333d7
2 changed files with 54 additions and 4 deletions
@@ -33,6 +33,8 @@ from typing_extensions import override
from ....agents.readonly_context import ReadonlyContext
from ....auth.auth_credential import AuthCredential
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_snake_case
from ...base_tool import BaseTool
@@ -221,10 +223,17 @@ class RestApiTool(BaseTool):
def _get_declaration(self) -> FunctionDeclaration:
"""Returns the function declaration in the Gemini Schema format."""
schema_dict = self._operation_parser.get_json_schema()
parameters = _to_gemini_schema(schema_dict)
function_decl = FunctionDeclaration(
name=self.name, description=self.description, parameters=parameters
)
if is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL):
function_decl = FunctionDeclaration(
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
def configure_auth_scheme(
@@ -29,6 +29,8 @@ from google.adk.auth.auth_credential import AuthCredential
from google.adk.auth.auth_credential import AuthCredentialTypes
from google.adk.auth.auth_credential import HttpAuth
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.sessions.state import State
from google.adk.tools.openapi_tool.auth.auth_helpers import token_to_scheme_credential
from google.adk.tools.openapi_tool.common.common import ApiParameter
@@ -204,6 +206,45 @@ class TestRestApiTool:
assert declaration.description == "Test description"
assert isinstance(declaration.parameters, Schema)
def test_get_declaration_with_json_schema_feature_enabled(
self, sample_endpoint, sample_operation
):
"""Test that _get_declaration uses parameters_json_schema when feature is enabled."""
mock_parser = MagicMock(spec=OperationParser)
mock_parser.get_json_schema.return_value = {
"type": "object",
"properties": {
"test_param": {"type": "string"},
},
"required": ["test_param"],
}
tool = RestApiTool(
name="test_tool",
description="Test description",
endpoint=sample_endpoint,
operation=sample_operation,
should_parse_operation=False,
)
tool._operation_parser = mock_parser
with temporary_feature_override(
FeatureName.JSON_SCHEMA_FOR_FUNC_DECL, True
):
declaration = tool._get_declaration()
assert isinstance(declaration, FunctionDeclaration)
assert declaration.name == "test_tool"
assert declaration.description == "Test description"
assert declaration.parameters is None
assert declaration.parameters_json_schema == {
"type": "object",
"properties": {
"test_param": {"type": "string"},
},
"required": ["test_param"],
}
@patch(
"google.adk.tools.openapi_tool.openapi_spec_parser.rest_api_tool.requests.request"
)