(2) stream_agent_response can handle agents sdk image gen output

Author: jiwon-oai

chatkit/agents.py

@@ -8,6 +8,8 @@
     Annotated,
     Any,
     AsyncGenerator,
+    Awaitable,
+    Callable,
     Generic,
     Sequence,
     TypeVar,
@@ -55,6 +57,8 @@
     DurationSummary,
     EndOfTurnItem,
     FileSource,
+    GeneratedImage,
+    GeneratedImageItem,
     HiddenContextItem,
     SDKHiddenContextItem,
     Task,
@@ -105,6 +109,7 @@ class AgentContext(BaseModel, Generic[TContext]):
     previous_response_id: str | None = None
     client_tool_call: ClientToolCall | None = None
     workflow_item: WorkflowItem | None = None
+    generated_image_item: GeneratedImageItem | None = None
     _events: asyncio.Queue[ThreadStreamEvent | _QueueCompleteSentinel] = asyncio.Queue()
 
     def generate_id(
@@ -357,7 +362,10 @@ class StreamingThoughtTracker(BaseModel):
 
 
 async def stream_agent_response(
-    context: AgentContext, result: RunResultStreaming
+    context: AgentContext,
+    result: RunResultStreaming,
+    *,
+    base64_to_generated_image: Callable[[str], Awaitable[GeneratedImage]] | None = None,
 ) -> AsyncIterator[ThreadStreamEvent]:
     """Convert a streamed Agents SDK run into ChatKit ThreadStreamEvents."""
     current_item_id = None
@@ -527,6 +535,15 @@ def end_workflow(item: WorkflowItem):
                             created_at=datetime.now(),
                         ),
                     )
+                elif item.type == "image_generation_call":
+                    ctx.generated_image_item = GeneratedImageItem(
+                        id=ctx.generate_id("message"),
+                        thread_id=thread.id,
+                        created_at=datetime.now(),
+                        image=None,
+                    )
+                    produced_items.add(ctx.generated_image_item.id)
+                    yield ThreadItemAddedEvent(item=ctx.generated_image_item)
             elif event.type == "response.reasoning_summary_text.delta":
                 if not ctx.workflow_item:
                     continue
@@ -604,6 +621,22 @@ def end_workflow(item: WorkflowItem):
                             created_at=datetime.now(),
                         ),
                     )
+                elif item.type == "image_generation_call" and item.result:
+                    if not ctx.generated_image_item:
+                        continue
+
+                    # Agents SDK only produces png and base64 output currently.
+                    if base64_to_generated_image:
+                        image = await base64_to_generated_image(item.result)
+                    else:
+                        image = GeneratedImage(
+                            id=item.id, url=f"data:image/png;base64,{item.result}"
+                        )
+
+                    ctx.generated_image_item.image = image
+                    yield ThreadItemDoneEvent(item=ctx.generated_image_item)
+
+                    ctx.generated_image_item = None
 
     except (InputGuardrailTripwireTriggered, OutputGuardrailTripwireTriggered):
         for item_id in produced_items:
@@ -694,6 +727,17 @@ async def tag_to_message_content(
             "A Tag was included in a UserMessageItem but Converter.tag_to_message_content is not implemented"
         )
 
+    async def generated_image_to_input(
+        self, item: GeneratedImageItem
+    ) -> TResponseInputItem | list[TResponseInputItem] | None:
+        """
+        Convert a GeneratedImageItem into a message content part to send to the model.
+        Required when generated images are enabled.
+        """
+        raise NotImplementedError(
+            "A GeneratedImageItem was included in a UserMessageItem but Converter.generated_image_to_message_content was not implemented"
+        )
+
     async def hidden_context_to_input(
         self, item: HiddenContextItem
     ) -> TResponseInputItem | list[TResponseInputItem] | None:
@@ -984,6 +1028,9 @@ async def _thread_item_to_input_item(
             case SDKHiddenContextItem():
                 out = await self.sdk_hidden_context_to_input(item) or []
                 return out if isinstance(out, list) else [out]
+            case GeneratedImageItem():
+                out = await self.generated_image_to_input(item) or []
+                return out if isinstance(out, list) else [out]
             case _:
                 assert_never(item)