Skip to content

Commit ab0730c

Browse files
committed
Review feedback re: openai
1 parent 1d8d19b commit ab0730c

File tree

4 files changed

+74
-5
lines changed

4 files changed

+74
-5
lines changed

docs/builtin-tools.md

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -291,6 +291,10 @@ assert isinstance(result.output, BinaryImage)
291291

292292
_(This example is complete, it can be run "as is")_
293293

294+
OpenAI Responses models also respect the `aspect_ratio` parameter. Because the OpenAI API only exposes discrete image sizes,
295+
PydanticAI maps `'1:1'` -> `1024x1024`, `'2:3'` -> `1024x1536`, and `'3:2'` -> `1536x1024`. Providing any other aspect ratio
296+
results in an error, and if you also set `size` it must match the computed value.
297+
294298
To control the aspect ratio when using Gemini image models, include the `ImageGenerationTool` explicitly:
295299

296300
```py {title="image_generation_google_aspect_ratio.py"}
@@ -322,7 +326,7 @@ For more details, check the [API documentation][pydantic_ai.builtin_tools.ImageG
322326
| `partial_images` |||
323327
| `quality` |||
324328
| `size` |||
325-
| `aspect_ratio` | ||
329+
| `aspect_ratio` | ✅ (1:1, 2:3, 3:2) ||
326330

327331
## URL Context Tool
328332

pydantic_ai_slim/pydantic_ai/builtin_tools.py

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -264,6 +264,7 @@ class ImageGenerationTool(AbstractBuiltinTool):
264264
Supported by:
265265
266266
* Google image-generation models (Gemini)
267+
* OpenAI Responses (maps '1:1', '2:3', and '3:2' to supported sizes)
267268
"""
268269

269270
kind: str = 'image_generation'

pydantic_ai_slim/pydantic_ai/models/openai.py

Lines changed: 33 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,7 @@
1818
from .._run_context import RunContext
1919
from .._thinking_part import split_content_into_text_and_thinking
2020
from .._utils import guard_tool_call_id as _guard_tool_call_id, now_utc as _now_utc, number_to_datetime
21-
from ..builtin_tools import CodeExecutionTool, ImageGenerationTool, MCPServerTool, WebSearchTool
21+
from ..builtin_tools import CodeExecutionTool, ImageAspectRatio, ImageGenerationTool, MCPServerTool, WebSearchTool
2222
from ..exceptions import UserError
2323
from ..messages import (
2424
AudioUrl,
@@ -134,6 +134,36 @@
134134
'failed': 'error',
135135
}
136136

137+
_OPENAI_ASPECT_RATIO_TO_SIZE: dict[ImageAspectRatio, Literal['1024x1024', '1024x1536', '1536x1024']] = {
138+
'1:1': '1024x1024',
139+
'2:3': '1024x1536',
140+
'3:2': '1536x1024',
141+
}
142+
143+
144+
def _resolve_openai_image_generation_size(
145+
tool: ImageGenerationTool,
146+
) -> Literal['auto', '1024x1024', '1024x1536', '1536x1024']:
147+
"""Map `ImageGenerationTool.aspect_ratio` to an OpenAI size string when provided."""
148+
aspect_ratio = tool.aspect_ratio
149+
if aspect_ratio is None:
150+
return tool.size
151+
152+
mapped_size = _OPENAI_ASPECT_RATIO_TO_SIZE.get(aspect_ratio)
153+
if mapped_size is None:
154+
supported = ', '.join(_OPENAI_ASPECT_RATIO_TO_SIZE)
155+
raise UserError(
156+
f'OpenAI image generation only supports `aspect_ratio` values: {supported}. '
157+
'Specify one of those values or omit `aspect_ratio`.'
158+
)
159+
160+
if tool.size not in ('auto', mapped_size):
161+
raise UserError(
162+
'`ImageGenerationTool` cannot combine `aspect_ratio` with a conflicting `size` when using OpenAI.'
163+
)
164+
165+
return mapped_size
166+
137167

138168
class OpenAIChatModelSettings(ModelSettings, total=False):
139169
"""Settings used for an OpenAI model request."""
@@ -1298,6 +1328,7 @@ def _get_builtin_tools(self, model_request_parameters: ModelRequestParameters) -
12981328
tools.append(mcp_tool)
12991329
elif isinstance(tool, ImageGenerationTool): # pragma: no branch
13001330
has_image_generating_tool = True
1331+
size = _resolve_openai_image_generation_size(tool)
13011332
tools.append(
13021333
responses.tool_param.ImageGeneration(
13031334
type='image_generation',
@@ -1308,7 +1339,7 @@ def _get_builtin_tools(self, model_request_parameters: ModelRequestParameters) -
13081339
output_format=tool.output_format or 'png',
13091340
partial_images=tool.partial_images,
13101341
quality=tool.quality,
1311-
size=tool.size,
1342+
size=size,
13121343
)
13131344
)
13141345
else:

tests/models/test_openai_responses.py

Lines changed: 35 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import json
22
import re
33
from dataclasses import replace
4-
from typing import Any, cast
4+
from typing import Any, Literal, cast
55

66
import pytest
77
from inline_snapshot import snapshot
@@ -32,17 +32,19 @@
3232
ToolCallPartDelta,
3333
ToolReturnPart,
3434
UnexpectedModelBehavior,
35+
UserError,
3536
UserPromptPart,
3637
capture_run_messages,
3738
)
3839
from pydantic_ai.agent import Agent
39-
from pydantic_ai.builtin_tools import CodeExecutionTool, MCPServerTool, WebSearchTool
40+
from pydantic_ai.builtin_tools import CodeExecutionTool, ImageAspectRatio, MCPServerTool, WebSearchTool
4041
from pydantic_ai.exceptions import ModelHTTPError, ModelRetry
4142
from pydantic_ai.messages import (
4243
BuiltinToolCallEvent, # pyright: ignore[reportDeprecated]
4344
BuiltinToolResultEvent, # pyright: ignore[reportDeprecated]
4445
)
4546
from pydantic_ai.models import ModelRequestParameters
47+
from pydantic_ai.models.openai import _resolve_openai_image_generation_size # pyright: ignore[reportPrivateUsage]
4648
from pydantic_ai.output import NativeOutput, PromptedOutput, TextOutput, ToolOutput
4749
from pydantic_ai.profiles.openai import openai_model_profile
4850
from pydantic_ai.tools import ToolDefinition
@@ -124,6 +126,37 @@ async def test_openai_responses_image_detail_vendor_metadata(allow_model_request
124126
assert all(part['detail'] == 'high' for part in image_parts)
125127

126128

129+
@pytest.mark.parametrize(
130+
('aspect_ratio', 'explicit_size', 'expected_size'),
131+
[
132+
('1:1', 'auto', '1024x1024'),
133+
('2:3', '1024x1536', '1024x1536'),
134+
('3:2', 'auto', '1536x1024'),
135+
],
136+
)
137+
def test_openai_responses_image_generation_tool_aspect_ratio_mapping(
138+
aspect_ratio: ImageAspectRatio,
139+
explicit_size: Literal['1024x1024', '1024x1536', '1536x1024', 'auto'],
140+
expected_size: Literal['1024x1024', '1024x1536', '1536x1024'],
141+
) -> None:
142+
tool = ImageGenerationTool(aspect_ratio=aspect_ratio, size=explicit_size)
143+
assert _resolve_openai_image_generation_size(tool) == expected_size
144+
145+
146+
def test_openai_responses_image_generation_tool_aspect_ratio_invalid() -> None:
147+
tool = ImageGenerationTool(aspect_ratio='16:9')
148+
149+
with pytest.raises(UserError, match='OpenAI image generation only supports `aspect_ratio` values'):
150+
_resolve_openai_image_generation_size(tool)
151+
152+
153+
def test_openai_responses_image_generation_tool_aspect_ratio_conflicts_with_size() -> None:
154+
tool = ImageGenerationTool(aspect_ratio='1:1', size='1536x1024')
155+
156+
with pytest.raises(UserError, match='cannot combine `aspect_ratio` with a conflicting `size`'):
157+
_resolve_openai_image_generation_size(tool)
158+
159+
127160
async def test_openai_responses_model_simple_response_with_tool_call(allow_model_requests: None, openai_api_key: str):
128161
model = OpenAIResponsesModel('gpt-4o', provider=OpenAIProvider(api_key=openai_api_key))
129162

0 commit comments

Comments
 (0)