diff --git a/python/packages/chatkit/AGENTS.md b/python/packages/chatkit/AGENTS.md index d11724b3ad..e99e011ecd 100644 --- a/python/packages/chatkit/AGENTS.md +++ b/python/packages/chatkit/AGENTS.md @@ -11,6 +11,9 @@ Integration with OpenAI ChatKit (Python) for building chat UIs. `ThreadItemConverter.structured_input_to_input()` is an async override hook returning `Message | list[Message] | None`; `to_agent_input()` normalizes its output and preserves thread order. +`ThreadItemConverter.generated_image_to_input()` is an async override hook returning +`Message | list[Message] | None`; data URI media types are inferred by `Content.from_uri()`. + ## Usage ```python diff --git a/python/packages/chatkit/README.md b/python/packages/chatkit/README.md index ac154fbec2..94516fb06c 100644 --- a/python/packages/chatkit/README.md +++ b/python/packages/chatkit/README.md @@ -16,6 +16,11 @@ Override the async `ThreadItemConverter.structured_input_to_input()` method to customize structured input context, such as formatting or redacting answers. Return a `Message`, a list of messages, or `None` to skip the item. +Override the async `ThreadItemConverter.generated_image_to_input()` method to +customize generated image context, such as resolving private image URLs. +Return a `Message`, a list of messages, or `None` to skip the item. +The default converter preserves data URI media types and leaves external URLs unchanged. + ## Installation ```bash diff --git a/python/packages/chatkit/agent_framework_chatkit/_converter.py b/python/packages/chatkit/agent_framework_chatkit/_converter.py index f6501178f9..d878da2f06 100644 --- a/python/packages/chatkit/agent_framework_chatkit/_converter.py +++ b/python/packages/chatkit/agent_framework_chatkit/_converter.py @@ -524,6 +524,36 @@ async def structured_input_to_input(self, item: StructuredInputItem) -> Message ) return Message(role="user", contents=[text]) + async def generated_image_to_input(self, item: GeneratedImageItem) -> Message | list[Message] | None: + """Convert a generated image to Agent Framework Message(s). + + This method is called internally by `to_agent_input()`. Override this method + to customize image conversion, such as resolving a private image URL. + + Args: + item: The ChatKit generated image item to convert. + + Returns: + A Message with user role, a list of messages, or None to skip. + + Raises: + ContentError: If the generated image URL is invalid. + + Note: + Use `to_agent_input()` to convert thread items with proper message ordering. + """ + if item.image is None: + return None + + uri = item.image.url + return Message( + role="user", + contents=[ + Content.from_text(text="The following image was generated by the agent."), + Content.from_uri(uri=uri, media_type=None if uri.startswith("data:") else "image/*"), + ], + ) + async def _thread_item_to_input_item( self, item: ThreadItem, @@ -567,8 +597,8 @@ async def _thread_item_to_input_item( out = self.hidden_context_to_input(item) or [] return out if isinstance(out, list) else [out] case GeneratedImageItem(): - # TODO(evmattso): Implement generated image handling in a future PR - return [] + out = await self.generated_image_to_input(item) or [] + return out if isinstance(out, list) else [out] case StructuredInputItem(): out = await self.structured_input_to_input(item) or [] return out if isinstance(out, list) else [out] diff --git a/python/packages/chatkit/tests/test_converter.py b/python/packages/chatkit/tests/test_converter.py index b9889f14db..2633abdda0 100644 --- a/python/packages/chatkit/tests/test_converter.py +++ b/python/packages/chatkit/tests/test_converter.py @@ -745,6 +745,118 @@ async def structured_input_to_input(self, item: StructuredInputItem) -> Message assert exception.value is error + @pytest.mark.parametrize("url", ["https://example.com/image.png", "https://example.com/asset?signature=sample"]) + async def test_to_agent_input_with_generated_image(self, converter, url): + """Test generated images are preserved as multimodal model context.""" + from chatkit.types import GeneratedImage, GeneratedImageItem + + input_item = GeneratedImageItem( + id="generated_1", + thread_id="thread_1", + created_at=datetime.now(), + type="generated_image", + image=GeneratedImage(id="image_1", url=url), + ) + + result = await converter.to_agent_input(input_item) + + assert len(result) == 1 + assert result[0].role == "user" + assert [content.type for content in result[0].contents] == ["text", "uri"] + assert result[0].contents[0].text == "The following image was generated by the agent." + assert result[0].contents[1].uri == url + assert result[0].contents[1].media_type == "image/*" + + @pytest.mark.parametrize("media_type", ["image/png", "image/jpeg", "image/webp"]) + async def test_to_agent_input_preserves_generated_image_data_uri_media_type(self, converter, media_type): + """Test data URI media types remain available to provider serializers.""" + from chatkit.types import GeneratedImage, GeneratedImageItem + + uri = f"data:{media_type};base64,YQ==" + input_item = GeneratedImageItem( + id="generated_1", + thread_id="thread_1", + created_at=datetime.now(), + image=GeneratedImage(id="image_1", url=uri), + ) + + result = await converter.to_agent_input(input_item) + + assert len(result) == 1 + assert result[0].role == "user" + assert [content.type for content in result[0].contents] == ["text", "data"] + assert result[0].contents[1].uri == uri + assert result[0].contents[1].media_type == media_type + + async def test_generated_image_to_input_without_image_returns_none(self, converter): + """Test unfinished image items are skipped by both conversion entry points.""" + from chatkit.types import GeneratedImageItem + + input_item = GeneratedImageItem(id="generated_1", thread_id="thread_1", created_at=datetime.now()) + + assert await converter.generated_image_to_input(input_item) is None + assert await converter.to_agent_input(input_item) == [] + + @pytest.mark.parametrize( + ("converted", "expected_texts"), + [ + (None, []), + ([], []), + (Message("user", ["image omitted"]), ["image omitted"]), + ([Message("user", ["first"]), Message("user", ["second"])], ["first", "second"]), + ], + ids=["skip", "empty", "single", "multiple"], + ) + async def test_to_agent_input_uses_generated_image_override( + self, converted: Message | list[Message] | None, expected_texts: list[str] + ) -> None: + """Test async image overrides can skip or expand an item while preserving order.""" + from chatkit.types import GeneratedImage, GeneratedImageItem, HiddenContextItem + + calls: list[GeneratedImageItem] = [] + + class CustomConverter(ThreadItemConverter): + async def generated_image_to_input(self, item: GeneratedImageItem) -> Message | list[Message] | None: + calls.append(item) + return converted + + input_item = GeneratedImageItem( + id="generated_1", + thread_id="thread_1", + created_at=datetime.now(), + image=GeneratedImage(id="image_1", url="https://example.com/image.png"), + ) + before = HiddenContextItem(id="before", thread_id="thread_1", created_at=datetime.now(), content="before") + after = HiddenContextItem(id="after", thread_id="thread_1", created_at=datetime.now(), content="after") + + result = await CustomConverter().to_agent_input([before, input_item, after]) + + assert calls == [input_item] + assert [message.text for message in result[1:-1]] == expected_texts + assert result[0].text == "before" + assert result[-1].text == "after" + + @pytest.mark.parametrize("error", [ValueError("conversion failed"), asyncio.CancelledError()]) + async def test_to_agent_input_propagates_generated_image_override_errors(self, error: BaseException) -> None: + """Test custom image conversion failures and cancellation propagate to the caller.""" + from chatkit.types import GeneratedImage, GeneratedImageItem + + class FailingConverter(ThreadItemConverter): + async def generated_image_to_input(self, item: GeneratedImageItem) -> Message | list[Message] | None: + raise error + + input_item = GeneratedImageItem( + id="generated_1", + thread_id="thread_1", + created_at=datetime.now(), + image=GeneratedImage(id="image_1", url="https://example.com/image.png"), + ) + + with pytest.raises(type(error)) as exception: + await FailingConverter().to_agent_input(input_item) + + assert exception.value is error + async def test_to_agent_input_dispatches_supported_variants(self, converter): """Test thread item dispatch converts supported items and skips unsupported variants.""" from chatkit.types import (