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