pydantic-ai-slim 1.0.13__py3-none-any.whl → 1.0.15__py3-none-any.whl

pydantic_ai/models/openai.py

@@ -17,14 +17,16 @@ from .._output import DEFAULT_OUTPUT_TOOL_NAME, OutputObjectDefinition
 from .._run_context import RunContext
 from .._thinking_part import split_content_into_text_and_thinking
 from .._utils import guard_tool_call_id as _guard_tool_call_id, now_utc as _now_utc, number_to_datetime
-from ..builtin_tools import CodeExecutionTool, WebSearchTool
+from ..builtin_tools import CodeExecutionTool, ImageGenerationTool, WebSearchTool
 from ..exceptions import UserError
 from ..messages import (
     AudioUrl,
     BinaryContent,
+    BinaryImage,
     BuiltinToolCallPart,
     BuiltinToolReturnPart,
     DocumentUrl,
+    FilePart,
     FinishReason,
     ImageUrl,
     ModelMessage,
@@ -678,6 +680,9 @@ class OpenAIChatModel(Model):
                     # OpenAI doesn't return built-in tool calls
                     elif isinstance(item, BuiltinToolCallPart | BuiltinToolReturnPart):  # pragma: no cover
                         pass
+                    elif isinstance(item, FilePart):  # pragma: no cover
+                        # Files generated by models are not sent back to models that don't themselves generate files.
+                        pass
                     else:
                         assert_never(item)
                 message_param = chat.ChatCompletionAssistantMessageParam(role='assistant')
@@ -782,31 +787,27 @@ class OpenAIChatModel(Model):
                                 identifier=item.identifier,
                             )
                         )
