feat: Support IntEnums when processing JSON schemas

PiperOrigin-RevId: 774980819

Author: Annhiluc

google/genai/_transformers.py

@@ -787,15 +787,25 @@ def _recurse(sub_schema: dict[str, Any]) -> dict[str, Any]:
 def _process_enum(
     enum: EnumMeta, client: _api_client.BaseApiClient
 ) -> types.Schema:
+  is_integer_enum = False
+
   for member in enum:  # type: ignore
-    if not isinstance(member.value, str):
+    if isinstance(member.value, int):
+      is_integer_enum = True
+    elif not isinstance(member.value, str):
       raise TypeError(
-          f'Enum member {member.name} value must be a string, got'
+          f'Enum member {member.name} value must be a string or integer, got'
           f' {type(member.value)}'
       )
 
+  enum_to_process = enum
+  if is_integer_enum:
+    str_members = [str(member.value) for member in enum]  # type: ignore
+    str_enum = Enum(enum.__name__, str_members, type=str)  # type: ignore
+    enum_to_process = str_enum
+
   class Placeholder(pydantic.BaseModel):
-    placeholder: enum  # type: ignore[valid-type]
+    placeholder: enum_to_process  # type: ignore[valid-type]
 
   enum_schema = Placeholder.model_json_schema()
   process_schema(enum_schema, client)

google/genai/tests/models/test_generate_content.py

@@ -1508,17 +1508,18 @@ class IntegerEnum(Enum):
     BRASS = 4
     KEYBOARD = 5
 
-  with pytest.raises(TypeError) as e:
-    client.models.generate_content(
-        model='gemini-1.5-flash',
-        contents='What instrument plays multiple notes at once?',
-        config={
-            'response_mime_type': 'text/x.enum',
-            'response_schema': IntegerEnum,
-        },
-    )
+  response =client.models.generate_content(
+      model='gemini-1.5-flash',
+      contents='What instrument plays multiple notes at once?',
+      config={
+          'response_mime_type': 'text/x.enum',
+          'response_schema': IntegerEnum,
+      },
+  )
+
+  instrument_values = {str(member.value) for member in IntegerEnum}
 
-    assert 'value must be a string' in str(e)
+  assert response.text in instrument_values
 
 
 def test_json_schema(client):