Merge pull request #22155 from BerriAI/litellm_fix_image

[Bug]Add ChatCompletionImageObject in OpenAIChatCompletionAssistantMessage

litellm/types/llms/openai.py

@@ -699,7 +699,15 @@ class OpenAIChatCompletionAssistantMessage(TypedDict, total=False):
     role: Required[Literal["assistant"]]
     content: Optional[
         Union[
-            str, Iterable[Union[ChatCompletionTextObject, ChatCompletionThinkingBlock]]
+            str,
+            Iterable[
+                Union[
+                    ChatCompletionTextObject,
+                    ChatCompletionThinkingBlock,
+                    ChatCompletionRedactedThinkingBlock,
+                    ChatCompletionImageObject,
+                ]
+            ],
         ]
     ]
     name: Optional[str]
@@ -786,17 +794,19 @@ ValidUserMessageContentTypes = [
     "file",
 ]  # used for validating user messages. Prevent users from accidentally sending anthropic messages.
 
-# Assistant message content types (text, thinking, redacted_thinking)
+# Assistant message content types (text, thinking, redacted_thinking, image_url)
 ValidAssistantMessageContentTypesLiteral = Literal[
     "text",
     "thinking",
     "redacted_thinking",
+    "image_url",
 ]
 
 ValidAssistantMessageContentTypes = [
     "text",
     "thinking",
     "redacted_thinking",
+    "image_url",
 ]
 
 # Combined valid content types for chat completion messages

tests/test_litellm/types/llms/test_types_llms_openai.py

@@ -169,3 +169,97 @@ class TestResponsesAPIResponseOutputText:
         )
 
         assert response.output_text == ""
+
+
+class TestAssistantMessageImageUrlContent:
+    """
+    Regression tests for image_url blocks in assistant message content.
+
+    Bug: ChatCompletionAssistantMessage.content did not include
+    ChatCompletionImageObject in its union, so Pydantic v2 silently dropped
+    image_url blocks (content → []) when serialising via AllMessageValues.
+    This affects users who store conversation history as JSON (e.g. in a DB)
+    and read it back typed as list[AllMessageValues].
+    """
+
+    ASSISTANT_MESSAGE_WITH_IMAGE = {
+        "role": "assistant",
+        "content": [
+            {"type": "text", "text": "Here is the image you requested:"},
+            {
+                "type": "image_url",
+                "image_url": {
+                    "url": (
+                        "data:image/png;base64,"
+                        "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAA"
+                        "DUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
+                    )
+                },
+            },
+        ],
+    }
+
+    def test_assistant_message_image_url_preserved_single(self):
+        """
+        TypeAdapter(ChatCompletionAssistantMessage): image_url block must survive
+        validate_python → dump_python without being dropped or raising an error.
+        """
+        from typing import List
+
+        from pydantic import TypeAdapter
+
+        from litellm.types.llms.openai import ChatCompletionAssistantMessage
+
+        adapter = TypeAdapter(ChatCompletionAssistantMessage)
+        validated = adapter.validate_python(self.ASSISTANT_MESSAGE_WITH_IMAGE)
+        dumped = adapter.dump_python(validated)
+
+        raw_content = dumped.get("content")
+        # Pydantic may return a lazy SerializationIterator for Iterable fields;
+        # convert to list to consume it — this must not raise ValidationError.
+        content_blocks = list(raw_content) if raw_content is not None else []
+
+        assert len(content_blocks) == 2, (
+            f"Expected 2 content blocks (text + image_url), got {len(content_blocks)}: {content_blocks}"
+        )
+        types = [b.get("type") for b in content_blocks if isinstance(b, dict)]
+        assert "image_url" in types, f"image_url block was silently dropped; blocks: {content_blocks}"
+
+    def test_assistant_message_image_url_preserved_in_all_message_values(self):
+        """
+        TypeAdapter(List[AllMessageValues]) DB round-trip: image_url blocks in an
+        assistant message must not be silently dropped during dump_python(mode='json').
+
+        This is the primary failing path: conversation history stored as JSON in a
+        database and read back typed as list[AllMessageValues].
+        """
+        from typing import List
+
+        from pydantic import TypeAdapter
+
+        from litellm.types.llms.openai import AllMessageValues
+
+        conversation = [
+            {
+                "role": "user",
+                "content": "Generate an image of a banana wearing a LiteLLM costume",
+            },
+            self.ASSISTANT_MESSAGE_WITH_IMAGE,
+        ]
+
+        adapter = TypeAdapter(List[AllMessageValues])
+        validated = adapter.validate_python(conversation)
+        dumped = adapter.dump_python(validated, mode="json")
+
+        assistant = next((m for m in dumped if m.get("role") == "assistant"), None)
+        assert assistant is not None, "Assistant message missing after serialisation"
+
+        content = assistant.get("content", [])
+        assert isinstance(content, list), f"content should be a list, got {type(content)}"
+        assert len(content) == 2, (
+            f"Expected 2 content blocks (text + image_url), got {len(content)}: {content}"
+        )
+        types = [b.get("type") for b in content if isinstance(b, dict)]
+        assert "image_url" in types, (
+            f"image_url block was silently dropped during AllMessageValues serialisation; blocks: {content}"
+        )