From 86a44873e9b2dfc7e62fa31a9ac3be57c0bbff7b Mon Sep 17 00:00:00 2001 From: "Xiang (Sean) Zhou" Date: Fri, 1 Aug 2025 10:20:43 -0700 Subject: [PATCH] fix: Annotate response type as None for transfer_to_agent tool and set empty Schema as response schema when tool has no response annotation 1. if a function has no return type annotation, we should treat it as returning any type 2. we use empty schema (with `type` as None) to indicate no type constraints and this is already supported by model server PiperOrigin-RevId: 789808104 --- .../tools/_automatic_function_calling_util.py | 23 ++- .../tools/_function_parameter_parse_util.py | 4 + .../adk/tools/transfer_to_agent_tool.py | 4 +- .../flows/llm_flows/test_agent_transfer.py | 8 +- .../tools/test_build_function_declaration.py | 14 +- .../tools/test_from_function_with_options.py | 26 ++- ...t_function_tool_with_import_annotations.py | 179 ++++++++++++++++++ 7 files changed, 242 insertions(+), 16 deletions(-) create mode 100644 tests/unittests/tools/test_function_tool_with_import_annotations.py diff --git a/src/google/adk/tools/_automatic_function_calling_util.py b/src/google/adk/tools/_automatic_function_calling_util.py index 3a26862e..5e32f68e 100644 --- a/src/google/adk/tools/_automatic_function_calling_util.py +++ b/src/google/adk/tools/_automatic_function_calling_util.py @@ -329,11 +329,28 @@ def from_function_with_options( return_annotation = inspect.signature(func).return_annotation - # Handle functions with no return annotation or that return None + # Handle functions with no return annotation + if return_annotation is inspect._empty: + # Functions with no return annotation can return any type + return_value = inspect.Parameter( + 'return_value', + inspect.Parameter.POSITIONAL_OR_KEYWORD, + annotation=typing.Any, + ) + declaration.response = ( + _function_parameter_parse_util._parse_schema_from_parameter( + variant, + return_value, + func.__name__, + ) + ) + return declaration + + # Handle functions that explicitly return None if ( - return_annotation is inspect._empty - or return_annotation is None + return_annotation is None or return_annotation is type(None) + or (isinstance(return_annotation, str) and return_annotation == 'None') ): # Create a response schema for None/null return return_value = inspect.Parameter( diff --git a/src/google/adk/tools/_function_parameter_parse_util.py b/src/google/adk/tools/_function_parameter_parse_util.py index ba1e3c9a..a0168fbe 100644 --- a/src/google/adk/tools/_function_parameter_parse_util.py +++ b/src/google/adk/tools/_function_parameter_parse_util.py @@ -38,6 +38,10 @@ _py_builtin_type_to_schema_type = { list: types.Type.ARRAY, dict: types.Type.OBJECT, None: types.Type.NULL, + # TODO requested google GenAI SDK to add a Type.ANY and do the mapping on + # their side, once new enum is added, replace the below one with + # Any: types.Type.ANY + Any: None, } logger = logging.getLogger('google_adk.' + __name__) diff --git a/src/google/adk/tools/transfer_to_agent_tool.py b/src/google/adk/tools/transfer_to_agent_tool.py index a16afca0..99ee234b 100644 --- a/src/google/adk/tools/transfer_to_agent_tool.py +++ b/src/google/adk/tools/transfer_to_agent_tool.py @@ -12,10 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. +from __future__ import annotations + from .tool_context import ToolContext -def transfer_to_agent(agent_name: str, tool_context: ToolContext): +def transfer_to_agent(agent_name: str, tool_context: ToolContext) -> None: """Transfer the question to another agent. This tool hands off control to another agent when it's more suitable to diff --git a/tests/unittests/flows/llm_flows/test_agent_transfer.py b/tests/unittests/flows/llm_flows/test_agent_transfer.py index 4cb48c84..5268d0ca 100644 --- a/tests/unittests/flows/llm_flows/test_agent_transfer.py +++ b/tests/unittests/flows/llm_flows/test_agent_transfer.py @@ -89,7 +89,7 @@ def test_auto_to_single(): ('sub_agent_1', 'response1'), ] - # root_agent should still be the current agent, becaues sub_agent_1 is single. + # root_agent should still be the current agent, because sub_agent_1 is single. assert testing_utils.simplify_events(runner.run('test2')) == [ ('root_agent', 'response2'), ] @@ -140,7 +140,7 @@ def test_auto_to_auto_to_single(): def test_auto_to_sequential(): response = [ transfer_call_part('sub_agent_1'), - # sub_agent_1 responds directly instead of transfering. + # sub_agent_1 responds directly instead of transferring. 'response1', 'response2', 'response3', @@ -189,7 +189,7 @@ def test_auto_to_sequential(): def test_auto_to_sequential_to_auto(): response = [ transfer_call_part('sub_agent_1'), - # sub_agent_1 responds directly instead of transfering. + # sub_agent_1 responds directly instead of transferring. 'response1', transfer_call_part('sub_agent_1_2_1'), 'response2', @@ -250,7 +250,7 @@ def test_auto_to_sequential_to_auto(): def test_auto_to_loop(): response = [ transfer_call_part('sub_agent_1'), - # sub_agent_1 responds directly instead of transfering. + # sub_agent_1 responds directly instead of transferring. 'response1', 'response2', 'response3', diff --git a/tests/unittests/tools/test_build_function_declaration.py b/tests/unittests/tools/test_build_function_declaration.py index 444fbd99..edf3c712 100644 --- a/tests/unittests/tools/test_build_function_declaration.py +++ b/tests/unittests/tools/test_build_function_declaration.py @@ -298,9 +298,10 @@ def test_function_no_return_annotation_vertex_ai(): assert function_decl.name == 'function_no_return' assert function_decl.parameters.type == 'OBJECT' assert function_decl.parameters.properties['param'].type == 'STRING' - # VERTEX_AI should have response schema for None return + # VERTEX_AI should have response schema for functions with no return annotation + # Changed: Now uses Any type instead of NULL for no return annotation assert function_decl.response is not None - assert function_decl.response.type == types.Type.NULL + assert function_decl.response.type is None # Any type maps to None in schema def test_function_explicit_none_return_vertex_ai(): @@ -359,8 +360,8 @@ def test_function_regular_return_type_vertex_ai(): assert function_decl.response.type == types.Type.STRING -def test_transfer_to_agent_like_function(): - """Test a function similar to transfer_to_agent that caused the original issue.""" +def test_fucntion_with_no_response_annotations(): + """Test a function that has no response annotations.""" def transfer_to_agent(agent_name: str, tool_context: ToolContext): """Transfer the question to another agent.""" @@ -376,6 +377,7 @@ def test_transfer_to_agent_like_function(): assert function_decl.parameters.type == 'OBJECT' assert function_decl.parameters.properties['agent_name'].type == 'STRING' assert 'tool_context' not in function_decl.parameters.properties - # This should now have a response schema for VERTEX_AI variant + # This function has no return annotation, so it gets Any type instead of NULL + # Changed: Now uses Any type instead of NULL for no return annotation assert function_decl.response is not None - assert function_decl.response.type == types.Type.NULL + assert function_decl.response.type is None # Any type maps to None in schema diff --git a/tests/unittests/tools/test_from_function_with_options.py b/tests/unittests/tools/test_from_function_with_options.py index 328eefab..3ae5e1f5 100644 --- a/tests/unittests/tools/test_from_function_with_options.py +++ b/tests/unittests/tools/test_from_function_with_options.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +from typing import Any from typing import Dict from google.adk.tools import _automatic_function_calling_util @@ -51,9 +52,10 @@ def test_from_function_with_options_no_return_annotation_vertex(): assert declaration.name == 'test_function' assert declaration.parameters.type == 'OBJECT' assert declaration.parameters.properties['param'].type == 'STRING' - # VERTEX_AI should have response schema for None return + # VERTEX_AI should have response schema for functions with no return annotation + # Changed: Now uses Any type instead of NULL for no return annotation assert declaration.response is not None - assert declaration.response.type == types.Type.NULL + assert declaration.response.type is None # Any type maps to None in schema def test_from_function_with_options_explicit_none_return_vertex(): @@ -150,6 +152,26 @@ def test_from_function_with_options_int_return_vertex(): assert declaration.response.type == types.Type.INTEGER +def test_from_function_with_options_any_annotation_vertex(): + """Test from_function_with_options with Any type annotation for VERTEX_AI.""" + + def test_function(param: Any) -> Any: + """A test function that uses Any type annotations.""" + return param + + declaration = _automatic_function_calling_util.from_function_with_options( + test_function, GoogleLLMVariant.VERTEX_AI + ) + + assert declaration.name == 'test_function' + assert declaration.parameters.type == 'OBJECT' + # Any type should map to None in schema (TYPE_UNSPECIFIED behavior) + assert declaration.parameters.properties['param'].type is None + # VERTEX_AI should have response schema for Any return + assert declaration.response is not None + assert declaration.response.type is None # Any type maps to None in schema + + def test_from_function_with_options_no_params(): """Test from_function_with_options with no parameters.""" diff --git a/tests/unittests/tools/test_function_tool_with_import_annotations.py b/tests/unittests/tools/test_function_tool_with_import_annotations.py new file mode 100644 index 00000000..99309a06 --- /dev/null +++ b/tests/unittests/tools/test_function_tool_with_import_annotations.py @@ -0,0 +1,179 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from typing import Any +from typing import Dict + +from google.adk.tools import _automatic_function_calling_util +from google.adk.utils.variant_utils import GoogleLLMVariant +from google.genai import types + + +def test_string_annotation_none_return_vertex(): + """Test function with string annotation 'None' return for VERTEX_AI.""" + + def test_function(_param: str) -> None: + """A test function that returns None with string annotation.""" + pass + + declaration = _automatic_function_calling_util.from_function_with_options( + test_function, GoogleLLMVariant.VERTEX_AI + ) + + assert declaration.name == 'test_function' + assert declaration.parameters.type == 'OBJECT' + assert declaration.parameters.properties['_param'].type == 'STRING' + # VERTEX_AI should have response schema for None return (stored as string) + assert declaration.response is not None + assert declaration.response.type == types.Type.NULL + + +def test_string_annotation_none_return_gemini(): + """Test function with string annotation 'None' return for GEMINI_API.""" + + def test_function(_param: str) -> None: + """A test function that returns None with string annotation.""" + pass + + declaration = _automatic_function_calling_util.from_function_with_options( + test_function, GoogleLLMVariant.GEMINI_API + ) + + assert declaration.name == 'test_function' + assert declaration.parameters.type == 'OBJECT' + assert declaration.parameters.properties['_param'].type == 'STRING' + # GEMINI_API should not have response schema + assert declaration.response is None + + +def test_string_annotation_str_return_vertex(): + """Test function with string annotation 'str' return for VERTEX_AI.""" + + def test_function(_param: str) -> str: + """A test function that returns a string with string annotation.""" + return _param + + declaration = _automatic_function_calling_util.from_function_with_options( + test_function, GoogleLLMVariant.VERTEX_AI + ) + + assert declaration.name == 'test_function' + assert declaration.parameters.type == 'OBJECT' + assert declaration.parameters.properties['_param'].type == 'STRING' + # VERTEX_AI should have response schema for string return (stored as string) + assert declaration.response is not None + assert declaration.response.type == types.Type.STRING + + +def test_string_annotation_int_return_vertex(): + """Test function with string annotation 'int' return for VERTEX_AI.""" + + def test_function(_param: str) -> int: + """A test function that returns an int with string annotation.""" + return 42 + + declaration = _automatic_function_calling_util.from_function_with_options( + test_function, GoogleLLMVariant.VERTEX_AI + ) + + assert declaration.name == 'test_function' + assert declaration.parameters.type == 'OBJECT' + assert declaration.parameters.properties['_param'].type == 'STRING' + # VERTEX_AI should have response schema for int return (stored as string) + assert declaration.response is not None + assert declaration.response.type == types.Type.INTEGER + + +def test_string_annotation_dict_return_vertex(): + """Test function with string annotation Dict return for VERTEX_AI.""" + + def test_function(_param: str) -> Dict[str, str]: + """A test function that returns a dict with string annotation.""" + return {'result': _param} + + declaration = _automatic_function_calling_util.from_function_with_options( + test_function, GoogleLLMVariant.VERTEX_AI + ) + + assert declaration.name == 'test_function' + assert declaration.parameters.type == 'OBJECT' + assert declaration.parameters.properties['_param'].type == 'STRING' + # VERTEX_AI should have response schema for dict return (stored as string) + assert declaration.response is not None + assert declaration.response.type == types.Type.OBJECT + + +def test_string_annotation_any_return_vertex(): + """Test function with string annotation 'Any' return for VERTEX_AI.""" + + def test_function(_param: Any) -> Any: + """A test function that uses Any type with string annotations.""" + return _param + + declaration = _automatic_function_calling_util.from_function_with_options( + test_function, GoogleLLMVariant.VERTEX_AI + ) + + assert declaration.name == 'test_function' + assert declaration.parameters.type == 'OBJECT' + # Any type should map to None in schema (TYPE_UNSPECIFIED behavior) + assert declaration.parameters.properties['_param'].type is None + # VERTEX_AI should have response schema for Any return (stored as string) + assert declaration.response is not None + assert declaration.response.type is None # Any type maps to None in schema + + +def test_string_annotation_mixed_parameters_vertex(): + """Test function with mixed string annotations for parameters.""" + + def test_function(str_param: str, int_param: int, any_param: Any) -> str: + """A test function with mixed parameter types as string annotations.""" + return f'{str_param}-{int_param}-{any_param}' + + declaration = _automatic_function_calling_util.from_function_with_options( + test_function, GoogleLLMVariant.VERTEX_AI + ) + + assert declaration.name == 'test_function' + assert declaration.parameters.type == 'OBJECT' + assert declaration.parameters.properties['str_param'].type == 'STRING' + assert declaration.parameters.properties['int_param'].type == 'INTEGER' + assert declaration.parameters.properties['any_param'].type is None # Any type + # VERTEX_AI should have response schema for string return (stored as string) + assert declaration.response is not None + assert declaration.response.type == types.Type.STRING + + +def test_string_annotation_no_params_vertex(): + """Test function with no parameters but string annotation return.""" + + def test_function() -> str: + """A test function with no parameters that returns string (string annotation).""" + return 'hello' + + declaration = _automatic_function_calling_util.from_function_with_options( + test_function, GoogleLLMVariant.VERTEX_AI + ) + + assert declaration.name == 'test_function' + # No parameters should result in no parameters field or empty parameters + assert ( + declaration.parameters is None + or len(declaration.parameters.properties) == 0 + ) + # VERTEX_AI should have response schema for string return (stored as string) + assert declaration.response is not None + assert declaration.response.type == types.Type.STRING