phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

fix(transformation.py): support mp4 + pdf url's for vertex ai

Browse source 
Fixes https://github.com/BerriAI/litellm/issues/6936
This commit is contained in:
Krrish Dholakia 2024-11-29 13:40:04 -08:00
parent 828bf909fe
commit 7624cc45e6
2 changed files with 19 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

4
litellm/llms/vertex_ai_and_google_ai_studio/gemini/transformation.py
View file

@ -107,6 +107,10 @@ def _get_image_mime_type_from_url(url: str) -> Optional[str]:
return "image/png" return "image/png"
elif url.endswith(".webp"): elif url.endswith(".webp"):
return "image/webp" return "image/webp"
elif url.endswith(".mp4"):
return "video/mp4"
elif url.endswith(".pdf"):
return "application/pdf"
return None return None

15
tests/llm_translation/test_vertex.py
View file

@ -1146,6 +1146,21 @@ def test_process_gemini_image():
mime_type="image/png", file_uri="https://example.com/image.png" 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 # Test base64 image
base64_image = "data:image/jpeg;base64,/9j/4AAQSkZJRg..." base64_image = "data:image/jpeg;base64,/9j/4AAQSkZJRg..."
base64_result = _process_gemini_image(base64_image) base64_result = _process_gemini_image(base64_image)