Skip to content

Commit e965fb6

Browse files
speedstorm1copybara-github
speedstorm1
authored and
copybara-github
committed
fix: Support numeric enum parameters in Schema.from_json_schema and AFC
fixes #2199 PiperOrigin-RevId: 1003376682
1 parent a216d95 commit e965fb6

7 files changed

Lines changed: 350 additions & 1 deletion

‎google/genai/_extra_utils.py‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717

1818
import asyncio
1919
import dataclasses
20+
import enum
2021
import inspect
2122
import io
2223
import logging
@@ -245,6 +246,34 @@ def is_annotation_pydantic_model(annotation: Any) -> bool:
245246
return False
246247

247248

249+
def _is_annotation_enum(annotation: Any) -> bool:
250+
try:
251+
return inspect.isclass(annotation) and issubclass(annotation, enum.Enum)
252+
except TypeError:
253+
return False
254+
255+
256+
def _convert_to_enum(
257+
value: Any, annotation: Type[enum.Enum], param_name: str, func_name: str
258+
) -> enum.Enum:
259+
"""Converts a function call argument value to a member of an Enum class."""
260+
try:
261+
return annotation(value)
262+
except ValueError:
263+
pass
264+
# Schema.enum holds the values as strings, so the model can return "5" for a
265+
# member whose value is 5.
266+
if isinstance(value, str):
267+
for member in annotation:
268+
if str(member.value) == value:
269+
return member
270+
raise errors.UnknownFunctionCallArgumentError(
271+
f'Failed to parse parameter {param_name} for function {func_name} from'
272+
f' function call part because function call argument value {value} is'
273+
f' not a valid value of enum {annotation}.'
274+
)
275+
276+
248277
def convert_if_exist_pydantic_model(
249278
value: Any, annotation: Any, param_name: str, func_name: str
250279
) -> Any:
@@ -270,10 +299,18 @@ def convert_if_exist_pydantic_model(
270299
k: convert_if_exist_pydantic_model(v, value_type, param_name, func_name)
271300
for k, v in value.items()
272301
}
302+
if _is_annotation_enum(annotation):
303+
return _convert_to_enum(value, annotation, param_name, func_name)
273304
# example 1: typing.Union[int, float]
274305
# example 2: int | float equivalent to UnionType[int, float]
275306
if get_origin(annotation) in (Union, UnionType):
276307
for arg in get_args(annotation):
308+
if _is_annotation_enum(arg):
309+
try:
310+
return _convert_to_enum(value, arg, param_name, func_name)
311+
# do not raise here because a later type in the union can match.
312+
except errors.UnknownFunctionCallArgumentError:
313+
continue
277314
if (
278315
(get_args(arg) and get_origin(arg) is list)
279316
or isinstance(value, arg)

‎google/genai/tests/afc/test_get_function_response_parts.py‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616

1717
"""Tests for get_function_response_parts."""
1818

19+
import enum
1920
import typing
2021
from typing import Any
2122
import pytest
@@ -160,6 +161,74 @@ def func_under_test(a: str) -> str:
160161
) == expected_part.model_dump_json(exclude_none=True)
161162

162163

164+
class _DaysEnum(enum.IntEnum):
165+
ONE = 1
166+
FIVE = 5
167+
168+
169+
def _days_enum_response() -> GenerateContentResponse:
170+
# Function call args arrive as JSON numbers, i.e. 5.0 for an INTEGER.
171+
return GenerateContentResponse(
172+
candidates=[
173+
Candidate(
174+
content=Content(
175+
parts=[
176+
Part(
177+
function_call=FunctionCall(
178+
name='func_under_test',
179+
args={'days': 5.0},
180+
)
181+
)
182+
]
183+
)
184+
)
185+
]
186+
)
187+
188+
189+
def test_int_enum_value():
190+
def func_under_test(days: _DaysEnum) -> str:
191+
assert isinstance(days, _DaysEnum)
192+
return days.name
193+
194+
function_map = {'func_under_test': func_under_test}
195+
expected_parts = [
196+
Part(
197+
function_response=FunctionResponse(
198+
name='func_under_test',
199+
response={'result': 'FIVE'},
200+
)
201+
)
202+
]
203+
actual_parts = get_function_response_parts(
204+
_days_enum_response(), function_map
205+
)
206+
207+
assert actual_parts == expected_parts
208+
209+
210+
@pytest.mark.asyncio
211+
async def test_int_enum_value_async():
212+
async def func_under_test(days: _DaysEnum) -> str:
213+
assert isinstance(days, _DaysEnum)
214+
return days.name
215+
216+
function_map = {'func_under_test': func_under_test}
217+
expected_parts = [
218+
Part(
219+
function_response=FunctionResponse(
220+
name='func_under_test',
221+
response={'result': 'FIVE'},
222+
)
223+
)
224+
]
225+
actual_parts = await get_function_response_parts_async(
226+
_days_enum_response(), function_map
227+
)
228+
229+
assert actual_parts == expected_parts
230+
231+
163232
@pytest.mark.asyncio
164233
async def test_mcp_tool():
165234
if not _is_mcp_imported:

