Skip to content

Commit a5f0d33

Browse files
xuanyang15copybara-github
authored andcommitted
feat: Use json schema for RestApiTool declaration when feature enabled
Co-authored-by: Xuan Yang <xygoogle@google.com> PiperOrigin-RevId: 855767527
1 parent fd2c0f5 commit a5f0d33

File tree

2 files changed

+54
-4
lines changed

2 files changed

+54
-4
lines changed

src/google/adk/tools/openapi_tool/openapi_spec_parser/rest_api_tool.py

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,8 @@
3333
from ....agents.readonly_context import ReadonlyContext
3434
from ....auth.auth_credential import AuthCredential
3535
from ....auth.auth_schemes import AuthScheme
36+
from ....features import FeatureName
37+
from ....features import is_feature_enabled
3638
from ..._gemini_schema_util import _to_gemini_schema
3739
from ..._gemini_schema_util import _to_snake_case
3840
from ...base_tool import BaseTool
@@ -221,10 +223,17 @@ def from_parsed_operation_str(
221223
def _get_declaration(self) -> FunctionDeclaration:
222224
"""Returns the function declaration in the Gemini Schema format."""
223225
schema_dict = self._operation_parser.get_json_schema()
224-
parameters = _to_gemini_schema(schema_dict)
225-
function_decl = FunctionDeclaration(
226-
name=self.name, description=self.description, parameters=parameters
227-
)
226+
if is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL):
227+
function_decl = FunctionDeclaration(
228+
name=self.name,
229+
description=self.description,
230+
parameters_json_schema=schema_dict,
231+
)
232+
else:
233+
parameters = _to_gemini_schema(schema_dict)
234+
function_decl = FunctionDeclaration(
235+
name=self.name, description=self.description, parameters=parameters
236+
)
228237
return function_decl
229238

230239
def configure_auth_scheme(

tests/unittests/tools/openapi_tool/openapi_spec_parser/test_rest_api_tool.py

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,8 @@
2929
from google.adk.auth.auth_credential import AuthCredentialTypes
3030
from google.adk.auth.auth_credential import HttpAuth
3131
from google.adk.auth.auth_credential import HttpCredentials
32+
from google.adk.features import FeatureName
33+
from google.adk.features._feature_registry import temporary_feature_override
3234
from google.adk.sessions.state import State
3335
from google.adk.tools.openapi_tool.auth.auth_helpers import token_to_scheme_credential
3436
from google.adk.tools.openapi_tool.common.common import ApiParameter
@@ -204,6 +206,45 @@ def test_get_declaration(
204206
assert declaration.description == "Test description"
205207
assert isinstance(declaration.parameters, Schema)
206208

209+
def test_get_declaration_with_json_schema_feature_enabled(
210+
self, sample_endpoint, sample_operation
211+
):
212+
"""Test that _get_declaration uses parameters_json_schema when feature is enabled."""
213+
mock_parser = MagicMock(spec=OperationParser)
214+
mock_parser.get_json_schema.return_value = {
215+
"type": "object",
216+
"properties": {
217+
"test_param": {"type": "string"},
218+
},
219+
"required": ["test_param"],
220+
}
221+
222+
tool = RestApiTool(
223+
name="test_tool",
224+
description="Test description",
225+
endpoint=sample_endpoint,
226+
operation=sample_operation,
227+
should_parse_operation=False,
228+
)
229+
tool._operation_parser = mock_parser
230+
231+
with temporary_feature_override(
232+
FeatureName.JSON_SCHEMA_FOR_FUNC_DECL, True
233+
):
234+
declaration = tool._get_declaration()
235+
236+
assert isinstance(declaration, FunctionDeclaration)
237+
assert declaration.name == "test_tool"
238+
assert declaration.description == "Test description"
239+
assert declaration.parameters is None
240+
assert declaration.parameters_json_schema == {
241+
"type": "object",
242+
"properties": {
243+
"test_param": {"type": "string"},
244+
},
245+
"required": ["test_param"],
246+
}
247+
207248
@patch(
208249
"google.adk.tools.openapi_tool.openapi_spec_parser.rest_api_tool.requests.request"
209250
)

0 commit comments

Comments
 (0)