From fc300219dc39a90f50b92f1fe5bb42f0a40f9013 Mon Sep 17 00:00:00 2001 From: Anish Mehta Date: Wed, 16 Sep 2026 16:21:45 +0530 Subject: [PATCH] fix(openai): emit image_generation_call result from streaming output_item.done With the default `partial_images=0` the Responses API delivers the final image only once, as a completed `image_generation_call` item on `response.output_item.done`. The streaming parser had no branch for that item type, so the image was silently dropped and `AgentResponse` contained no image content, while the non-streaming parser mapped it correctly. Factor the non-streaming mapping into `_image_generation_item_to_contents` and use it from both parsers so streaming emits the same `image_generation_tool_call` + `image_generation_tool_result` contents. Fixes #8422 Co-Authored-By: Claude Fable 5.1 --- .../agent_framework_openai/_chat_client.py | 58 ++++++---- .../tests/openai/test_openai_chat_client.py | 108 ++++++++++++++++++ 2 files changed, 143 insertions(+), 23 deletions(-) diff --git a/python/packages/openai/agent_framework_openai/_chat_client.py b/python/packages/openai/agent_framework_openai/_chat_client.py index f6e18a761ec..cb836e90046 100644 --- a/python/packages/openai/agent_framework_openai/_chat_client.py +++ b/python/packages/openai/agent_framework_openai/_chat_client.py @@ -2209,6 +2209,36 @@ def _to_shell_call_output_payload(content: Content) -> list[dict[str, Any]]: } ] + def _image_generation_item_to_contents(self, item: Any) -> list[Content]: + """Convert a completed ``image_generation_call`` output item into framework ``Content`` objects. + + Used by both the non-streaming parser and the streaming + ``response.output_item.done`` handler. With the default ``partial_images=0`` + no ``response.image_generation_call.partial_image`` events are emitted, so the + completed item is the only place the final base64 image is delivered. + """ + image_output: Content | None = None + image_result = getattr(item, "result", None) + if image_result is not None: + # item.result contains raw base64 string + # so we call detect_media_type_from_base64 to get the media type and fallback to image/png + image_output = Content.from_uri( + uri=f"data:{detect_media_type_from_base64(data_str=image_result) or 'image/png'};base64,{image_result}", + raw_representation=image_result, + ) + image_id = getattr(item, "id", None) + return [ + Content.from_image_generation_tool_call( + image_id=image_id, + raw_representation=item, + ), + Content.from_image_generation_tool_result( + image_id=image_id, + outputs=image_output, + raw_representation=item, + ), + ] + def _shell_item_to_contents(self, item: Any, local_shell_tool_name: str | None) -> list[Content]: """Convert a shell output item into framework ``Content`` objects. @@ -2983,29 +3013,7 @@ def _parse_response_from_openai( ) ) case "image_generation_call": # ResponseOutputImageGenerationCall - image_output: Content | None = None - if item.result is not None: - # item.result contains raw base64 string - # so we call detect_media_type_from_base64 to get the media type and fallback to image/png - image_output = Content.from_uri( - uri=f"data:{detect_media_type_from_base64(data_str=item.result) or 'image/png'}" - f";base64,{item.result}", - raw_representation=item.result, - ) - image_id = item.id - contents.append( - Content.from_image_generation_tool_call( - image_id=image_id, - raw_representation=item, - ) - ) - contents.append( - Content.from_image_generation_tool_result( - image_id=image_id, - outputs=image_output, - raw_representation=item, - ) - ) + contents.extend(self._image_generation_item_to_contents(item)) case "shell_call" | "local_shell_call" | "shell_call_output": contents.extend(self._shell_item_to_contents(item, local_shell_tool_name)) case _: @@ -3597,6 +3605,10 @@ def _get_ann_value(key: str) -> Any: # Shell items are parsed here (not on `response.output_item.added`) because the # command/output is only populated on the completed item. contents.extend(self._shell_item_to_contents(done_item, local_shell_tool_name)) + elif getattr(done_item, "type", None) == "image_generation_call": + # The final image is only delivered on the completed item; with the default + # `partial_images=0` there are no `partial_image` events at all. + contents.extend(self._image_generation_item_to_contents(done_item)) elif getattr(done_item, "type", None) == "custom_tool_call": custom_tool_call = cast(ResponseCustomToolCall, done_item) contents.append( diff --git a/python/packages/openai/tests/openai/test_openai_chat_client.py b/python/packages/openai/tests/openai/test_openai_chat_client.py index 80ae0aa9ef9..0754c035cdf 100644 --- a/python/packages/openai/tests/openai/test_openai_chat_client.py +++ b/python/packages/openai/tests/openai/test_openai_chat_client.py @@ -3480,6 +3480,114 @@ def test_parse_chunk_from_openai_with_file_search_call_done() -> None: assert content.result == {"results": [{"file_id": "file_1", "text": "Seattle was cloudy."}]} +def test_parse_chunk_from_openai_image_generation_call_done_emits_final_image() -> None: + """Test that response.output_item.done for image_generation_call emits the final image. + + With the default ``partial_images=0`` there are no ``partial_image`` events; the completed + item on ``response.output_item.done`` is the only place the base64 result is delivered. + """ + client = OpenAIChatClient(model="test-model", api_key="test-key") + chat_options: dict[str, Any] = {} + function_call_ids: dict[int, tuple[str, str]] = {} + + png_signature = b"\x89PNG\r\n\x1a\n" + image_base64 = base64.b64encode(png_signature + b"fake_png_data_here").decode() + + mock_event = MagicMock() + mock_event.type = "response.output_item.done" + + mock_item = MagicMock() + mock_item.type = "image_generation_call" + mock_item.id = "ig_123" + mock_item.status = "completed" + mock_item.result = image_base64 + mock_event.item = mock_item + + update = client._parse_chunk_from_openai(mock_event, options=chat_options, function_call_ids=function_call_ids) + + assert len(update.contents) == 2 + call_content, result_content = update.contents + assert call_content.type == "image_generation_tool_call" + assert call_content.image_id == "ig_123" + assert call_content.raw_representation is mock_item + assert result_content.type == "image_generation_tool_result" + assert result_content.image_id == "ig_123" + assert result_content.raw_representation is mock_item + data_out = result_content.outputs + assert isinstance(data_out, Content) + assert data_out.type == "data" + assert data_out.media_type == "image/png" + assert data_out.uri == f"data:image/png;base64,{image_base64}" + + +def test_parse_chunk_from_openai_image_generation_call_done_without_result() -> None: + """Test that a completed image_generation_call without a result still emits call + empty result. + + Mirrors the non-streaming parser, which emits both contents with ``outputs=None``. + """ + client = OpenAIChatClient(model="test-model", api_key="test-key") + chat_options: dict[str, Any] = {} + function_call_ids: dict[int, tuple[str, str]] = {} + + mock_event = MagicMock() + mock_event.type = "response.output_item.done" + + mock_item = MagicMock() + mock_item.type = "image_generation_call" + mock_item.id = "ig_456" + mock_item.status = "failed" + mock_item.result = None + mock_event.item = mock_item + + update = client._parse_chunk_from_openai(mock_event, options=chat_options, function_call_ids=function_call_ids) + + assert len(update.contents) == 2 + call_content, result_content = update.contents + assert call_content.type == "image_generation_tool_call" + assert call_content.image_id == "ig_456" + assert result_content.type == "image_generation_tool_result" + assert result_content.image_id == "ig_456" + assert result_content.outputs is None + + +def test_parse_chunk_from_openai_image_generation_call_done_matches_non_streaming() -> None: + """Test that streaming and non-streaming parsers produce the same contents for a completed image item.""" + client = OpenAIChatClient(model="test-model", api_key="test-key") + + jpeg_signature = b"\xff\xd8\xff" + image_base64 = base64.b64encode(jpeg_signature + b"fake_jpeg_data").decode() + + mock_item = MagicMock() + mock_item.type = "image_generation_call" + mock_item.id = "ig_789" + mock_item.status = "completed" + mock_item.result = image_base64 + + mock_event = MagicMock() + mock_event.type = "response.output_item.done" + mock_event.item = mock_item + + update = client._parse_chunk_from_openai(mock_event, options={}, function_call_ids={}) + + mock_response = MagicMock() + mock_response.output_parsed = None + mock_response.metadata = {} + mock_response.usage = None + mock_response.id = "test-response-id" + mock_response.model = "test-model" + mock_response.created_at = 1234567890 + mock_response.output = [mock_item] + + with patch.object(client, "_get_metadata_from_response", return_value={}): + response = client._parse_response_from_openai(mock_response, options={}) # type: ignore + + streamed = [c.to_dict() for c in update.contents] + non_streamed = [c.to_dict() for c in response.messages[0].contents] + assert streamed == non_streamed + assert [c["type"] for c in streamed] == ["image_generation_tool_call", "image_generation_tool_result"] + assert streamed[1]["outputs"]["uri"] == f"data:image/jpeg;base64,{image_base64}" + + def test_parse_chunk_from_openai_shell_call_added_defers_command() -> None: """An in-progress shell_call on output_item.added has no command yet, so it must emit nothing.