-                    else:
-                        base64_encoded = base64.b64encode(item.data).decode('utf-8')
-                        if item.is_image:
-                            image_url: ImageURL = {'url': f'data:{item.media_type};base64,{base64_encoded}'}
-                            if metadata := item.vendor_metadata:
-                                image_url['detail'] = metadata.get('detail', 'auto')
-                            content.append(ChatCompletionContentPartImageParam(image_url=image_url, type='image_url'))
-                        elif item.is_audio:
-                            assert item.format in ('wav', 'mp3')
-                            audio = InputAudio(data=base64_encoded, format=item.format)
-                            content.append(
-                                ChatCompletionContentPartInputAudioParam(input_audio=audio, type='input_audio')
-                            )
-                        elif item.is_document:
-                            content.append(
-                                File(
-                                    file=FileFile(
-                                        file_data=f'data:{item.media_type};base64,{base64_encoded}',
-                                        filename=f'filename.{item.format}',
-                                    ),
-                                    type='file',
-                                )
+                    elif item.is_image:
+                        image_url = ImageURL(url=item.data_uri)
+                        if metadata := item.vendor_metadata:
+                            image_url['detail'] = metadata.get('detail', 'auto')
+                        content.append(ChatCompletionContentPartImageParam(image_url=image_url, type='image_url'))
+                    elif item.is_audio:
+                        assert item.format in ('wav', 'mp3')
+                        audio = InputAudio(data=base64.b64encode(item.data).decode('utf-8'), format=item.format)
+                        content.append(ChatCompletionContentPartInputAudioParam(input_audio=audio, type='input_audio'))
+                    elif item.is_document:
+                        content.append(
+                            File(
+                                file=FileFile(
+                                    file_data=item.data_uri,
+                                    filename=f'filename.{item.format}',
+                                ),
+                                type='file',
                             )
-                        else:  # pragma: no cover
-                            raise RuntimeError(f'Unsupported binary content type: {item.media_type}')
+                        )
+                    else:  # pragma: no cover
+                        raise RuntimeError(f'Unsupported binary content type: {item.media_type}')
                 elif isinstance(item, AudioUrl):
                     downloaded_item = await download_item(item, data_format='base64', type_format='extension')
                     assert downloaded_item['data_type'] in (
@@ -941,7 +942,7 @@ class OpenAIResponsesModel(Model):
         response = await self._responses_create(
             messages, False, cast(OpenAIResponsesModelSettings, model_settings or {}), model_request_parameters
         )
-        return self._process_response(response)
+        return self._process_response(response, model_request_parameters)
 
     @asynccontextmanager
     async def request_stream(
@@ -962,7 +963,9 @@ class OpenAIResponsesModel(Model):
         async with response:
             yield await self._process_streamed_response(response, model_request_parameters)
 
-    def _process_response(self, response: responses.Response) -> ModelResponse:  # noqa: C901
+    def _process_response(  # noqa: C901
+        self, response: responses.Response, model_request_parameters: ModelRequestParameters
+    ) -> ModelResponse:
         """Process a non-streamed response, and prepare a message to return."""
         timestamp = number_to_datetime(response.created_at)
         items: list[ModelResponsePart] = []
@@ -1002,19 +1005,24 @@ class OpenAIResponsesModel(Model):
                     ToolCallPart(item.name, item.arguments, tool_call_id=_combine_tool_call_ids(item.call_id, item.id))
                 )
             elif isinstance(item, responses.ResponseCodeInterpreterToolCall):
-                call_part, return_part = _map_code_interpreter_tool_call(item, self.system)
+                call_part, return_part, file_parts = _map_code_interpreter_tool_call(item, self.system)
                 items.append(call_part)
+                if file_parts:
+                    items.extend(file_parts)
                 items.append(return_part)
             elif isinstance(item, responses.ResponseFunctionWebSearch):
                 call_part, return_part = _map_web_search_tool_call(item, self.system)
                 items.append(call_part)
                 items.append(return_part)
+            elif isinstance(item, responses.response_output_item.ImageGenerationCall):
+                call_part, return_part, file_part = _map_image_generation_tool_call(item, self.system)
+                items.append(call_part)
+                if file_part:  # pragma: no branch
+                    items.append(file_part)
+                items.append(return_part)
             elif isinstance(item, responses.ResponseComputerToolCall):  # pragma: no cover
                 # Pydantic AI doesn't yet support the ComputerUse built-in tool
                 pass
-            elif isinstance(item, responses.response_output_item.ImageGenerationCall):  # pragma: no cover
-                # Pydantic AI doesn't yet support the ImageGeneration built-in tool
-                pass
             elif isinstance(item, responses.ResponseCustomToolCall):  # pragma: no cover
                 # Support is being implemented in https://github.com/pydantic/pydantic-ai/pull/2572
                 pass
@@ -1204,6 +1212,7 @@ class OpenAIResponsesModel(Model):
 
     def _get_builtin_tools(self, model_request_parameters: ModelRequestParameters) -> list[responses.ToolParam]:
         tools: list[responses.ToolParam] = []
+        has_image_generating_tool = False
         for tool in model_request_parameters.builtin_tools:
             if isinstance(tool, WebSearchTool):
                 web_search_tool = responses.WebSearchToolParam(
@@ -1214,12 +1223,31 @@ class OpenAIResponsesModel(Model):
                         type='approximate', **tool.user_location
                     )
                 tools.append(web_search_tool)
-            elif isinstance(tool, CodeExecutionTool):  # pragma: no branch
+            elif isinstance(tool, CodeExecutionTool):
+                has_image_generating_tool = True
                 tools.append({'type': 'code_interpreter', 'container': {'type': 'auto'}})
+            elif isinstance(tool, ImageGenerationTool):  # pragma: no branch
+                has_image_generating_tool = True
+                tools.append(
+                    responses.tool_param.ImageGeneration(
+                        type='image_generation',
+                        background=tool.background,
+                        input_fidelity=tool.input_fidelity,
+                        moderation=tool.moderation,
+                        output_compression=tool.output_compression,
+                        output_format=tool.output_format or 'png',
+                        partial_images=tool.partial_images,
+                        quality=tool.quality,
+                        size=tool.size,
+                    )
+                )
             else:
                 raise UserError(  # pragma: no cover
                     f'`{tool.__class__.__name__}` is not supported by `OpenAIResponsesModel`. If it should be, please file an issue.'
                 )
+
+        if model_request_parameters.allow_image_output and not has_image_generating_tool:
+            tools.append({'type': 'image_generation'})
         return tools
 
     def _map_tool_definition(self, f: ToolDefinition) -> responses.FunctionToolParam:
@@ -1282,8 +1310,7 @@ class OpenAIResponsesModel(Model):
                         )
                         openai_messages.append(item)
                     elif isinstance(part, RetryPromptPart):
-                        # TODO(Marcelo): How do we test this conditional branch?
-                        if part.tool_name is None:  # pragma: no cover
+                        if part.tool_name is None:
                             openai_messages.append(
                                 Message(role='user', content=[{'type': 'input_text', 'text': part.model_response()}])
                             )
@@ -1342,7 +1369,7 @@ class OpenAIResponsesModel(Model):
                             param['id'] = id
                         openai_messages.append(param)
                     elif isinstance(item, BuiltinToolCallPart):
-                        if item.provider_name == self.system:
+                        if item.provider_name == self.system and send_item_ids:
                             if (
                                 item.tool_name == CodeExecutionTool.kind
                                 and item.tool_call_id
@@ -1353,7 +1380,7 @@ class OpenAIResponsesModel(Model):
                                     id=item.tool_call_id,
                                     code=args.get('code'),
                                     container_id=container_id,
-                                    outputs=None,
+                                    outputs=None,  # These can be read server-side
                                     status='completed',
                                     type='code_interpreter_call',
                                 )
@@ -1362,7 +1389,7 @@ class OpenAIResponsesModel(Model):
                                 item.tool_name == WebSearchTool.kind
                                 and item.tool_call_id
                                 and (args := item.args_as_dict())
-                            ):  # pragma: no branch
+                            ):
                                 web_search_item = responses.ResponseFunctionWebSearchParam(
                                     id=item.tool_call_id,
                                     action=cast(responses.response_function_web_search_param.Action, args),
@@ -1370,8 +1397,18 @@ class OpenAIResponsesModel(Model):
                                     type='web_search_call',
                                 )
                                 openai_messages.append(web_search_item)
