Description
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 string
s. In our case, we are supplying {"type": "integer", "enum": [20, ...] }
A possible workaround is to replace such enums with anyOf
s 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
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.