Tests

Author: jiwon-oai

tests/test_agents.py

@@ -78,6 +78,8 @@
     CustomTask,
     DurationSummary,
     FileSource,
+    GeneratedImage,
+    GeneratedImageItem,
     HiddenContextItem,
     InferenceOptions,
     Page,
@@ -1191,6 +1193,115 @@ async def test_stream_agent_response_assistant_message_content_types():
     assert message.id == "1"
 
 
+async def test_stream_agent_response_image_generation_events_default_data_url():
+    context = AgentContext(
+        previous_response_id=None, thread=thread, store=mock_store, request_context=None
+    )
+    result = make_result()
+
+    result.add_event(
+        RawResponsesStreamEvent(
+            type="raw_response_event",
+            data=Mock(
+                type="response.output_item.added",
+                item=Mock(type="image_generation_call", id="img_call_1"),
+                output_index=0,
+                sequence_number=0,
+            ),
+        )
+    )
+    result.add_event(
+        RawResponsesStreamEvent(
+            type="raw_response_event",
+            data=Mock(
+                type="response.output_item.done",
+                item=Mock(
+                    type="image_generation_call", id="img_call_1", result="dGVzdA=="
+                ),
+                output_index=0,
+                sequence_number=1,
+            ),
+        )
+    )
+    result.done()
+
+    stream = stream_agent_response(context, result)
+    event1 = await stream.__anext__()
+    assert isinstance(event1, ThreadItemAddedEvent)
+    assert isinstance(event1.item, GeneratedImageItem)
+    assert event1.item.type == "generated_image"
+    assert event1.item.id == "message_id"
+    assert event1.item.image is None
+
+    event2 = await stream.__anext__()
+    assert isinstance(event2, ThreadItemDoneEvent)
+    assert isinstance(event2.item, GeneratedImageItem)
+    assert event2.item.id == event1.item.id
+    assert event2.item.image == GeneratedImage(
+        id="img_call_1", url="data:image/png;base64,dGVzdA=="
+    )
+
+    with pytest.raises(StopAsyncIteration):
+        await stream.__anext__()
+
+
+async def test_stream_agent_response_image_generation_events_with_async_converter():
+    context = AgentContext(
+        previous_response_id=None, thread=thread, store=mock_store, request_context=None
+    )
+    result = make_result()
+
+    result.add_event(
+        RawResponsesStreamEvent(
+            type="raw_response_event",
+            data=Mock(
+                type="response.output_item.added",
+                item=Mock(type="image_generation_call", id="img_call_1"),
+                output_index=0,
+                sequence_number=0,
+            ),
+        )
+    )
+    result.add_event(
+        RawResponsesStreamEvent(
+            type="raw_response_event",
+            data=Mock(
+                type="response.output_item.done",
+                item=Mock(
+                    type="image_generation_call", id="img_call_1", result="dGVzdA=="
+                ),
+                output_index=0,
+                sequence_number=1,
+            ),
+        )
+    )
+    result.done()
+
+    calls: list[str] = []
+
+    async def base64_to_generated_image(b64: str) -> GeneratedImage:
+        calls.append(b64)
+        return GeneratedImage(id="custom_img_1", url=f"https://example.com/{b64}")
+
+    stream = stream_agent_response(
+        context, result, base64_to_generated_image=base64_to_generated_image
+    )
+    event1 = await stream.__anext__()
+    assert isinstance(event1, ThreadItemAddedEvent)
+    assert isinstance(event1.item, GeneratedImageItem)
+    assert event1.item.image is None
+
+    event2 = await stream.__anext__()
+    assert isinstance(event2, ThreadItemDoneEvent)
+    assert isinstance(event2.item, GeneratedImageItem)
+    assert calls == ["dGVzdA=="]
+    assert event2.item.image == GeneratedImage(
+        id="custom_img_1", url="https://example.com/dGVzdA=="
+    )
+    with pytest.raises(StopAsyncIteration):
+        await stream.__anext__()
+
+
 async def test_workflow_streams_first_thought():
     context = AgentContext(
         previous_response_id=None, thread=thread, store=mock_store, request_context=None