‎google/genai/tests/afc/test_invoke_function_from_dict_args.py‎

Lines changed: 86 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,8 @@
1616

1717
"""Test invoke_function_from_dict_args."""
1818

19-
from typing import Union
19+
import enum
20+
from typing import Optional, Union
2021
import pydantic
2122
import pytest
2223
import sys
@@ -233,6 +234,90 @@ def func_under_test(x: int):
233234
invoke_function_from_dict_args({'x': {}}, func_under_test)
234235

235236

237+
class _DaysEnum(enum.IntEnum):
238+
ONE = 1
239+
FIVE = 5
240+
TEN = 10
241+
242+
243+
class _ColorEnum(enum.Enum):
244+
RED = 'red'
245+
BLUE = 'blue'
246+
247+
248+
def test_int_enum():
249+
def func_under_test(days: _DaysEnum):
250+
assert isinstance(days, _DaysEnum)
251+
return days.name
252+
253+
assert invoke_function_from_dict_args({'days': 5}, func_under_test) == 'FIVE'
254+
255+
256+
def test_int_enum_string_value():
257+
"""Schema.enum sends the values as strings, so "5" must also convert."""
258+
259+
def func_under_test(days: _DaysEnum):
260+
assert isinstance(days, _DaysEnum)
261+
return days.name
262+
263+
assert (
264+
invoke_function_from_dict_args({'days': '5'}, func_under_test) == 'FIVE'
265+
)
266+
267+
268+
def test_str_enum():
269+
def func_under_test(color: _ColorEnum):
270+
assert isinstance(color, _ColorEnum)
271+
return color.name
272+
273+
assert (
274+
invoke_function_from_dict_args({'color': 'blue'}, func_under_test)
275+
== 'BLUE'
276+
)
277+
278+
279+
def test_optional_int_enum():
280+
def func_under_test(days: Optional[_DaysEnum] = None):
281+
return days
282+
283+
assert (
284+
invoke_function_from_dict_args({'days': 10}, func_under_test)
285+
is _DaysEnum.TEN
286+
)
287+
assert invoke_function_from_dict_args({'days': None}, func_under_test) is None
288+
289+
290+
def test_union_enum_falls_through_to_later_type():
291+
def func_under_test(x: Union[_DaysEnum, str]):
292+
return x
293+
294+
assert invoke_function_from_dict_args({'x': 1}, func_under_test) is (
295+
_DaysEnum.ONE
296+
)
297+
assert invoke_function_from_dict_args({'x': 'abc'}, func_under_test) == 'abc'
298+
299+
300+
def test_list_of_int_enum():
301+
def func_under_test(days: list[_DaysEnum]):
302+
return days
303+
304+
assert invoke_function_from_dict_args(
305+
{'days': [1, '10']}, func_under_test
306+
) == [_DaysEnum.ONE, _DaysEnum.TEN]
307+
308+
309+
def test_invalid_enum_value():
310+
def func_under_test(days: _DaysEnum):
311+
return days
312+
313+
with pytest.raises(errors.UnknownFunctionCallArgumentError):
314+
invoke_function_from_dict_args({'days': 7}, func_under_test)
315+
with pytest.raises(errors.UnknownFunctionCallArgumentError):
316+
invoke_function_from_dict_args({'days': 'FIVE'}, func_under_test)
317+
with pytest.raises(errors.UnknownFunctionCallArgumentError):
318+
invoke_function_from_dict_args({'days': None}, func_under_test)
319+
320+
236321
def test_function_invocation_error():
237322
def func_under_test(x: int):
238323
return x / 0

‎google/genai/tests/mcp/test_mcp_to_gemini_tools.py‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -124,6 +124,44 @@ def test_items_conversion():
124124
]
125125

126126

