diff --git a/src/google/adk/tools/_function_tool_declarations.py b/src/google/adk/tools/_function_tool_declarations.py index d50b12efee..f5383cf66a 100644 --- a/src/google/adk/tools/_function_tool_declarations.py +++ b/src/google/adk/tools/_function_tool_declarations.py @@ -63,7 +63,7 @@ def _get_function_fields( # Get type hints with forward reference resolution try: - type_hints = get_type_hints(func) + type_hints = get_type_hints(func, include_extras=True) except TypeError: # Can happen with mock objects or complex annotations type_hints = {} @@ -160,7 +160,7 @@ def _build_response_json_schema( # Handle string annotations (forward references) if isinstance(return_annotation, str): try: - type_hints = get_type_hints(func) + type_hints = get_type_hints(func, include_extras=True) return_annotation = type_hints.get('return', return_annotation) except TypeError: pass diff --git a/src/google/adk/tools/function_tool.py b/src/google/adk/tools/function_tool.py index 382dcbcc1e..aa8c7b7a0f 100644 --- a/src/google/adk/tools/function_tool.py +++ b/src/google/adk/tools/function_tool.py @@ -20,6 +20,7 @@ import inspect import logging from types import UnionType +from typing import Annotated from typing import Any from typing import Awaitable from typing import Callable @@ -189,15 +190,20 @@ def _preprocess_args(self, args: dict[str, Any]) -> dict[str, Any]: for param_name, param in signature.parameters.items(): if param_name in args: target_type = type_hints.get(param_name, param.annotation) + if get_origin(target_type) is Annotated: + # Strip Annotated to get actual type + target_type = get_args(target_type)[0] if target_type != inspect.Parameter.empty: - # Handle Optional/Union types (e.g. Optional[PydanticModel], PydanticModel | None) origin = get_origin(target_type) if origin is Union or origin is UnionType: union_args = get_args(target_type) - # Find the non-None type in Optional[T] (which is Union[T, None]) + # Find the non-None type in Optional[T] (which is Union[T, None]). + # Handle Optional[Annotated(...)] non_none_types = [ - arg for arg in union_args if arg is not type(None) + get_args(arg)[0] if get_origin(arg) is Annotated else arg + for arg in union_args + if arg is not type(None) ] if len(non_none_types) == 1: target_type = non_none_types[0] diff --git a/tests/unittests/tools/test_build_function_declaration.py b/tests/unittests/tools/test_build_function_declaration.py index 599341c90b..1f5107718b 100644 --- a/tests/unittests/tools/test_build_function_declaration.py +++ b/tests/unittests/tools/test_build_function_declaration.py @@ -13,6 +13,7 @@ # limitations under the License. from enum import Enum +from typing import Annotated from typing import Any from google.adk.features import FeatureName @@ -25,6 +26,7 @@ # TODO: crewai requires python 3.10 as minimum # from crewai_tools import FileReadTool from pydantic import BaseModel +from pydantic import Field import pytest @@ -648,6 +650,22 @@ def __call__(self, a: int, b: int): assert function_decl.name == 'Calc' assert function_decl.response is not None + def test_annotated_field_metadata_preserved(self): + """Test Annotated[T, Field(...)] metadata reaches the schema.""" + + def legacy_annotated_function( + count: Annotated[int, Field(description='How many widgets', ge=1)], + ) -> str: + return str(count) + + function_decl = _automatic_function_calling_util.build_function_declaration( + func=legacy_annotated_function + ) + + count_schema = function_decl.parameters.properties['count'] + assert count_schema.description == 'How many widgets' + assert count_schema.minimum == 1 + class TestBuildFunctionDeclarationWithJsonSchema: """Tests for build_function_declaration when JSON_SCHEMA_FOR_FUNC_DECL is enabled.""" @@ -918,6 +936,22 @@ def greet(name: str = 'World') -> str: assert schema['properties']['name']['default'] == 'World' assert 'name' not in schema.get('required', []) + def test_annotated_field_metadata_preserved(self): + """Test Annotated[T, Field(...)] metadata reaches the schema.""" + + def json_schema_annotated_function( + count: Annotated[int, Field(description='How many widgets', ge=1)], + ) -> str: + return str(count) + + function_decl = _automatic_function_calling_util.build_function_declaration( + func=json_schema_annotated_function + ) + + count_schema = function_decl.parameters_json_schema['properties']['count'] + assert count_schema['description'] == 'How many widgets' + assert count_schema['minimum'] == 1 + class TestBuildFunctionDeclarationFromSchemaDict: """Tests for the declaration builders that take a JSON schema dict. diff --git a/tests/unittests/tools/test_function_tool_declarations.py b/tests/unittests/tools/test_function_tool_declarations.py index 1efa438f33..29e33496b7 100644 --- a/tests/unittests/tools/test_function_tool_declarations.py +++ b/tests/unittests/tools/test_function_tool_declarations.py @@ -23,6 +23,7 @@ from collections.abc import Sequence import dataclasses from enum import Enum +from typing import Annotated from typing import Any from typing import AsyncGenerator from typing import Generator @@ -639,6 +640,103 @@ def get_status() -> StandardReturnDataclass: self.assertIn("status", decl.response_json_schema["properties"]) +class TestAnnotatedMetadata(parameterized.TestCase): + """Tests that Annotated[T, Field(...)] metadata reaches the schema.""" + + def test_annotated_field_metadata_in_schema(self): + """Test descriptions, constraints and defaults attached via Annotated.""" + + def configure( + count: Annotated[ + int, Field(description="How many widgets", ge=1, le=10) + ], + zip_code: Annotated[str, Field(pattern=r"^\d{5}$")], + retries: Annotated[int, Field(description="Retry attempts")] = 3, + ) -> str: + return "ok" + + decl = build_function_declaration_with_json_schema(configure) + schema = decl.parameters_json_schema + + self.assertEqual( + schema["properties"], + { + "count": { + "description": "How many widgets", + "maximum": 10, + "minimum": 1, + "title": "Count", + "type": "integer", + }, + "zip_code": { + "pattern": r"^\d{5}$", + "title": "Zip Code", + "type": "string", + }, + "retries": { + "default": 3, + "description": "Retry attempts", + "title": "Retries", + "type": "integer", + }, + }, + ) + self.assertEqual(set(schema["required"]), {"count", "zip_code"}) + + def test_annotated_optional_model(self): + """Test Annotated[Optional[Model], Field(...)] keeps its description.""" + + def save( + address: Annotated[ + Optional[Address], Field(description="Where to ship") + ] = None, + ) -> str: + return "ok" + + decl = build_function_declaration_with_json_schema(save) + address_schema = decl.parameters_json_schema["properties"]["address"] + + self.assertEqual(address_schema["description"], "Where to ship") + self.assertIsNone(address_schema["default"]) + self.assertEqual( + address_schema["anyOf"], + [{"$ref": "#/$defs/Address"}, {"type": "null"}], + ) + + def test_annotated_nested_list_constraints(self): + """Test Annotated metadata on both a list and its item type.""" + + def tally( + scores: Annotated[ + list[Annotated[int, Field(ge=0)]], Field(description="Scores") + ], + ) -> int: + return sum(scores) + + decl = build_function_declaration_with_json_schema(tally) + scores_schema = decl.parameters_json_schema["properties"]["scores"] + + self.assertEqual(scores_schema["description"], "Scores") + self.assertEqual(scores_schema["type"], "array") + self.assertEqual(scores_schema["items"]["type"], "integer") + self.assertEqual(scores_schema["items"]["minimum"], 0) + + def test_annotated_return_type_metadata(self): + """Test Annotated metadata on the return type.""" + + def count_items( + items: list[str], + ) -> Annotated[int, Field(description="Item count")]: + return len(items) + + decl = build_function_declaration_with_json_schema(count_items) + + self.assertEqual( + decl.response_json_schema, + {"description": "Item count", "type": "integer"}, + ) + + class TestSpecialCases(parameterized.TestCase): """Tests for special cases and edge cases.""" diff --git a/tests/unittests/tools/test_function_tool_pydantic.py b/tests/unittests/tools/test_function_tool_pydantic.py index 02328e0452..960dbcfba8 100644 --- a/tests/unittests/tools/test_function_tool_pydantic.py +++ b/tests/unittests/tools/test_function_tool_pydantic.py @@ -14,6 +14,7 @@ # Pydantic model conversion tests +from typing import Annotated from typing import Optional from typing import Union from unittest.mock import MagicMock @@ -23,6 +24,7 @@ from google.adk.tools.function_tool import FunctionTool from google.adk.tools.tool_context import ToolContext import pydantic +from pydantic import Field import pytest @@ -521,3 +523,192 @@ def create_entity_profile( tool_context=tool_context_mock, ) assert company_result == {"entity_type": "company", "name": "Acme Corp"} + + +# Annotated parameters, with type hints resolving normally. + + +def test_preprocess_args_with_annotated_pydantic_model(): + """Test _preprocess_args converts a dict for Annotated[Model, Field].""" + + def fn(user: Annotated[UserModel, Field(description="A user")]): + return user.name + + processed_args = FunctionTool(fn)._preprocess_args( + {"user": {"name": "Alice", "age": 30}} + ) + + assert isinstance(processed_args["user"], UserModel) + assert processed_args["user"].name == "Alice" + + +def test_preprocess_args_with_annotated_optional_model(): + """Test _preprocess_args converts a dict for Annotated[Optional[Model], Field].""" + + def fn( + preferences: Annotated[ + Optional[PreferencesModel], Field(description="Prefs") + ] = None, + ): + return preferences + + processed_args = FunctionTool(fn)._preprocess_args( + {"preferences": {"theme": "dark"}} + ) + + assert isinstance(processed_args["preferences"], PreferencesModel) + assert processed_args["preferences"].theme == "dark" + + +def test_preprocess_args_with_optional_annotated_model(): + """Test _preprocess_args converts a dict for Optional[Annotated[Model, Field]]. + + Here the Annotated sits inside the union, so unwrapping the outer annotation + is not enough. + """ + + def fn( + user: Optional[Annotated[UserModel, Field(description="A user")]] = None, + ): + return user + + processed_args = FunctionTool(fn)._preprocess_args( + {"user": {"name": "Bob", "age": 25}} + ) + + assert isinstance(processed_args["user"], UserModel) + assert processed_args["user"].name == "Bob" + + +def test_preprocess_args_with_annotated_union_of_basemodels(): + """Test _preprocess_args picks the right member of a union of Annotated models.""" + + def fn( + entity: Union[ + Annotated[UserModel, Field(description="A user")], + Annotated[CompanyModel, Field(description="A company")], + ], + ): + return entity + + processed_args = FunctionTool(fn)._preprocess_args({ + "entity": { + "company_name": "Acme Corp", + "industry": "tech", + "employee_count": 50, + } + }) + + assert isinstance(processed_args["entity"], CompanyModel) + assert processed_args["entity"].company_name == "Acme Corp" + + +def test_preprocess_args_with_annotated_list_of_models(): + """Test _preprocess_args converts dicts for Annotated[list[Model], Field].""" + + def fn( + users: Annotated[list[UserModel], Field(description="Users")], + ): + return users + + processed_args = FunctionTool(fn)._preprocess_args( + {"users": [{"name": "Alice", "age": 30}, {"name": "Bob", "age": 25}]} + ) + + assert all(isinstance(user, UserModel) for user in processed_args["users"]) + assert processed_args["users"][1].name == "Bob" + + +def test_preprocess_args_with_annotated_primitive_unchanged(): + """Test _preprocess_args leaves an Annotated primitive alone.""" + + def fn(count: Annotated[int, Field(ge=1)]): + return count + + processed_args = FunctionTool(fn)._preprocess_args({"count": 7}) + + assert processed_args["count"] == 7 + + +# In each test below, the locally-scoped recursive alias can't be resolved from +# module globals, so get_type_hints() raises NameError and _preprocess_args +# falls back to the raw, still-Annotated param.annotation. + + +def test_preprocess_args_annotated_model_unresolvable_signature(): + """Test Annotated[Model, Field] converts on the param.annotation fallback.""" + Recursive = Union[int, str, list["Recursive"]] + + def fn( + user: Annotated[UserModel, Field(description="A user")], + data: Recursive = None, + ) -> dict: + return {"name": user.name, "type": type(user).__name__} + + processed_args = FunctionTool(fn)._preprocess_args( + {"user": {"name": "Alice", "age": 30}} + ) + + assert isinstance(processed_args["user"], UserModel) + assert processed_args["user"].name == "Alice" + + +def test_preprocess_args_optional_annotated_model_unresolvable_signature(): + """Test Optional[Annotated[Model, Field]] converts on the fallback too.""" + Recursive = Union[int, str, list["Recursive"]] + + def fn( + user: Optional[Annotated[UserModel, Field(description="A user")]] = None, + data: Recursive = None, + ): + return user + + processed_args = FunctionTool(fn)._preprocess_args( + {"user": {"name": "Bob", "age": 25}} + ) + + assert isinstance(processed_args["user"], UserModel) + assert processed_args["user"].name == "Bob" + + +def test_preprocess_args_bare_model_unresolvable_signature(): + """Control: a bare model on the same signature shape already converted.""" + Recursive = Union[int, str, list["Recursive"]] + + def fn( + user: UserModel, + data: Recursive = None, + ): + return user + + processed_args = FunctionTool(fn)._preprocess_args( + {"user": {"name": "Charlie", "age": 35}} + ) + + assert isinstance(processed_args["user"], UserModel) + + +async def test_run_async_with_annotated_model_unresolvable_signature(): + """run_async end-to-end passes a model instance, not a dict, to the function.""" + Recursive = Union[int, str, list["Recursive"]] + + def fn( + user: Annotated[UserModel, Field(description="A user")], + data: Recursive = None, + ) -> dict: + return {"name": user.name, "type": type(user).__name__} + + tool = FunctionTool(fn) + + tool_context_mock = MagicMock(spec=ToolContext) + invocation_context_mock = MagicMock(spec=InvocationContext) + session_mock = MagicMock(spec=Session) + invocation_context_mock.session = session_mock + tool_context_mock.invocation_context = invocation_context_mock + + result = await tool.run_async( + args={"user": {"name": "Diana", "age": 32}}, + tool_context=tool_context_mock, + ) + + assert result == {"name": "Diana", "type": "UserModel"}