diff --git a/src/google/adk/tools/_function_tool_declarations.py b/src/google/adk/tools/_function_tool_declarations.py index ebb6af7de5..f766c8245b 100644 --- a/src/google/adk/tools/_function_tool_declarations.py +++ b/src/google/adk/tools/_function_tool_declarations.py @@ -40,6 +40,7 @@ from pydantic import create_model from pydantic import fields as pydantic_fields +from ..utils._callable_utils import unwrap_callable from ..utils.variant_utils import get_google_llm_variant from ..utils.variant_utils import GoogleLLMVariant @@ -66,7 +67,7 @@ def _get_function_fields( # Get type hints with forward reference resolution try: - type_hints = get_type_hints(func) + type_hints = get_type_hints(unwrap_callable(func)) except TypeError: # Can happen with mock objects or complex annotations type_hints = {} @@ -234,7 +235,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(unwrap_callable(func)) return_annotation = type_hints.get('return', return_annotation) except TypeError: pass @@ -264,7 +265,7 @@ def _build_response_json_schema( logging.debug( 'Failed to build schema with config, retrying without config for' ' %s: %s', - func.__name__, + get_callable_name(func), e, ) adapter = pydantic.TypeAdapter(return_annotation) @@ -272,7 +273,7 @@ def _build_response_json_schema( except Exception: logging.warning( 'Failed to build response JSON schema for %s', - func.__name__, + get_callable_name(func), exc_info=True, ) # Fall back to untyped response diff --git a/tests/unittests/tools/test_function_tool_with_import_annotations.py b/tests/unittests/tools/test_function_tool_with_import_annotations.py index c917bc2ca6..2fae009905 100644 --- a/tests/unittests/tools/test_function_tool_with_import_annotations.py +++ b/tests/unittests/tools/test_function_tool_with_import_annotations.py @@ -14,6 +14,7 @@ from __future__ import annotations +import functools from typing import Any from typing import Dict from typing import Optional @@ -23,6 +24,7 @@ from google.adk.utils.variant_utils import GoogleLLMVariant from google.genai import types import pydantic +import pytest def test_string_annotation_none_return_vertex(): @@ -207,6 +209,33 @@ class ItemModel(pydantic.BaseModel): quantity: int +@pytest.mark.parametrize('wrapper', ['function', 'partial', 'callable']) +def test_wrapped_tool_resolves_model_annotations(wrapper): + """Wrapped tools advertise the same resolved input and output model types.""" + + def lookup(prefix: str, item: ItemModel) -> ItemModel: + return item + + class Lookup: + + def __call__(self, item: ItemModel) -> ItemModel: + return item + + functions = { + 'function': lookup, + 'partial': functools.partial(lookup, 'catalog'), + 'callable': Lookup(), + } + declaration = _automatic_function_calling_util.build_function_declaration( + functions[wrapper], variant=GoogleLLMVariant.VERTEX_AI + ) + + schema = declaration.parameters_json_schema + assert schema['properties']['item']['$ref'] == '#/$defs/ItemModel' + assert schema['$defs']['ItemModel'] == ItemModel.model_json_schema() + assert declaration.response_json_schema == ItemModel.model_json_schema() + + def test_preprocess_args_with_list_of_pydantic_models_and_annotations(): """Test _preprocess_args converts dict to Pydantic model with string annotations."""