+                            elif item.tool_name == ImageGenerationTool.kind and item.tool_call_id:  # pragma: no branch
+                                # The cast is necessary because of https://github.com/openai/openai-python/issues/2648
+                                image_generation_item = cast(
+                                    responses.response_input_item_param.ImageGenerationCall,
+                                    {
+                                        'id': item.tool_call_id,
+                                        'type': 'image_generation_call',
+                                    },
+                                )
+                                openai_messages.append(image_generation_item)
                     elif isinstance(item, BuiltinToolReturnPart):
-                        if item.provider_name == self.system:
+                        if item.provider_name == self.system and send_item_ids:
                             if (
                                 item.tool_name == CodeExecutionTool.kind
                                 and code_interpreter_item is not None
@@ -1379,7 +1416,6 @@ class OpenAIResponsesModel(Model):
                                 and (content := cast(dict[str, Any], item.content))  # pyright: ignore[reportUnknownMemberType]
                                 and (status := content.get('status'))
                             ):
-                                code_interpreter_item['outputs'] = content.get('outputs')
                                 code_interpreter_item['status'] = status
                             elif (
                                 item.tool_name == WebSearchTool.kind
@@ -1387,8 +1423,16 @@ class OpenAIResponsesModel(Model):
                                 and isinstance(item.content, dict)  # pyright: ignore[reportUnknownMemberType]
                                 and (content := cast(dict[str, Any], item.content))  # pyright: ignore[reportUnknownMemberType]
                                 and (status := content.get('status'))
-                            ):  # pragma: no branch
+                            ):
                                 web_search_item['status'] = status
