diff --git a/src/google/adk/tools/_function_parameter_parse_util.py b/src/google/adk/tools/_function_parameter_parse_util.py index a0168fbe219..9e91bb59caf 100644 --- a/src/google/adk/tools/_function_parameter_parse_util.py +++ b/src/google/adk/tools/_function_parameter_parse_util.py @@ -15,6 +15,7 @@ from __future__ import annotations +from enum import Enum import inspect import logging import types as typing_types @@ -145,6 +146,17 @@ def _parse_schema_from_parameter( schema.type = _py_builtin_type_to_schema_type[param.annotation] _raise_if_schema_unsupported(variant, schema) return schema + + if isinstance(param.annotation, type) and issubclass(param.annotation, Enum): + schema.type = types.Type.STRING + schema.enum = [e.value for e in param.annotation] + if param.default is not inspect.Parameter.empty: + if not _is_default_value_compatible(param.default, param.annotation): + raise ValueError(default_value_error_msg) + schema.default = param.default + _raise_if_schema_unsupported(variant, schema) + return schema + if ( get_origin(param.annotation) is Union # only parse simple UnionType, example int | str | float | bool diff --git a/tests/unittests/tools/test_build_function_declaration.py b/tests/unittests/tools/test_build_function_declaration.py index edf3c7128eb..67e80dd21ae 100644 --- a/tests/unittests/tools/test_build_function_declaration.py +++ b/tests/unittests/tools/test_build_function_declaration.py @@ -102,6 +102,26 @@ def simple_function(input_str: Dict[str, str]) -> str: assert function_decl.parameters.properties['input_str'].type == 'OBJECT' +def test_enum_input(): + from enum import Enum + + class Color(str, Enum): + RED = 'red' + GREEN = 'green' + BLUE = 'blue' + + def simple_function(input_str: Color) -> str: + return {'result': input_str} + + function_decl = _automatic_function_calling_util.build_function_declaration( + func=simple_function + ) + + assert function_decl.name == 'simple_function' + assert function_decl.parameters.type == 'OBJECT' + assert function_decl.parameters.properties['input_str'].type == 'STRING' + + def test_basemodel_input(): class CustomInput(BaseModel): input_str: str