127+
def test_integer_enum_conversion():
128+
"""Test that integer enum values in MCP tools become strings."""
129+
mcp_tools = [
130+
mcp_types.Tool(
131+
name='tool',
132+
description='tool-description',
133+
inputSchema={
134+
'type': 'object',
135+
'properties': {
136+
'days': {
137+
'type': 'integer',
138+
'enum': [1, 5, 10],
139+
},
140+
},
141+
},
142+
),
143+
]
144+
result = _mcp_utils.mcp_to_gemini_tools(mcp_tools)
145+
assert result == [
146+
types.Tool(
147+
function_declarations=[
148+
types.FunctionDeclaration(
149+
name='tool',
150+
description='tool-description',
151+
parameters=types.Schema(
152+
type='OBJECT',
153+
properties={
154+
'days': types.Schema(
155+
type='INTEGER', enum=['1', '5', '10']
156+
),
157+
},
158+
),
159+
),
160+
],
161+
),
162+
]
163+
164+
127165
def test_any_of_conversion():
128166
"""Test conversion of MCP tools with any_of to Gemini tools."""
129167
mcp_tools = [

‎google/genai/tests/types/test_schema_from_json_schema.py‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,7 +13,10 @@
1313
# limitations under the License.
1414
#
1515

16+
import enum
1617
import logging
18+
import warnings
19+
1720
import pydantic
1821

1922
from ... import types
@@ -275,6 +278,70 @@ def test_union_like_type_conversion_suite2():
275278
assert actual_vertex_ai_schema == expected_schema
276279

277280

281+
def test_integer_enum_conversion():
282+
"""Test that numeric JSONSchema.enum values become strings in Schema.enum."""
283+
json_schema = types.JSONSchema(type='integer', enum=[1, 5, 10])
284+
gemini_api_schema = types.Schema.from_json_schema(json_schema=json_schema)
285+
vertex_ai_schema = types.Schema.from_json_schema(
286+
json_schema=json_schema, api_option='VERTEX_AI'
287+
)
288+
expected_schema = types.Schema(type='INTEGER', enum=['1', '5', '10'])
289+
290+
assert gemini_api_schema == expected_schema
291+
assert vertex_ai_schema == expected_schema
292+
with warnings.catch_warnings():
293+
warnings.simplefilter('error')
294+
gemini_api_schema.model_dump(mode='json')
295+
296+
297+
def test_number_enum_conversion():
298+
"""Test that float and int enum values of a NUMBER schema become strings."""
299+
json_schema = types.JSONSchema(type='number', enum=[1.5, 2])
300+
schema = types.Schema.from_json_schema(json_schema=json_schema)
301+
302+
assert schema == types.Schema(type='NUMBER', enum=['1.5', '2'])
303+
304+
305+
def test_nullable_integer_enum_conversion():
306+
"""Test that a nullable integer enum keeps nullable and gets string values."""
307+
json_schema = types.JSONSchema(type=['integer', 'null'], enum=[1, 2])
308+
schema = types.Schema.from_json_schema(json_schema=json_schema)
309+
310+
assert schema == types.Schema(type='INTEGER', nullable=True, enum=['1', '2'])
311+
312+
313+
def test_union_like_type_integer_enum_conversion():
314+
"""Test that the enum copied into an any_of sub-schema becomes strings."""
315+
json_schema = types.JSONSchema(
316+
type=['integer', 'array'],
317+
enum=[1, 2],
318+
items=types.JSONSchema(type='integer'),
319+
)
320+
schema = types.Schema.from_json_schema(json_schema=json_schema)
321+
322+
assert schema == types.Schema(
323+
any_of=[
324+
types.Schema(type='INTEGER', enum=['1', '2']),
325+
types.Schema(type='ARRAY', items=types.Schema(type='INTEGER')),
326+
]
327+
)
328+
329+
330+
def test_int_enum_member_enum_conversion():
331+
"""Test that IntEnum members convert to their value, not their name."""
332+
333+
class DaysEnum(enum.IntEnum):
334+
ONE = 1
335+
FIVE = 5
336+
337+
json_schema = types.JSONSchema(
338+
type='integer', enum=[DaysEnum.ONE, DaysEnum.FIVE]
339+
)
340+
schema = types.Schema.from_json_schema(json_schema=json_schema)
341+
342+
assert schema == types.Schema(type='INTEGER', enum=['1', '5'])
343+
344+
278345
def test_array_type_conversion():
279346
"""Test conversion of JSONSchema.items to Schema.items"""
280347
json_schema = types.JSONSchema(

0 commit comments

Comments
 (0)