+                            elif item.tool_name == ImageGenerationTool.kind:  # pragma: no branch
+                                # Image generation result does not need to be sent back, just the `id` off of `BuiltinToolCallPart`.
+                                pass
+                    elif isinstance(item, FilePart):
+                        # This was generated by the `ImageGenerationTool` or `CodeExecutionTool`,
+                        # and does not need to be sent back separately from the corresponding `BuiltinToolReturnPart`.
+                        # If `send_item_ids` is false, we won't send the `BuiltinToolReturnPart`, but OpenAI does not have a type for files from the assistant.
+                        pass
                     elif isinstance(item, ThinkingPart):
                         if item.id and send_item_ids:
                             signature: str | None = None
@@ -1454,7 +1498,6 @@ class OpenAIResponsesModel(Model):
                 if isinstance(item, str):
                     content.append(responses.ResponseInputTextParam(text=item, type='input_text'))
                 elif isinstance(item, BinaryContent):
-                    base64_encoded = base64.b64encode(item.data).decode('utf-8')
                     if item.is_image:
                         detail: Literal['auto', 'low', 'high'] = 'auto'
                         if metadata := item.vendor_metadata:
@@ -1464,7 +1507,7 @@ class OpenAIResponsesModel(Model):
                             )
                         content.append(
                             responses.ResponseInputImageParam(
-                                image_url=f'data:{item.media_type};base64,{base64_encoded}',
+                                image_url=item.data_uri,
                                 type='input_image',
                                 detail=detail,
                             )
@@ -1473,7 +1516,7 @@ class OpenAIResponsesModel(Model):
                         content.append(
                             responses.ResponseInputFileParam(
                                 type='input_file',
-                                file_data=f'data:{item.media_type};base64,{base64_encoded}',
+                                file_data=item.data_uri,
                                 # NOTE: Type wise it's not necessary to include the filename, but it's required by the
                                 # API itself. If we add empty string, the server sends a 500 error - which OpenAI needs
                                 # to fix. In any case, we add a placeholder name.
@@ -1532,7 +1575,6 @@ class OpenAIStreamedResponse(StreamedResponse):
 
     async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
         async for chunk in self._response:
-            print(chunk)
             self._usage += _map_usage(chunk)
 
             if chunk.id:  # pragma: no branch
@@ -1682,18 +1724,18 @@ class OpenAIResponsesStreamedResponse(StreamedResponse):
                     pass
                 elif isinstance(chunk.item, responses.ResponseFunctionWebSearch):
                     call_part, _ = _map_web_search_tool_call(chunk.item, self.provider_name)
-                    yield self._parts_manager.handle_builtin_tool_call_part(
+                    yield self._parts_manager.handle_part(
                         vendor_part_id=f'{chunk.item.id}-call', part=replace(call_part, args=None)
                     )
                 elif isinstance(chunk.item, responses.ResponseCodeInterpreterToolCall):
-                    call_part, _ = _map_code_interpreter_tool_call(chunk.item, self.provider_name)
+                    call_part, _, _ = _map_code_interpreter_tool_call(chunk.item, self.provider_name)
 
                     args_json = call_part.args_as_json_str()
                     # Drop the final `"}` so that we can add code deltas
                     args_json_delta = args_json[:-2]
                     assert args_json_delta.endswith('code":"')
 
-                    yield self._parts_manager.handle_builtin_tool_call_part(
+                    yield self._parts_manager.handle_part(
                         vendor_part_id=f'{chunk.item.id}-call', part=replace(call_part, args=None)
                     )
                     maybe_event = self._parts_manager.handle_tool_call_delta(
@@ -1702,6 +1744,10 @@ class OpenAIResponsesStreamedResponse(StreamedResponse):
                     )
                     if maybe_event is not None:  # pragma: no branch
                         yield maybe_event
+                elif isinstance(chunk.item, responses.response_output_item.ImageGenerationCall):
+                    call_part, _, _ = _map_image_generation_tool_call(chunk.item, self.provider_name)
+                    yield self._parts_manager.handle_part(vendor_part_id=f'{chunk.item.id}-call', part=call_part)
+
                 else:
                     warnings.warn(  # pragma: no cover
                         f'Handling of this item type is not yet implemented. Please report on our GitHub: {chunk}',
@@ -1719,10 +1765,12 @@ class OpenAIResponsesStreamedResponse(StreamedResponse):
                             provider_name=self.provider_name,
                         )
                 elif isinstance(chunk.item, responses.ResponseCodeInterpreterToolCall):
-                    _, return_part = _map_code_interpreter_tool_call(chunk.item, self.provider_name)
-                    yield self._parts_manager.handle_builtin_tool_return_part(
-                        vendor_part_id=f'{chunk.item.id}-return', part=return_part
-                    )
+                    _, return_part, file_parts = _map_code_interpreter_tool_call(chunk.item, self.provider_name)
+                    for i, file_part in enumerate(file_parts):
+                        yield self._parts_manager.handle_part(
+                            vendor_part_id=f'{chunk.item.id}-file-{i}', part=file_part
+                        )
+                    yield self._parts_manager.handle_part(vendor_part_id=f'{chunk.item.id}-return', part=return_part)
                 elif isinstance(chunk.item, responses.ResponseFunctionWebSearch):
                     call_part, return_part = _map_web_search_tool_call(chunk.item, self.provider_name)
 
@@ -1733,9 +1781,12 @@ class OpenAIResponsesStreamedResponse(StreamedResponse):
                     if maybe_event is not None:  # pragma: no branch
                         yield maybe_event
 
-                    yield self._parts_manager.handle_builtin_tool_return_part(
-                        vendor_part_id=f'{chunk.item.id}-return', part=return_part
-                    )
+                    yield self._parts_manager.handle_part(vendor_part_id=f'{chunk.item.id}-return', part=return_part)
+                elif isinstance(chunk.item, responses.response_output_item.ImageGenerationCall):
+                    _, return_part, file_part = _map_image_generation_tool_call(chunk.item, self.provider_name)
+                    if file_part:  # pragma: no branch
+                        yield self._parts_manager.handle_part(vendor_part_id=f'{chunk.item.id}-file', part=file_part)
+                    yield self._parts_manager.handle_part(vendor_part_id=f'{chunk.item.id}-return', part=return_part)
 
             elif isinstance(chunk, responses.ResponseReasoningSummaryPartAddedEvent):
                 yield self._parts_manager.handle_thinking_delta(
@@ -1757,8 +1808,8 @@ class OpenAIResponsesStreamedResponse(StreamedResponse):
                     id=chunk.item_id,
                 )
 
-            # TODO(Marcelo): We should support annotations in the future.
-            elif isinstance(chunk, responses.ResponseOutputTextAnnotationAddedEvent):  # pragma: no cover
+            elif isinstance(chunk, responses.ResponseOutputTextAnnotationAddedEvent):
+                # TODO(Marcelo): We should support annotations in the future.
                 pass  # there's nothing we need to do here
 
             elif isinstance(chunk, responses.ResponseTextDeltaEvent):
@@ -1809,6 +1860,28 @@ class OpenAIResponsesStreamedResponse(StreamedResponse):
             elif isinstance(chunk, responses.ResponseCodeInterpreterCallInterpretingEvent):
                 pass  # there's nothing we need to do here
 
+            elif isinstance(chunk, responses.ResponseImageGenCallCompletedEvent):  # pragma: no cover
+                pass  # there's nothing we need to do here
+
+            elif isinstance(chunk, responses.ResponseImageGenCallGeneratingEvent):
+                pass  # there's nothing we need to do here
+
+            elif isinstance(chunk, responses.ResponseImageGenCallInProgressEvent):
+                pass  # there's nothing we need to do here
+
+            elif isinstance(chunk, responses.ResponseImageGenCallPartialImageEvent):
+                # Not present on the type, but present on the actual object.
+                # See https://github.com/openai/openai-python/issues/2649
+                output_format = getattr(chunk, 'output_format', 'png')
+                file_part = FilePart(
+                    content=BinaryImage(
+                        data=base64.b64decode(chunk.partial_image_b64),
+                        media_type=f'image/{output_format}',
+                    ),
+                    id=chunk.item_id,
+                )
+                yield self._parts_manager.handle_part(vendor_part_id=f'{chunk.item_id}-file', part=file_part)
+
             else:  # pragma: no cover
                 warnings.warn(
                     f'Handling of this event type is not yet implemented. Please report on our GitHub: {chunk}',
@@ -1898,12 +1971,29 @@ def _split_combined_tool_call_id(combined_id: str) -> tuple[str, str | None]:
 
 def _map_code_interpreter_tool_call(
     item: responses.ResponseCodeInterpreterToolCall, provider_name: str
-) -> tuple[BuiltinToolCallPart, BuiltinToolReturnPart]:
+) -> tuple[BuiltinToolCallPart, BuiltinToolReturnPart, list[FilePart]]:
     result: dict[str, Any] = {
         'status': item.status,
     }
+
+    file_parts: list[FilePart] = []
+    logs: list[str] = []
     if item.outputs:
-        result['outputs'] = [output.model_dump(mode='json') for output in item.outputs]
+        for output in item.outputs:
+            if isinstance(output, responses.response_code_interpreter_tool_call.OutputImage):
+                file_parts.append(
+                    FilePart(
+                        content=BinaryImage.from_data_uri(output.url),
+                        id=item.id,
+                    )
+                )
+            elif isinstance(output, responses.response_code_interpreter_tool_call.OutputLogs):
+                logs.append(output.logs)
+            else:
+                assert_never(output)
+
+    if logs:
+        result['logs'] = logs
 
     return (
         BuiltinToolCallPart(
@@ -1921,6 +2011,7 @@ def _map_code_interpreter_tool_call(
             content=result,
             provider_name=provider_name,
         ),
+        file_parts,
     )
 
 
@@ -1954,3 +2045,51 @@ def _map_web_search_tool_call(
             provider_name=provider_name,
         ),
     )
+
+
+def _map_image_generation_tool_call(
+    item: responses.response_output_item.ImageGenerationCall, provider_name: str
+) -> tuple[BuiltinToolCallPart, BuiltinToolReturnPart, FilePart | None]:
+    result = {
+        'status': item.status,
+    }
+
+    # Not present on the type, but present on the actual object.
+    # See https://github.com/openai/openai-python/issues/2649
+    if background := getattr(item, 'background', None):
+        result['background'] = background
+    if quality := getattr(item, 'quality', None):
+        result['quality'] = quality
+    if size := getattr(item, 'size', None):
+        result['size'] = size
+    if revised_prompt := getattr(item, 'revised_prompt', None):
+        result['revised_prompt'] = revised_prompt
+    output_format = getattr(item, 'output_format', 'png')
+
+    file_part: FilePart | None = None
+    if item.result:
+        file_part = FilePart(
+            content=BinaryImage(
+                data=base64.b64decode(item.result),
+                media_type=f'image/{output_format}',
+            ),
+            id=item.id,
+        )
+
+        # For some reason, the streaming API leaves `status` as `generating` even though generation has completed.
+        result['status'] = 'completed'
+
+    return (
+        BuiltinToolCallPart(
+            tool_name=ImageGenerationTool.kind,
+            tool_call_id=item.id,
+            provider_name=provider_name,
+        ),
+        BuiltinToolReturnPart(
+            tool_name=ImageGenerationTool.kind,
+            tool_call_id=item.id,
+            content=result,
+            provider_name=provider_name,
+        ),
+        file_part,
+    )

pydantic_ai/models/test.py

@@ -17,6 +17,7 @@ from ..exceptions import UserError
 from ..messages import (
     BuiltinToolCallPart,
     BuiltinToolReturnPart,
+    FilePart,
     ModelMessage,
     ModelRequest,
     ModelResponse,
@@ -327,6 +328,9 @@ class TestStreamedResponse(StreamedResponse):
             elif isinstance(part, ThinkingPart):  # pragma: no cover
                 # NOTE: There's no way to reach this part of the code, since we don't generate ThinkingPart on TestModel.
                 assert False, "This should be unreachable — we don't generate ThinkingPart on TestModel."
+            elif isinstance(part, FilePart):  # pragma: no cover
+                # NOTE: There's no way to reach this part of the code, since we don't generate FilePart on TestModel.
+                assert False, "This should be unreachable — we don't generate FilePart on TestModel."
             else:
                 assert_never(part)
 

pydantic_ai/output.py

@@ -37,8 +37,11 @@ T_co = TypeVar('T_co', covariant=True)
 OutputDataT = TypeVar('OutputDataT', default=str, covariant=True)
 """Covariant type variable for the output data type of a run."""
 
-OutputMode = Literal['text', 'tool', 'native', 'prompted', 'tool_or_text']
-"""All output modes."""
+OutputMode = Literal['text', 'tool', 'native', 'prompted', 'tool_or_text', 'image']
+"""All output modes.
+
+`tool_or_text` is deprecated and no longer in use.
+"""
 StructuredOutputMode = Literal['tool', 'native', 'prompted']
 """Output modes that can be used for structured output. Used by ModelProfile.default_structured_output_mode"""
 

pydantic_ai/profiles/__init__.py

@@ -28,6 +28,8 @@ class ModelProfile:
     """Whether the model supports JSON schema output."""
     supports_json_object_output: bool = False
     """Whether the model supports JSON object output."""
+    supports_image_output: bool = False
+    """Whether the model supports image output."""
     default_structured_output_mode: StructuredOutputMode = 'tool'
     """The default structured output mode to use for the model."""
     prompted_output_template: str = dedent(

pydantic_ai/profiles/google.py

@@ -10,10 +10,13 @@ from . import ModelProfile
 
 def google_model_profile(model_name: str) -> ModelProfile | None:
     """Get the model profile for a Google model."""
+    is_image_model = 'image' in model_name
     return ModelProfile(
         json_schema_transformer=GoogleJsonSchemaTransformer,
-        supports_json_schema_output=True,
-        supports_json_object_output=True,
+        supports_image_output=is_image_model,
+        supports_json_schema_output=not is_image_model,
+        supports_json_object_output=not is_image_model,
+        supports_tools=not is_image_model,
     )
 
 

pydantic_ai/profiles/openai.py

@@ -26,7 +26,7 @@ class OpenAIModelProfile(ModelProfile):
     """Turn off to don't send sampling settings like `temperature` and `top_p` to models that don't support them, like OpenAI's o-series reasoning models."""
 
     openai_unsupported_model_settings: Sequence[str] = ()
-    """A list of model settings that are not supported by the model."""
+    """A list of model settings that are not supported by this model."""
 
     # Some OpenAI-compatible providers (e.g. MoonshotAI) currently do **not** accept
     # `tool_choice="required"`.  This flag lets the calling model know whether it's
@@ -84,6 +84,7 @@ def openai_model_profile(model_name: str) -> ModelProfile:
         json_schema_transformer=OpenAIJsonSchemaTransformer,
         supports_json_schema_output=True,
         supports_json_object_output=True,
+        supports_image_output=is_reasoning_model or '4.1' in model_name or '4o' in model_name,
         openai_unsupported_model_settings=openai_unsupported_model_settings,
         openai_system_prompt_role=openai_system_prompt_role,
         openai_chat_supports_web_search=supports_web_search,