mirror of
https://github.com/encounter/adk-python.git
synced 2026-07-09 18:19:28 -07:00
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:
committed by
Copybara-Service
parent
fd2c0f556b
commit
a5f0d333d7
@@ -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"
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user