Add support for non-string enums in Gemini (#1564)

Co-authored-by: Thomas Boser <thomasboser@gmail.com>
Co-authored-by: Douwe Maan <douwe@pydantic.dev>

Author: 3 people

pydantic_ai_slim/pydantic_ai/models/gemini.py

@@ -831,6 +831,12 @@ def transform(self, schema: JsonSchema) -> JsonSchema:
         schema.pop('exclusiveMaximum', None)
         schema.pop('exclusiveMinimum', None)
 
+        # Gemini only supports string enums, so we need to convert any enum values to strings.
+        # Pydantic will take care of transforming the transformed string values to the correct type.
+        if enum := schema.get('enum'):
+            schema['type'] = 'string'
+            schema['enum'] = [str(val) for val in enum]
+
         type_ = schema.get('type')
         if 'oneOf' in schema and 'type' not in schema:  # pragma: no cover
             # This gets hit when we have a discriminated union

tests/models/test_gemini.py

@@ -6,6 +6,7 @@
 from collections.abc import AsyncIterator, Callable, Sequence
 from dataclasses import dataclass
 from datetime import timezone
+from enum import IntEnum
 from typing import Annotated
 
 import httpx
@@ -228,6 +229,64 @@ class Locations(BaseModel):
     )
 
 
+async def test_json_def_enum(allow_model_requests: None):
+    class ProgressEnum(IntEnum):
+        DONE = 100
+        ALMOST_DONE = 80
+        IN_PROGRESS = 60
+        BARELY_STARTED = 40
+        NOT_STARTED = 20
+
+    class QueryDetails(BaseModel):
+        progress: list[ProgressEnum] | None = None
+
+    json_schema = QueryDetails.model_json_schema()
+    assert json_schema == snapshot(
+        {
+            '$defs': {'ProgressEnum': {'enum': [100, 80, 60, 40, 20], 'title': 'ProgressEnum', 'type': 'integer'}},
+            'properties': {
+                'progress': {
+                    'anyOf': [{'items': {'$ref': '#/$defs/ProgressEnum'}, 'type': 'array'}, {'type': 'null'}],
+                    'default': None,
+                    'title': 'Progress',
+                }
+            },
+            'title': 'QueryDetails',
+            'type': 'object',
+        }
+    )
+    m = GeminiModel('gemini-1.5-flash', provider=GoogleGLAProvider(api_key='via-arg'))
+    output_tool = ToolDefinition(
+        'result',
+        'This is the tool for the final Result',
+        json_schema,
+    )
+    mrp = ModelRequestParameters(function_tools=[], allow_text_output=True, output_tools=[output_tool])
+    mrp = m.customize_request_parameters(mrp)
+
+    # This tests that the enum values are properly converted to strings for Gemini
+    assert m._get_tools(mrp) == snapshot(
+        _GeminiTools(
+            function_declarations=[
+                _GeminiFunction(
+                    name='result',
+                    description='This is the tool for the final Result',
+                    parameters={
+                        'properties': {
+                            'progress': {
+                                'items': {'enum': ['100', '80', '60', '40', '20'], 'type': 'string'},
+                                'type': 'array',
+                                'nullable': True,
+                            }
+                        },
+                        'type': 'object',
+                    },
+                )
+            ]
+        )
+    )
+
+
 async def test_json_def_replaced_any_of(allow_model_requests: None):
     class Location(BaseModel):
         lat: float