fix(vertex_ai/gemini/transformation.py): handle 'http://' image urls

BerriAI · Jan 16, 2025 · 370d669 · 370d669
1 parent 4ec2c2c
commit 370d669
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 8 deletions.
diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py
@@ -76,10 +76,9 @@ def _process_gemini_image(image_url: str) -> PartType:
             file_data = FileDataType(mime_type=mime_type, file_uri=image_url)
 
             return PartType(file_data=file_data)
-        elif (
-            "https://" in image_url
-            and (image_type := _get_image_mime_type_from_url(image_url)) is not None
-        ):
+        elif ("http://" in image_url or "https://" in image_url) and (
+            image_type := _get_image_mime_type_from_url(image_url)
+        ) is not None:
             file_data = FileDataType(file_uri=image_url, mime_type=image_type)
             return PartType(file_data=file_data)
         elif "http://" in image_url or "https://" in image_url or "base64" in image_url:

diff --git a/tests/llm_translation/test_vertex.py b/tests/llm_translation/test_vertex.py
@@ -1289,10 +1289,9 @@ def test_process_gemini_image_http_url(
         http_url: Test HTTP URL
         mock_convert_to_anthropic: Mocked convert_to_anthropic_image_obj function
         mock_blob: Mocked BlobType instance
-    """
-    # Arrange
-    expected_image_data = "data:image/jpeg;base64,/9j/4AAQSkZJRg..."
-    mock_convert_url_to_base64.return_value = expected_image_data
 
+    Vertex AI supports image urls. Ensure no network requests are made.
+    """
     # Act
     result = _process_gemini_image(http_url)
+    assert result["file_data"]["file_uri"] == http_url