update base64_image_to_url to take image id and partial image index as input to url generation

Author: jiwon-oai

chatkit/agents.py

@@ -386,12 +386,25 @@ def __init__(self, partial_images: int | None = None):
         """
         self.partial_images = partial_images
 
-    async def base64_image_to_url(self, base64_image: str) -> str:
+    async def base64_image_to_url(
+        self,
+        image_id: str,
+        base64_image: str,
+        partial_image_index: int | None = None,
+    ) -> str:
         """
         Convert a base64-encoded image into a URL.
 
         This method is used to produce the URL stored on thread items for image
         generation results.
+
+        Args:
+            image_id: The ID of the image generation call. This stays stable across partial image updates.
+            base64_image: The base64-encoded image.
+            partial_image_index: The index of the partial image update, starting from 0.
+
+        Returns:
+            A URL string.
         """
         return f"data:image/png;base64,{base64_image}"
 
@@ -695,7 +708,7 @@ def end_workflow(item: WorkflowItem):
                     if not ctx.generated_image_item:
                         continue
 
-                    url = await converter.base64_image_to_url(item.result)
+                    url = await converter.base64_image_to_url(item.id, item.result)
                     image = GeneratedImage(id=item.id, url=url)
 
                     ctx.generated_image_item.image = image
@@ -706,7 +719,11 @@ def end_workflow(item: WorkflowItem):
                 if not ctx.generated_image_item:
                     continue
 
-                url = await converter.base64_image_to_url(event.partial_image_b64)
+                url = await converter.base64_image_to_url(
+                    event.item_id,
+                    event.partial_image_b64,
+                    event.partial_image_index,
+                )
                 progress = converter.partial_image_index_to_progress(
                     event.partial_image_index
                 )

tests/test_agents.py

@@ -1282,11 +1282,16 @@ async def test_stream_agent_response_image_generation_events_with_custom_convert
     class CustomResponseStreamConverter(ResponseStreamConverter):
         def __init__(self):
             super().__init__()
-            self.calls: list[str] = []
+            self.calls: list[tuple[str, str, int | None]] = []
 
-        async def base64_image_to_url(self, base64_image: str) -> str:
-            self.calls.append(base64_image)
-            return f"https://example.com/{base64_image}"
+        async def base64_image_to_url(
+            self,
+            image_id: str,
+            base64_image: str,
+            partial_image_index: int | None = None,
+        ) -> str:
+            self.calls.append((image_id, base64_image, partial_image_index))
+            return f"https://example.com/{image_id}"
 
     converter = CustomResponseStreamConverter()
     stream = stream_agent_response(context, result, converter=converter)
@@ -1298,9 +1303,9 @@ async def base64_image_to_url(self, base64_image: str) -> str:
     event2 = await stream.__anext__()
     assert isinstance(event2, ThreadItemDoneEvent)
     assert isinstance(event2.item, GeneratedImageItem)
-    assert converter.calls == ["dGVzdA=="]
+    assert converter.calls == [("img_call_1", "dGVzdA==", None)]
     assert event2.item.image == GeneratedImage(
-        id="img_call_1", url="https://example.com/dGVzdA=="
+        id="img_call_1", url="https://example.com/img_call_1"
     )
     with pytest.raises(StopAsyncIteration):
         await stream.__anext__()