Skip to content

[Gemini] IntEnum output not supported #1553

Closed
@zaidhaan

Description

@zaidhaan

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:

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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions