Python: OpenAI Responses Image Generation Stream Support, Sample and Unit Tests (#1853)

* support for image gen streaming

* small fixes

* fixes

* added comment

python/packages/core/agent_framework/_types.py

@@ -1050,6 +1050,50 @@ class DataContent(BaseContent):
     def has_top_level_media_type(self, top_level_media_type: Literal["application", "audio", "image", "text"]) -> bool:
         return _has_top_level_media_type(self.media_type, top_level_media_type)
 
+    @staticmethod
+    def detect_image_format_from_base64(image_base64: str) -> str:
+        """Detect image format from base64 data by examining the binary header.
+
+        Args:
+            image_base64: Base64 encoded image data
+
+        Returns:
+            Image format as string (png, jpeg, webp, gif) with png as fallback
+        """
+        try:
+            # Constants for image format detection
+            # ~75 bytes of binary data should be enough to detect most image formats
+            FORMAT_DETECTION_BASE64_CHARS = 100
+
+            # Decode a small portion to detect format
+            decoded_data = base64.b64decode(image_base64[:FORMAT_DETECTION_BASE64_CHARS])
+            if decoded_data.startswith(b"\x89PNG"):
+                return "png"
+            if decoded_data.startswith(b"\xff\xd8\xff"):
+                return "jpeg"
+            if decoded_data.startswith(b"RIFF") and b"WEBP" in decoded_data[:12]:
+                return "webp"
+            if decoded_data.startswith(b"GIF87a") or decoded_data.startswith(b"GIF89a"):
+                return "gif"
+            return "png"  # Default fallback
+        except Exception:
+            return "png"  # Fallback if decoding fails
+
+    @classmethod
+    def create_data_uri_from_base64(cls, image_base64: str) -> tuple[str, str]:
+        """Create a data URI and media type from base64 image data.
+
+        Args:
+            image_base64: Base64 encoded image data
+
+        Returns:
+            Tuple of (data_uri, media_type)
+        """
+        format_type = cls.detect_image_format_from_base64(image_base64)
+        uri = f"data:image/{format_type};base64,{image_base64}"
+        media_type = f"image/{format_type}"
+        return uri, media_type
+
 
 class UriContent(BaseContent):
     """Represents a URI content.

python/packages/core/agent_framework/openai/_responses_client.py

@@ -293,6 +293,14 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
                             # Map the parameter name and remove the old one
                             mapped_tool[api_param] = mapped_tool.pop(user_param)
 
+                    # Validate partial_images parameter for streaming image generation
+                    # OpenAI API requires partial_images to be between 0-3 (inclusive) for image_generation tool
+                    # Reference: https://platform.openai.com/docs/api-reference/responses/create#responses_create-tools-image_generation_tool-partial_images
+                    if "partial_images" in mapped_tool:
+                        partial_images = mapped_tool["partial_images"]
+                        if not isinstance(partial_images, int) or partial_images < 0 or partial_images > 3:
+                            raise ValueError("partial_images must be an integer between 0 and 3 (inclusive).")
+
                     response_tools.append(mapped_tool)
                 else:
                     response_tools.append(tool_dict)
@@ -695,29 +703,8 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
                         uri = item.result
                         media_type = None
                         if not uri.startswith("data:"):
-                            # Raw base64 string - convert to proper data URI format
-                            # Detect format from base64 data
-                            import base64
-
-                            try:
-                                # Decode a small portion to detect format
-                                decoded_data = base64.b64decode(uri[:100])  # First ~75 bytes should be enough
-                                if decoded_data.startswith(b"\x89PNG"):
-                                    format_type = "png"
-                                elif decoded_data.startswith(b"\xff\xd8\xff"):
-                                    format_type = "jpeg"
-                                elif decoded_data.startswith(b"RIFF") and b"WEBP" in decoded_data[:12]:
-                                    format_type = "webp"
-                                elif decoded_data.startswith(b"GIF87a") or decoded_data.startswith(b"GIF89a"):
-                                    format_type = "gif"
-                                else:
-                                    # Default to png if format cannot be detected
-                                    format_type = "png"
-                            except Exception:
-                                # Fallback to png if decoding fails
-                                format_type = "png"
-                            uri = f"data:image/{format_type};base64,{uri}"
-                            media_type = f"image/{format_type}"
+                            # Raw base64 string - convert to proper data URI format using helper
+                            uri, media_type = DataContent.create_data_uri_from_base64(uri)
                         else:
                             # Parse media type from existing data URI
                             try:
@@ -933,6 +920,25 @@ class OpenAIBaseResponsesClient(OpenAIBase, BaseChatClient):
                             raw_representation=event,
                         )
                     )
+            case "response.image_generation_call.partial_image":
+                # Handle streaming partial image generation
+                image_base64 = event.partial_image_b64
+                partial_index = event.partial_image_index
+
+                # Use helper function to create data URI from base64
+                uri, media_type = DataContent.create_data_uri_from_base64(image_base64)
+
+                contents.append(
+                    DataContent(
+                        uri=uri,
+                        media_type=media_type,
+                        additional_properties={
+                            "partial_image_index": partial_index,
+                            "is_partial_image": True,
+                        },
+                        raw_representation=event,
+                    )
+                )
             case _:
                 logger.debug("Unparsed event of type: %s: %s", event.type, event)
 

python/packages/core/tests/core/test_types.py

@@ -1,5 +1,6 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+import base64
 from collections.abc import AsyncIterable
 from typing import Any
 
@@ -166,6 +167,57 @@ def test_data_content_empty():
         DataContent(uri="")
 
 
+def test_data_content_detect_image_format_from_base64():
+    """Test the detect_image_format_from_base64 static method."""
+    # Test each supported format
+    png_data = b"\x89PNG\r\n\x1a\n" + b"fake_data"
+    assert DataContent.detect_image_format_from_base64(base64.b64encode(png_data).decode()) == "png"
+
+    jpeg_data = b"\xff\xd8\xff\xe0" + b"fake_data"
+    assert DataContent.detect_image_format_from_base64(base64.b64encode(jpeg_data).decode()) == "jpeg"
+
+    webp_data = b"RIFF" + b"1234" + b"WEBP" + b"fake_data"
+    assert DataContent.detect_image_format_from_base64(base64.b64encode(webp_data).decode()) == "webp"
+
+    gif_data = b"GIF89a" + b"fake_data"
+    assert DataContent.detect_image_format_from_base64(base64.b64encode(gif_data).decode()) == "gif"
+
+    # Test fallback behavior
+    unknown_data = b"UNKNOWN_FORMAT"
+    assert DataContent.detect_image_format_from_base64(base64.b64encode(unknown_data).decode()) == "png"
+
+    # Test error handling
+    assert DataContent.detect_image_format_from_base64("invalid_base64!") == "png"
+    assert DataContent.detect_image_format_from_base64("") == "png"
+
+
+def test_data_content_create_data_uri_from_base64():
+    """Test the create_data_uri_from_base64 class method."""
+    # Test with PNG data
+    png_data = b"\x89PNG\r\n\x1a\n" + b"fake_data"
+    png_base64 = base64.b64encode(png_data).decode()
+    uri, media_type = DataContent.create_data_uri_from_base64(png_base64)
+
+    assert uri == f"data:image/png;base64,{png_base64}"
+    assert media_type == "image/png"
+
+    # Test with different format
+    jpeg_data = b"\xff\xd8\xff\xe0" + b"fake_data"
+    jpeg_base64 = base64.b64encode(jpeg_data).decode()
+    uri, media_type = DataContent.create_data_uri_from_base64(jpeg_base64)
+
+    assert uri == f"data:image/jpeg;base64,{jpeg_base64}"
+    assert media_type == "image/jpeg"
+
+    # Test fallback for unknown format
+    unknown_data = b"UNKNOWN_FORMAT"
+    unknown_base64 = base64.b64encode(unknown_data).decode()
+    uri, media_type = DataContent.create_data_uri_from_base64(unknown_base64)
+
+    assert uri == f"data:image/png;base64,{unknown_base64}"
+    assert media_type == "image/png"
+
+
 # region UriContent
 
 

python/samples/getting_started/agents/openai/README.md

@@ -23,6 +23,7 @@ This folder contains examples demonstrating different ways to create and use age
 | [`openai_responses_client_image_analysis.py`](openai_responses_client_image_analysis.py) | Demonstrates how to use vision capabilities with agents to analyze images. |
 | [`openai_responses_client_image_generation.py`](openai_responses_client_image_generation.py) | Demonstrates how to use image generation capabilities with OpenAI agents to create images based on text descriptions. Requires PIL (Pillow) for image display. |
 | [`openai_responses_client_reasoning.py`](openai_responses_client_reasoning.py) | Demonstrates how to use reasoning capabilities with OpenAI agents, showing how the agent can provide detailed reasoning for its responses. |
+| [`openai_responses_client_streaming_image_generation.py`](openai_responses_client_streaming_image_generation.py) | Demonstrates streaming image generation with partial images for real-time image creation feedback and improved user experience. |
 | [`openai_responses_client_with_code_interpreter.py`](openai_responses_client_with_code_interpreter.py) | Shows how to use the HostedCodeInterpreterTool with OpenAI agents to write and execute Python code. Includes helper methods for accessing code interpreter data from response chunks. |
 | [`openai_responses_client_with_explicit_settings.py`](openai_responses_client_with_explicit_settings.py) | Shows how to initialize an agent with a specific responses client, configuring settings explicitly including API key and model ID. |
 | [`openai_responses_client_with_file_search.py`](openai_responses_client_with_file_search.py) | Demonstrates how to use file search capabilities with OpenAI agents, allowing the agent to search through uploaded files to answer questions. |

python/samples/getting_started/agents/openai/openai_responses_client_streaming_image_generation.py

@@ -0,0 +1,96 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+import asyncio
+import base64
+
+import anyio
+from agent_framework import DataContent
+from agent_framework.openai import OpenAIResponsesClient
+
+"""OpenAI Responses Client Streaming Image Generation Example
+
+Demonstrates streaming partial image generation using OpenAI's image generation tool.
+Shows progressive image rendering with partial images for improved user experience.
+
+Note: The number of partial images received depends on generation speed:
+- High quality/complex images: More partials (generation takes longer)
+- Low quality/simple images: Fewer partials (generation completes quickly)
+- You may receive fewer partial images than requested if generation is fast
+
+Important: The final partial image IS the complete, full-quality image. Each partial
+represents a progressive refinement, with the last one being the finished result.
+"""
+
+
+async def save_image_from_data_uri(data_uri: str, filename: str) -> None:
+    """Save an image from a data URI to a file."""
+    try:
+        if data_uri.startswith("data:image/"):
+            # Extract base64 data
+            base64_data = data_uri.split(",", 1)[1]
+            image_bytes = base64.b64decode(base64_data)
+
+            # Save to file
+            await anyio.Path(filename).write_bytes(image_bytes)
+            print(f"    Saved: {filename} ({len(image_bytes) / 1024:.1f} KB)")
+    except Exception as e:
+        print(f"    Error saving {filename}: {e}")
+
+
+async def main():
+    """Demonstrate streaming image generation with partial images."""
+    print("=== OpenAI Streaming Image Generation Example ===\n")
+
+    # Create agent with streaming image generation enabled
+    agent = OpenAIResponsesClient().create_agent(
+        instructions="You are a helpful agent that can generate images.",
+        tools=[
+            {
+                "type": "image_generation",
+                "size": "1024x1024",
+                "quality": "high",
+                "partial_images": 3,
+            }
+        ],
+    )
+
+    query = "Draw a beautiful sunset over a calm ocean with sailboats"
+    print(f" User: {query}")
+    print()
+
+    # Track partial images
+    image_count = 0
+
+    # Create output directory
+    output_dir = anyio.Path("generated_images")
+    await output_dir.mkdir(exist_ok=True)
+
+    print(" Streaming response:")
+    async for update in agent.run_stream(query):
+        for content in update.contents:
+            # Handle partial images
+            # The final partial image IS the complete, full-quality image. Each partial
+            # represents a progressive refinement, with the last one being the finished result.
+            if isinstance(content, DataContent) and content.additional_properties.get("is_partial_image"):
+                print(f"     Image {image_count} received")
+
+                # Extract file extension from media_type (e.g., "image/png" -> "png")
+                extension = "png"  # Default fallback
+                if content.media_type and "/" in content.media_type:
+                    extension = content.media_type.split("/")[-1]
+
+                # Save images with correct extension
+                filename = output_dir / f"image{image_count}.{extension}"
+                await save_image_from_data_uri(content.uri, str(filename))
+
+                image_count += 1
+
+    # Summary
+    print("\n Summary:")
+    print(f"    Images received: {image_count}")
+    print("    Output directory: generated_images")
+    print("\n Streaming image generation completed!")
+
+
+if __name__ == "__main__":
+    asyncio.run(main())