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

fix(vertex_ai.py): support video + pdf for vertex ai

Browse source
This commit is contained in:
Krrish Dholakia 2024-06-06 08:05:32 -07:00
parent e01c0d9a44
commit 58bd2b4ea6
1 changed files with 14 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

23
litellm/llms/vertex_ai.py
View file

@ -297,24 +297,29 @@ def _convert_gemini_role(role: str) -> Literal["user", "model"]:
def _process_gemini_image(image_url: str) -> PartType: def _process_gemini_image(image_url: str) -> PartType:
try: try:
if "gs://" in image_url: if ".mp4" in image_url and "gs://" in image_url:
# Case 1: Images with Cloud Storage URIs # Case 1: Videos with Cloud Storage URIs
part_mime = "video/mp4"
_file_data = FileDataType(mime_type=part_mime, file_uri=image_url)
return PartType(file_data=_file_data)
elif ".pdf" in image_url and "gs://" in image_url:
# Case 2: PDF's with Cloud Storage URIs
part_mime = "application/pdf"
_file_data = FileDataType(mime_type=part_mime, file_uri=image_url)
return PartType(file_data=_file_data)
elif "gs://" in image_url:
# Case 3: Images with Cloud Storage URIs
# The supported MIME types for images include image/png and image/jpeg. # The supported MIME types for images include image/png and image/jpeg.
part_mime = "image/png" if "png" in image_url else "image/jpeg" part_mime = "image/png" if "png" in image_url else "image/jpeg"
_file_data = FileDataType(mime_type=part_mime, file_uri=image_url) _file_data = FileDataType(mime_type=part_mime, file_uri=image_url)
return PartType(file_data=_file_data) return PartType(file_data=_file_data)
elif "https:/" in image_url: elif "https:/" in image_url:
# Case 2: Images with direct links # Case 4: Images with direct links
image = _load_image_from_url(image_url) image = _load_image_from_url(image_url)
_blob = BlobType(data=image.data, mime_type=image._mime_type) _blob = BlobType(data=image.data, mime_type=image._mime_type)
return PartType(inline_data=_blob) return PartType(inline_data=_blob)
elif ".mp4" in image_url and "gs://" in image_url:
# Case 3: Videos with Cloud Storage URIs
part_mime = "video/mp4"
_file_data = FileDataType(mime_type=part_mime, file_uri=image_url)
return PartType(file_data=_file_data)
elif "base64" in image_url: elif "base64" in image_url:
# Case 4: Images with base64 encoding # Case 5: Images with base64 encoding
import base64, re import base64, re
# base 64 is passed as data:image/jpeg;base64,<base-64-encoded-image> # base 64 is passed as data:image/jpeg;base64,<base-64-encoded-image>