[Gemini] IntEnum output not supported #1553

zaidhaan · 2025-04-19T23:11:53Z

I'm not sure if this should qualify as a bug or a feature request, so I'm not marking it as either for now.

pydantic-ai version: 0.1.3.
Python version: 3.12.9.

TL;DR:

from enum import IntEnum
from typing import Annotated, Optional
from pydantic_ai import Agent
from pydantic import BaseModel, Field

class ProgressEnum(IntEnum):
    DONE = 100
    ALMOST_DONE = 80
    IN_PROGRESS = 60
    BARELY_STARTED = 40
    NOT_STARTED = 20

class QueryDetails(BaseModel):
    progress: Optional[
        Annotated[list[ProgressEnum], Field(description="The progress of an item.")]
    ] = None

agent = Agent(
    'google-gla:gemini-2.0-flash',
    output_type=QueryDetails,
    system_prompt='You are a useful assistant that helps translate the user\'s request into a search query',
)

if __name__ == "__main__":
    result = agent.run_sync("Items with more than 50% progress")
    print(result.output)

Current result:

Traceback (most recent call last):
  File "/Users/me/test/packages/ai/ai/test.py", line 25, in <module>
    result = agent.run_sync("Items with more than 50% progress")
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/me/test/.venv/lib/python3.12/site-packages/pydantic_ai/agent.py", line 766, in run_sync
    return get_event_loop().run_until_complete(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Cellar/python@3.12/3.12.9/Frameworks/Python.framework/Versions/3.12/lib/python3.12/asyncio/base_events.py", line 691, in run_until_complete
    return future.result()
           ^^^^^^^^^^^^^^^
  File "/Users/me/test/.venv/lib/python3.12/site-packages/pydantic_ai/agent.py", line 436, in run
    async for _ in agent_run:
  File "/Users/me/test/.venv/lib/python3.12/site-packages/pydantic_ai/agent.py", line 1745, in __anext__
    next_node = await self._graph_run.__anext__()
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/me/test/.venv/lib/python3.12/site-packages/pydantic_graph/graph.py", line 800, in __anext__
    return await self.next(self._next_node)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/me/test/.venv/lib/python3.12/site-packages/pydantic_graph/graph.py", line 773, in next
    self._next_node = await node.run(ctx)
                      ^^^^^^^^^^^^^^^^^^^
  File "/Users/me/test/.venv/lib/python3.12/site-packages/pydantic_ai/_agent_graph.py", line 271, in run
    return await self._make_request(ctx)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/me/test/.venv/lib/python3.12/site-packages/pydantic_ai/_agent_graph.py", line 325, in _make_request
    model_response, request_usage = await ctx.deps.model.request(
                                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/me/test/.venv/lib/python3.12/site-packages/pydantic_ai/models/gemini.py", line 138, in request
    async with self._make_request(
               ^^^^^^^^^^^^^^^^^^^
  File "/opt/homebrew/Cellar/python@3.12/3.12.9/Frameworks/Python.framework/Versions/3.12/lib/python3.12/contextlib.py", line 210, in __aenter__
    return await anext(self.gen)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/Users/me/test/.venv/lib/python3.12/site-packages/pydantic_ai/models/gemini.py", line 245, in _make_request
    raise ModelHTTPError(status_code=status_code, model_name=self.model_name, body=r.text)
pydantic_ai.exceptions.ModelHTTPError: status_code: 400, model_name: gemini-2.0-flash, body: {
  "error": {
    "code": 400,
    "message": "Invalid value at 'tools.function_declarations[0].parameters.properties[0].value.items.enum[0]' (TYPE_STRING), 100\nInvalid value at 'tools.function_declarations[0].parameters.properties[0].value.items.enum[1]' (TYPE_STRING), 80\nInvalid value at 'tools.function_declarations[0].parameters.properties[0].value.items.enum[2]' (TYPE_STRING), 60\nInvalid value at 'tools.function_declarations[0].parameters.properties[0].value.items.enum[3]' (TYPE_STRING), 40\nInvalid value at 'tools.function_declarations[0].parameters.properties[0].value.items.enum[4]' (TYPE_STRING), 20",
    "status": "INVALID_ARGUMENT",
    "details": [
      {
        "@type": "type.googleapis.com/google.rpc.BadRequest",
        "fieldViolations": [
          {
            "field": "tools.function_declarations[0].parameters.properties[0].value.items.enum[0]",
            "description": "Invalid value at 'tools.function_declarations[0].parameters.properties[0].value.items.enum[0]' (TYPE_STRING), 100"
          },
          {
            "field": "tools.function_declarations[0].parameters.properties[0].value.items.enum[1]",
            "description": "Invalid value at 'tools.function_declarations[0].parameters.properties[0].value.items.enum[1]' (TYPE_STRING), 80"
          },
          {
            "field": "tools.function_declarations[0].parameters.properties[0].value.items.enum[2]",
            "description": "Invalid value at 'tools.function_declarations[0].parameters.properties[0].value.items.enum[2]' (TYPE_STRING), 60"
          },
          {
            "field": "tools.function_declarations[0].parameters.properties[0].value.items.enum[3]",
            "description": "Invalid value at 'tools.function_declarations[0].parameters.properties[0].value.items.enum[3]' (TYPE_STRING), 40"
          },
          {
            "field": "tools.function_declarations[0].parameters.properties[0].value.items.enum[4]",
            "description": "Invalid value at 'tools.function_declarations[0].parameters.properties[0].value.items.enum[4]' (TYPE_STRING), 20"
          }
        ]
      }
    ]
  }
}

Expected result:

progress=[<ProgressEnum.IN_PROGRESS: 60>, <ProgressEnum.ALMOST_DONE: 80>, <ProgressEnum.DONE: 100>]

Analysis
The Gemini API Schema expects that enum is an optional array of strings. In our case, we are supplying {"type": "integer", "enum": [20, ...] }

A possible workaround is to replace such enums with anyOfs consisting of the the number type, with its value fixed using minimum and maximum.

This can be done below here:

pydantic-ai/pydantic_ai_slim/pydantic_ai/models/gemini.py

Lines 818 to 824 in 7487ab4

    
           if type_ == 'string' and (fmt := schema.pop('format', None)): 
        
               description = schema.get('description') 
        
               if description: 
        
                   schema['description'] = f'{description} (format: {fmt})' 
        
               else: 
        
                   schema['description'] = f'Format: {fmt}'

with this:

if type_ in ('integer', 'number'):
    if enum := schema.get('enum'):
	schema.pop('enum')
	schema.pop('type')
	schema['anyOf'] = [
	    {'type': type_, 'minimum': val, 'maximum': val} for val in enum
	]

Admittedly, it's a workaround that's a bit hacky, though that is mostly due to limitations of the Gemini function calling schema (I'm not sure how other models compare with regards to this specific issue).

As far as I can tell, there's a few ways to go about this:

Throw an error when an Enum consisting of anything other than a string is received. (i.e. WONTFIX)
- In other words, users will need to implement a workaround (e.g. creating a string enum that wraps around the IntEnum). This is a bit cumbersome if the user's application code is written for the IntEnum, but having a minor translation method for such cases isn't a big deal I suppose.
Do the enum -> anyOf conversion proposed above. This is a working solution, though inelegant. I'm happy to translate my patch above into a PR if deemed the only viable solution.

The text was updated successfully, but these errors were encountered:

DouweM · 2025-04-21T19:32:49Z

@zaidhaan Transforming the schema to be Gemini-compatible in _GeminiJsonSchema makes sense, but there may be a better way. Could you try just changing the enum type and possible values to strings instead? We'd rely on Pydantic turning the returned values into (enum)ints again, so the user code (your code) should not need any changes.

zaidhaan · 2025-04-21T19:48:16Z

@DouweM well damn, you're right!

With this patch, everything works fine:

        if enum := schema.get('enum'):
            schema['type'] = 'string'
            schema['enum'] = [str(val) for val in enum]

Let me make a PR and this can be further reviewed there.

zaidhaan linked a pull request Apr 21, 2025 that will close this issue

Add support for non-string enums in Gemini #1564

Draft

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Gemini] IntEnum output not supported #1553

[Gemini] IntEnum output not supported #1553

zaidhaan commented Apr 19, 2025

DouweM commented Apr 21, 2025

zaidhaan commented Apr 21, 2025

[Gemini] IntEnum output not supported #1553

[Gemini] IntEnum output not supported #1553

Comments

zaidhaan commented Apr 19, 2025

DouweM commented Apr 21, 2025

zaidhaan commented Apr 21, 2025