From cb3244bb58904ab508f77069b436f85b442d3299 Mon Sep 17 00:00:00 2001 From: Xuan Yang Date: Thu, 11 Dec 2025 13:24:07 -0800 Subject: [PATCH] feat: Use json schema for function tool declaration when feature enabled Co-authored-by: Xuan Yang PiperOrigin-RevId: 843350330 --- .../tools/_automatic_function_calling_util.py | 17 ++ .../tools/test_build_function_declaration.py | 249 +++++++++++++++++- 2 files changed, 260 insertions(+), 6 deletions(-) diff --git a/src/google/adk/tools/_automatic_function_calling_util.py b/src/google/adk/tools/_automatic_function_calling_util.py index 99e6c9f3..92df8871 100644 --- a/src/google/adk/tools/_automatic_function_calling_util.py +++ b/src/google/adk/tools/_automatic_function_calling_util.py @@ -30,6 +30,9 @@ from pydantic import create_model from pydantic import fields as pydantic_fields from . import _function_parameter_parse_util +from . import _function_tool_declarations +from ..features import FeatureName +from ..features import is_feature_enabled from ..utils.variant_utils import GoogleLLMVariant _py_type_2_schema_type = { @@ -196,6 +199,20 @@ def build_function_declaration( ignore_params: Optional[list[str]] = None, variant: GoogleLLMVariant = GoogleLLMVariant.GEMINI_API, ) -> types.FunctionDeclaration: + # ========== Pydantic-based function tool declaration (new feature) ========== + if is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL): + declaration = ( + _function_tool_declarations.build_function_declaration_with_json_schema( + func, ignore_params=ignore_params + ) + ) + # Add response schema only for VERTEX_AI + # TODO(b/421991354): Remove this check once the bug is fixed. + if variant != GoogleLLMVariant.VERTEX_AI: + declaration.response_json_schema = None + return declaration + + # ========== ADK defined function tool declaration (old behavior) ========== signature = inspect.signature(func) should_update_signature = False new_func = None diff --git a/tests/unittests/tools/test_build_function_declaration.py b/tests/unittests/tools/test_build_function_declaration.py index 8a562c7c..dd85b20c 100644 --- a/tests/unittests/tools/test_build_function_declaration.py +++ b/tests/unittests/tools/test_build_function_declaration.py @@ -13,9 +13,9 @@ # limitations under the License. from enum import Enum -from typing import Dict -from typing import List +from unittest import mock +from google.adk.features import FeatureName from google.adk.tools import _automatic_function_calling_util from google.adk.tools.tool_context import ToolContext from google.adk.utils.variant_utils import GoogleLLMVariant @@ -79,7 +79,7 @@ def test_bool_input(): def test_array_input(): - def simple_function(input_str: List[str]) -> str: + def simple_function(input_str: list[str]) -> str: return {'result': input_str} function_decl = _automatic_function_calling_util.build_function_declaration( @@ -92,7 +92,7 @@ def test_array_input(): def test_dict_input(): - def simple_function(input_str: Dict[str, str]) -> str: + def simple_function(input_str: dict[str, str]) -> str: return {'result': input_str} function_decl = _automatic_function_calling_util.build_function_declaration( @@ -206,7 +206,7 @@ def test_basemodel_with_nested_basemodel(): def test_list(): def simple_function( - input_str: List[str], input_dir: List[Dict[str, str]] + input_str: list[str], input_dir: list[dict[str, str]] ) -> str: return {'result': input_str} @@ -257,7 +257,7 @@ def test_basemodel_list(): class CustomInput(BaseModel): child: ChildInput - def simple_function(input_str: List[CustomInput]) -> str: + def simple_function(input_str: list[CustomInput]) -> str: return {'result': input_str} function_decl = _automatic_function_calling_util.build_function_declaration( @@ -427,3 +427,240 @@ def test_transfer_to_agent_tool_with_enum_constraint(): assert function_decl.parameters.properties['agent_name'].type == 'STRING' assert function_decl.parameters.properties['agent_name'].enum == agent_names assert 'tool_context' not in function_decl.parameters.properties + + +class TestJsonSchemaFeatureFlagEnabled: + """Tests for build_function_declaration when JSON_SCHEMA_FOR_FUNC_DECL is enabled.""" + + @pytest.fixture(autouse=True) + def enable_feature_flag(self): + """Enable the JSON_SCHEMA_FOR_FUNC_DECL feature flag for all tests.""" + with mock.patch.object( + _automatic_function_calling_util, + 'is_feature_enabled', + autospec=True, + side_effect=lambda f: f == FeatureName.JSON_SCHEMA_FOR_FUNC_DECL, + ): + yield + + def test_basic_string_parameter(self): + """Test basic string parameter with feature flag enabled.""" + + def greet(name: str) -> str: + """Greet someone.""" + return f'Hello, {name}!' + + decl = _automatic_function_calling_util.build_function_declaration(greet) + + assert decl.name == 'greet' + assert decl.description == 'Greet someone.' + assert decl.parameters_json_schema == { + 'properties': {'name': {'title': 'Name', 'type': 'string'}}, + 'required': ['name'], + 'title': 'greetParams', + 'type': 'object', + } + + def test_multiple_parameter_types(self): + """Test multiple parameter types with feature flag enabled.""" + + def create_user(name: str, age: int, active: bool) -> str: + """Create a new user.""" + return f'Created {name}' + + decl = _automatic_function_calling_util.build_function_declaration( + create_user + ) + + schema = decl.parameters_json_schema + assert schema['properties'] == { + 'name': {'title': 'Name', 'type': 'string'}, + 'age': {'title': 'Age', 'type': 'integer'}, + 'active': {'title': 'Active', 'type': 'boolean'}, + } + assert set(schema['required']) == {'name', 'age', 'active'} + + def test_list_parameter(self): + """Test list parameter with feature flag enabled.""" + + def sum_numbers(numbers: list[int]) -> int: + """Sum a list of numbers.""" + return sum(numbers) + + decl = _automatic_function_calling_util.build_function_declaration( + sum_numbers + ) + + schema = decl.parameters_json_schema + assert schema['properties']['numbers'] == { + 'items': {'type': 'integer'}, + 'title': 'Numbers', + 'type': 'array', + } + + def test_dict_parameter(self): + """Test dict parameter with feature flag enabled.""" + + def process_data(data: dict[str, str]) -> str: + """Process a dictionary.""" + return str(data) + + decl = _automatic_function_calling_util.build_function_declaration( + process_data + ) + + schema = decl.parameters_json_schema + assert schema['properties']['data'] == { + 'additionalProperties': {'type': 'string'}, + 'title': 'Data', + 'type': 'object', + } + + def test_optional_parameter(self): + """Test optional parameter with feature flag enabled.""" + + def search(query: str, limit: int | None = None) -> str: + """Search for something.""" + return query + + decl = _automatic_function_calling_util.build_function_declaration(search) + + schema = decl.parameters_json_schema + assert schema['required'] == ['query'] + assert 'query' in schema['properties'] + assert 'limit' in schema['properties'] + + def test_enum_parameter(self): + """Test enum parameter with feature flag enabled.""" + + class Color(Enum): + RED = 'red' + GREEN = 'green' + BLUE = 'blue' + + def set_color(color: Color) -> str: + """Set the color.""" + return color.value + + decl = _automatic_function_calling_util.build_function_declaration( + set_color + ) + + schema = decl.parameters_json_schema + assert schema['properties']['color'] == { + '$ref': '#/$defs/Color', + } + assert schema['$defs']['Color'] == { + 'enum': ['red', 'green', 'blue'], + 'title': 'Color', + 'type': 'string', + } + + def test_tool_context_ignored(self): + """Test that tool_context is ignored.""" + + def my_tool(query: str, tool_context: ToolContext) -> str: + """A tool that uses context.""" + return query + + decl = _automatic_function_calling_util.build_function_declaration( + my_tool, ignore_params=['tool_context'] + ) + + schema = decl.parameters_json_schema + assert set(schema['properties'].keys()) == {'query'} + assert 'tool_context' not in schema['properties'] + + def test_gemini_api_no_response_schema(self): + """Test that GEMINI_API variant does not include response schema.""" + + def get_data() -> dict[str, int]: + """Get some data.""" + return {'count': 42} + + decl = _automatic_function_calling_util.build_function_declaration( + get_data, variant=GoogleLLMVariant.GEMINI_API + ) + + # GEMINI_API should not have response_json_schema due to bug b/421991354 + assert decl.response_json_schema is None + + @pytest.mark.parametrize( + 'variant, expect_response_schema', + [ + (GoogleLLMVariant.GEMINI_API, False), + (GoogleLLMVariant.VERTEX_AI, True), + ], + ) + def test_response_schema_by_variant(self, variant, expect_response_schema): + """Test response schema generation based on the LLM variant.""" + + def get_data() -> dict[str, int]: + """Get some data.""" + return {'count': 42} + + decl = _automatic_function_calling_util.build_function_declaration( + get_data, variant=variant + ) + + assert (decl.response_json_schema is not None) == expect_response_schema + + def test_pydantic_model_parameter(self): + """Test Pydantic model parameter with feature flag enabled.""" + + class Address(BaseModel): + street: str + city: str + + def save_address(address: Address) -> str: + """Save an address.""" + return f'Saved address in {address.city}' + + decl = _automatic_function_calling_util.build_function_declaration( + save_address + ) + + assert decl.parameters_json_schema is not None + assert 'address' in decl.parameters_json_schema['properties'] + + def test_no_parameters(self): + """Test function with no parameters.""" + + def get_time() -> str: + """Get current time.""" + return '12:00' + + decl = _automatic_function_calling_util.build_function_declaration(get_time) + + assert decl.name == 'get_time' + assert decl.parameters_json_schema is None + + def test_docstring_preserved(self): + """Test that docstring is preserved as description.""" + + def well_documented(x: int) -> int: + """This is a well-documented function. + + It does something useful. + """ + return x + + decl = _automatic_function_calling_util.build_function_declaration( + well_documented + ) + + assert 'well-documented function' in decl.description + assert 'something useful' in decl.description + + def test_default_values(self): + """Test parameters with default values.""" + + def greet(name: str = 'World') -> str: + """Greet someone.""" + return f'Hello, {name}!' + + decl = _automatic_function_calling_util.build_function_declaration(greet) + + schema = decl.parameters_json_schema + assert schema['properties']['name']['default'] == 'World' + assert 'name' not in schema.get('required', [])