Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 4 additions & 1 deletion src/google/adk/tools/_function_parameter_parse_util.py
Original file line number Diff line number Diff line change
Expand Up @@ -304,14 +304,17 @@ def _parse_schema_from_parameter(
_raise_if_schema_unsupported(variant, schema)
return schema
if isinstance(param.annotation, type) and issubclass(param.annotation, Enum):
# `schema.type` is always STRING here, so every enum value must be a
# string too (e.g. IntEnum members have int `.value`s otherwise).
schema.type = types.Type.STRING
schema.enum = [e.value for e in param.annotation]
schema.enum = [str(e.value) for e in param.annotation]
if param.default is not inspect.Parameter.empty:
default_value = (
param.default.value
if isinstance(param.default, Enum)
else param.default
)
default_value = str(default_value)
if default_value not in schema.enum:
raise ValueError(default_value_error_msg)
schema.default = default_value
Expand Down
22 changes: 22 additions & 0 deletions tests/unittests/tools/test_build_function_declaration.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
# limitations under the License.

from enum import Enum
from enum import IntEnum
from typing import Any

from google.adk.features import FeatureName
Expand Down Expand Up @@ -439,6 +440,27 @@ def simple_function_with_wrong_enum(input: InputEnum = 'WRONG_ENUM'):
func=simple_function_with_wrong_enum
)

def test_int_enum(self):

class Level(IntEnum):
LOW = 1
HIGH = 2

def set_level(level: Level = Level.LOW):
return level.value

function_decl = _automatic_function_calling_util.build_function_declaration(
func=set_level
)

level_schema = function_decl.parameters.properties['level']
assert level_schema.type == 'STRING'
assert level_schema.enum == ['1', '2']
assert level_schema.default == '1'
# The declaration must be valid for types.Schema's own contract: an
# `enum` field on a STRING schema must contain only strings.
types.Schema(type=level_schema.type, enum=level_schema.enum)

def test_basemodel_list(self):
class ChildInput(BaseModel):
input_str: str
Expand Down