Add Regression tests for image_url blocks in assistant message content.

This commit is contained in:
Sameer Kankute
2026-02-27 12:04:12 +05:30
parent f1080a7e30
commit 596437b3b9
@@ -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}"
)