tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-16 10:23:14 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #20845 from BerriAI/litellm_gemini_image_handling

Browse Source 
Handle image in assitant message for gemini
This commit is contained in:
Sameer Kankute
2026-02-10 18:24:09 +05:30
committed by GitHub
parent 3de892b8ca f6228fda3e
commit 0f01802dde
2 changed files with 273 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/llms/vertex_ai/gemini/transformation.py
+21
View File
@@ -437,6 +437,27 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915
else:
assistant_content.append(PartType(text=assistant_text)) # type: ignore
## HANDLE ASSISTANT IMAGES FIELD
# Process images field if present (for generated images from assistant)
assistant_images = assistant_msg.get("images")
if assistant_images is not None and isinstance(assistant_images, list):
for image_item in assistant_images:
if isinstance(image_item, dict):
image_url_obj = image_item.get("image_url")
if isinstance(image_url_obj, dict):
assistant_image_url = image_url_obj.get("url")
format = image_url_obj.get("format")
detail = image_url_obj.get("detail")
media_resolution_enum = _convert_detail_to_media_resolution_enum(detail)
if assistant_image_url:
_part = _process_gemini_media(
image_url=assistant_image_url,
format=format,
media_resolution_enum=media_resolution_enum,
model=model,
)
assistant_content.append(_part)
## HANDLE ASSISTANT FUNCTION CALL
if (
assistant_msg.get("tool_calls", []) is not None
tests/test_litellm/llms/vertex_ai/gemini/test_vertex_ai_gemini_transformation.py
+252 -14
View File
@@ -7,6 +7,7 @@ from litellm.llms.vertex_ai.gemini.transformation import (
check_if_part_exists_in_parts,
)
from litellm.types.llms.vertex_ai import BlobType
from litellm.types.utils import Message
def test_check_if_part_exists_in_parts():
@@ -735,8 +736,9 @@ def test_file_data_field_order():
Related issue: Gemini API returns 400 INVALID_ARGUMENT when fields are in wrong order.
"""
import json
from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_media
# Test with HTTPS URL and explicit format (audio file)
file_url = "https://generativelanguage.googleapis.com/v1beta/files/test123"
format = "audio/mpeg"
@@ -770,8 +772,9 @@ def test_file_data_field_order():
def test_file_data_field_order_gcs_urls():
"""Test that GCS URLs also maintain correct field order."""
import json
from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_media
# Test with GCS URL
gcs_url = "gs://bucket/audio.mp3"
@@ -802,11 +805,14 @@ def test_extract_file_data_with_path_object():
Related issue: Files uploaded with wrong MIME type cause Gemini API to reject
requests where the specified format doesn't match the uploaded file's MIME type.
"""
from pathlib import Path
from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data
import tempfile
import os
import tempfile
from pathlib import Path
from litellm.litellm_core_utils.prompt_templates.common_utils import (
extract_file_data,
)
# Create a temporary MP3 file
with tempfile.NamedTemporaryFile(suffix=".mp3", delete=False) as tmp:
tmp.write(b"fake mp3 content")
@@ -835,10 +841,13 @@ def test_extract_file_data_with_path_object():
def test_extract_file_data_with_string_path():
"""Test that filename is correctly extracted from string paths."""
from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data
import tempfile
import os
import tempfile
from litellm.litellm_core_utils.prompt_templates.common_utils import (
extract_file_data,
)
# Create a temporary WAV file
with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
tmp.write(b"fake wav content")
@@ -866,8 +875,10 @@ def test_extract_file_data_with_string_path():
def test_extract_file_data_with_tuple_format():
"""Test that tuple format (with explicit content_type) still works correctly."""
from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data
from litellm.litellm_core_utils.prompt_templates.common_utils import (
extract_file_data,
)
# Test with tuple format: (filename, content, content_type)
filename = "test_audio.mp3"
content = b"test audio content"
@@ -883,10 +894,13 @@ def test_extract_file_data_with_tuple_format():
def test_extract_file_data_fallback_to_octet_stream():
"""Test that unknown file types fall back to application/octet-stream."""
from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data
import tempfile
import os
import tempfile
from litellm.litellm_core_utils.prompt_templates.common_utils import (
extract_file_data,
)
# Create a temporary file with unknown extension
with tempfile.NamedTemporaryFile(suffix=".xyz123", delete=False) as tmp:
tmp.write(b"unknown content")
@@ -1086,3 +1100,227 @@ def test_convert_tool_response_with_nested_file_object():
assert "mime_type" in inline_data
assert inline_data["mime_type"] == "application/pdf"
assert inline_data["data"] == test_pdf_base64
def test_assistant_message_with_images_field():
"""
Test that assistant messages with images field are properly converted to Gemini format.
This handles the case where an assistant message contains generated images in the
`images` field (e.g., from image generation models like gemini-2.5-flash-image).
The images should be converted to inline_data parts in the Gemini format.
"""
# Create a small test image (1x1 red pixel PNG)
test_image_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
image_data_uri = f"data:image/png;base64,{test_image_base64}"
# Create messages with assistant message containing images field
messages = [
{
"role": "user",
"content": "Generate an image of a banana wearing a costume that says LiteLLM"
},
{
"role": "assistant",
"content": "Here's your banana in a LiteLLM costume!",
"images": [
{
"image_url": {
"url": image_data_uri,
"detail": "auto"
},
"index": 0,
"type": "image_url"
}
]
}
]
# Convert messages to Gemini format
contents = _gemini_convert_messages_with_history(messages=messages)
# Verify structure
assert len(contents) == 2, f"Expected 2 content blocks, got {len(contents)}"
# Verify user message
assert contents[0]["role"] == "user"
assert len(contents[0]["parts"]) == 1
assert contents[0]["parts"][0]["text"] == "Generate an image of a banana wearing a costume that says LiteLLM"
# Verify assistant message
assert contents[1]["role"] == "model"
assert len(contents[1]["parts"]) == 2, f"Expected 2 parts (text + image), got {len(contents[1]['parts'])}"
# Find text part and inline_data part
text_part = None
inline_data_part = None
for part in contents[1]["parts"]:
if "text" in part:
text_part = part
elif "inline_data" in part:
inline_data_part = part
# Verify text part
assert text_part is not None, "Missing text part in assistant message"
assert text_part["text"] == "Here's your banana in a LiteLLM costume!"
# Verify inline_data part (image)
assert inline_data_part is not None, "Missing inline_data part in assistant message"
inline_data: BlobType = inline_data_part["inline_data"]
assert "data" in inline_data
assert "mime_type" in inline_data
assert inline_data["mime_type"] == "image/png"
assert inline_data["data"] == test_image_base64
def test_assistant_message_with_multiple_images():
"""Test that assistant messages with multiple images are properly converted."""
# Create two test images
test_image1_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
test_image2_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8DwHwAFBQIAX8jx0gAAAABJRU5ErkJggg=="
image1_data_uri = f"data:image/png;base64,{test_image1_base64}"
image2_data_uri = f"data:image/jpeg;base64,{test_image2_base64}"
messages = [
{
"role": "user",
"content": "Generate two images"
},
{
"role": "assistant",
"content": "Here are your images:",
"images": [
{
"image_url": {
"url": image1_data_uri,
"detail": "auto"
},
"index": 0,
"type": "image_url"
},
{
"image_url": {
"url": image2_data_uri,
"detail": "high"
},
"index": 1,
"type": "image_url"
}
]
}
]
# Convert messages to Gemini format
contents = _gemini_convert_messages_with_history(messages=messages)
# Verify assistant message has 3 parts (1 text + 2 images)
assert contents[1]["role"] == "model"
assert len(contents[1]["parts"]) == 3, f"Expected 3 parts (text + 2 images), got {len(contents[1]['parts'])}"
# Count inline_data parts
inline_data_parts = [part for part in contents[1]["parts"] if "inline_data" in part]
assert len(inline_data_parts) == 2, f"Expected 2 inline_data parts, got {len(inline_data_parts)}"
# Verify first image
assert inline_data_parts[0]["inline_data"]["mime_type"] == "image/png"
assert inline_data_parts[0]["inline_data"]["data"] == test_image1_base64
# Verify second image
assert inline_data_parts[1]["inline_data"]["mime_type"] == "image/jpeg"
assert inline_data_parts[1]["inline_data"]["data"] == test_image2_base64
def test_assistant_message_with_images_using_message_object():
"""Test that Message objects with images field are properly converted."""
# Create a small test image
test_image_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
image_data_uri = f"data:image/png;base64,{test_image_base64}"
# Create messages using Message object (as returned by LiteLLM)
user_message = {
"role": "user",
"content": "Generate an image"
}
assistant_message = Message(
content="Here's your image!",
role="assistant",
tool_calls=None,
function_call=None,
images=[
{
"image_url": {
"url": image_data_uri,
"detail": "auto"
},
"index": 0,
"type": "image_url"
}
]
)
messages = [user_message, assistant_message]
# Convert messages to Gemini format
contents = _gemini_convert_messages_with_history(messages=messages)
# Verify assistant message has both text and image
assert contents[1]["role"] == "model"
assert len(contents[1]["parts"]) == 2
# Verify image was converted
inline_data_parts = [part for part in contents[1]["parts"] if "inline_data" in part]
assert len(inline_data_parts) == 1
assert inline_data_parts[0]["inline_data"]["mime_type"] == "image/png"
assert inline_data_parts[0]["inline_data"]["data"] == test_image_base64
def test_assistant_message_with_images_in_conversation_history():
"""
Test multi-turn conversation where assistant message with images is in history.
This simulates the real use case where:
1. User asks for image generation
2. Assistant generates image (with images field)
3. User asks follow-up question about the image
"""
test_image_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNk+M9QDwADhgGAWjR9awAAAABJRU5ErkJggg=="
image_data_uri = f"data:image/png;base64,{test_image_base64}"
messages = [
{
"role": "user",
"content": "Generate an image of a cat"
},
{
"role": "assistant",
"content": "Here's a cat image:",
"images": [
{
"image_url": {
"url": image_data_uri,
"detail": "auto"
},
"index": 0,
"type": "image_url"
}
]
},
{
"role": "user",
"content": "Can you make it more colorful?"
}
]
# Convert messages to Gemini format
contents = _gemini_convert_messages_with_history(messages=messages)
# Verify structure: user -> model (with image) -> user
assert len(contents) == 3
assert contents[0]["role"] == "user"
assert contents[1]["role"] == "model"
assert contents[2]["role"] == "user"
# Verify assistant message has image in history
inline_data_parts = [part for part in contents[1]["parts"] if "inline_data" in part]
assert len(inline_data_parts) == 1
assert inline_data_parts[0]["inline_data"]["mime_type"] == "image/png"