From 7624cc45e6ce47303985d894d0fa27f473306e62 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 29 Nov 2024 13:40:04 -0800 Subject: [PATCH] fix(transformation.py): support mp4 + pdf url's for vertex ai Fixes https://github.com/BerriAI/litellm/issues/6936 --- .../gemini/transformation.py | 4 ++++ tests/llm_translation/test_vertex.py | 15 +++++++++++++++ 2 files changed, 19 insertions(+) diff --git a/litellm/llms/vertex_ai_and_google_ai_studio/gemini/transformation.py b/litellm/llms/vertex_ai_and_google_ai_studio/gemini/transformation.py index 4b5b7281b..c9fe6e3f4 100644 --- a/litellm/llms/vertex_ai_and_google_ai_studio/gemini/transformation.py +++ b/litellm/llms/vertex_ai_and_google_ai_studio/gemini/transformation.py @@ -107,6 +107,10 @@ def _get_image_mime_type_from_url(url: str) -> Optional[str]: return "image/png" elif url.endswith(".webp"): return "image/webp" + elif url.endswith(".mp4"): + return "video/mp4" + elif url.endswith(".pdf"): + return "application/pdf" return None diff --git a/tests/llm_translation/test_vertex.py b/tests/llm_translation/test_vertex.py index 1ea2514c9..425b6f9f4 100644 --- a/tests/llm_translation/test_vertex.py +++ b/tests/llm_translation/test_vertex.py @@ -1146,6 +1146,21 @@ def test_process_gemini_image(): mime_type="image/png", file_uri="https://example.com/image.png" ) + # Test HTTPS VIDEO URL + https_result = _process_gemini_image("https://cloud-samples-data/video/animals.mp4") + print("https_result PNG", https_result) + assert https_result["file_data"] == FileDataType( + mime_type="video/mp4", file_uri="https://cloud-samples-data/video/animals.mp4" + ) + + # Test HTTPS PDF URL + https_result = _process_gemini_image("https://cloud-samples-data/pdf/animals.pdf") + print("https_result PDF", https_result) + assert https_result["file_data"] == FileDataType( + mime_type="application/pdf", + file_uri="https://cloud-samples-data/pdf/animals.pdf", + ) + # Test base64 image base64_image = "data:image/jpeg;base64,/9j/4AAQSkZJRg..." base64_result = _process_gemini_image(base64_image)