mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-05 02:41:56 +00:00
Fix Gemini MIME detection for extensionless GCS URIs.
Propagate litellm_params across Gemini transformers, resolve MIME from GCS metadata for extensionless gs:// objects, normalize file MIME aliases (format/mime_type/content_type), improve error mapping to BadRequestError, and add regression coverage for explicit MIME, metadata success, and file_data error-message behavior. Made-with: Cursor
This commit is contained in:
parent
b770372555
commit
29020e9316
4 changed files with 415 additions and 26 deletions
|
|
@ -103,7 +103,10 @@ class GoogleAIStudioGeminiConfig(VertexGeminiConfig):
|
|||
return supported_params
|
||||
|
||||
def _transform_messages(
|
||||
self, messages: List[AllMessageValues], model: Optional[str] = None
|
||||
self,
|
||||
messages: List[AllMessageValues],
|
||||
model: Optional[str] = None,
|
||||
litellm_params: Optional[dict] = None,
|
||||
) -> List[ContentType]:
|
||||
"""
|
||||
Google AI Studio Gemini does not support HTTP/HTTPS URLs for files.
|
||||
|
|
@ -160,4 +163,6 @@ class GoogleAIStudioGeminiConfig(VertexGeminiConfig):
|
|||
except Exception:
|
||||
# If conversion fails, leave as is and let the API handle it
|
||||
pass
|
||||
return _gemini_convert_messages_with_history(messages=messages, model=model)
|
||||
return _gemini_convert_messages_with_history(
|
||||
messages=messages, model=model, litellm_params=litellm_params
|
||||
)
|
||||
|
|
|
|||
|
|
@ -6,7 +6,9 @@ Why separate file? Make it easy to see how transformation works
|
|||
|
||||
import json
|
||||
import os
|
||||
from typing import TYPE_CHECKING, Dict, List, Literal, Optional, Tuple, Union, cast
|
||||
import re
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union, cast
|
||||
from urllib.parse import quote
|
||||
|
||||
import httpx
|
||||
from pydantic import BaseModel
|
||||
|
|
@ -26,6 +28,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
|
|||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
from litellm.llms.vertex_ai.common_utils import pop_vertex_request_labels
|
||||
from litellm.types.files import (
|
||||
get_file_extension_from_mime_type,
|
||||
get_file_mime_type_for_file_type,
|
||||
get_file_type_from_extension,
|
||||
is_gemini_1_5_accepted_file_type,
|
||||
|
|
@ -56,6 +59,12 @@ from ..common_utils import (
|
|||
get_supports_response_schema,
|
||||
get_supports_system_message,
|
||||
)
|
||||
from ..vertex_llm_base import VertexBase
|
||||
|
||||
_GCS_METADATA_VERTEX_BASE = VertexBase()
|
||||
_GEMINI_MIME_TYPE_ALIASES: Dict[str, str] = {
|
||||
"image/jpg": "image/jpeg",
|
||||
}
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj
|
||||
|
|
@ -171,12 +180,125 @@ def _apply_gemini_metadata(
|
|||
return cast(PartType, part_dict)
|
||||
|
||||
|
||||
def _parse_gs_uri(gs_uri: str) -> Tuple[str, str]:
|
||||
if not gs_uri.startswith("gs://"):
|
||||
raise ValueError(f"Invalid gs URI: {gs_uri}")
|
||||
uri_without_scheme = gs_uri[5:] # drop gs://
|
||||
uri_parts = uri_without_scheme.split("/", 1)
|
||||
if len(uri_parts) != 2 or not uri_parts[0] or not uri_parts[1]:
|
||||
raise ValueError(f"Invalid gs URI: {gs_uri}")
|
||||
return uri_parts[0], uri_parts[1]
|
||||
|
||||
|
||||
def _is_valid_gcs_bucket_name(bucket: str) -> bool:
|
||||
"""
|
||||
Validate bucket name against core GCS naming constraints.
|
||||
"""
|
||||
bucket_length = len(bucket)
|
||||
max_bucket_length = 222 if "." in bucket else 63
|
||||
if bucket_length < 3 or bucket_length > max_bucket_length:
|
||||
return False
|
||||
if "." in bucket and any(
|
||||
len(label) == 0 or len(label) > 63 for label in bucket.split(".")
|
||||
):
|
||||
return False
|
||||
if not re.fullmatch(r"[a-z0-9][a-z0-9._-]*[a-z0-9]", bucket):
|
||||
return False
|
||||
if ".." in bucket:
|
||||
return False
|
||||
if re.fullmatch(r"\d+\.\d+\.\d+\.\d+", bucket):
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
def _get_gcs_object_content_type(
|
||||
image_url: str,
|
||||
vertex_project: Optional[str] = None,
|
||||
vertex_credentials: Optional[Any] = None,
|
||||
) -> Optional[str]:
|
||||
"""
|
||||
Resolve content type from GCS object metadata.
|
||||
"""
|
||||
try:
|
||||
bucket, object_name = _parse_gs_uri(image_url)
|
||||
except ValueError:
|
||||
return None
|
||||
if not _is_valid_gcs_bucket_name(bucket):
|
||||
return None
|
||||
|
||||
headers: Dict[str, str] = {}
|
||||
explicit_vertex_auth_provided = (
|
||||
vertex_project is not None or vertex_credentials is not None
|
||||
)
|
||||
try:
|
||||
access_token, _ = _GCS_METADATA_VERTEX_BASE.get_access_token(
|
||||
credentials=vertex_credentials,
|
||||
project_id=vertex_project,
|
||||
)
|
||||
headers["Authorization"] = f"Bearer {access_token}"
|
||||
except Exception as e:
|
||||
if explicit_vertex_auth_provided:
|
||||
raise litellm.BadRequestError(
|
||||
message=(
|
||||
"Unable to fetch GCS metadata with provided Vertex credentials/project. "
|
||||
f"Original error: {str(e)}"
|
||||
),
|
||||
model=None,
|
||||
llm_provider="vertex_ai",
|
||||
)
|
||||
# Metadata may still be readable for public objects.
|
||||
pass
|
||||
|
||||
object_path = quote(object_name, safe="")
|
||||
metadata_url = (
|
||||
f"https://storage.googleapis.com/storage/v1/b/{bucket}/o/{object_path}"
|
||||
"?fields=contentType"
|
||||
)
|
||||
try:
|
||||
response = httpx.get(url=metadata_url, headers=headers, timeout=5.0)
|
||||
response.raise_for_status()
|
||||
content_type = response.json().get("contentType")
|
||||
if isinstance(content_type, str) and len(content_type) > 0:
|
||||
return content_type
|
||||
except Exception:
|
||||
return None
|
||||
return None
|
||||
|
||||
|
||||
def _normalize_and_validate_gemini_mime_type(
|
||||
mime_type: str, model: Optional[str]
|
||||
) -> str:
|
||||
normalized_mime_type = _GEMINI_MIME_TYPE_ALIASES.get(
|
||||
mime_type.strip().lower(), mime_type.strip().lower()
|
||||
)
|
||||
try:
|
||||
file_extension = get_file_extension_from_mime_type(normalized_mime_type)
|
||||
file_type = get_file_type_from_extension(file_extension)
|
||||
except ValueError:
|
||||
raise litellm.BadRequestError(
|
||||
message=f"File type not supported by gemini - {normalized_mime_type}",
|
||||
model=model,
|
||||
llm_provider="vertex_ai",
|
||||
)
|
||||
|
||||
if not is_gemini_1_5_accepted_file_type(file_type):
|
||||
raise litellm.BadRequestError(
|
||||
message=f"File type not supported by gemini - {file_type}",
|
||||
model=model,
|
||||
llm_provider="vertex_ai",
|
||||
)
|
||||
|
||||
return get_file_mime_type_for_file_type(file_type)
|
||||
|
||||
|
||||
def _process_gemini_media(
|
||||
image_url: str,
|
||||
format: Optional[str] = None,
|
||||
media_resolution_enum: Optional[Dict[str, str]] = None,
|
||||
model: Optional[str] = None,
|
||||
video_metadata: Optional[Dict[str, Any]] = None,
|
||||
vertex_project: Optional[str] = None,
|
||||
vertex_credentials: Optional[Any] = None,
|
||||
) -> PartType:
|
||||
"""
|
||||
Given a media URL (image, audio, or video), return the appropriate PartType for Gemini
|
||||
|
|
@ -193,20 +315,55 @@ def _process_gemini_media(
|
|||
try:
|
||||
# GCS URIs
|
||||
if "gs://" in image_url:
|
||||
# Figure out file type
|
||||
extension_with_dot = os.path.splitext(image_url)[-1] # Ex: ".png"
|
||||
extension = extension_with_dot[1:] # Ex: "png"
|
||||
|
||||
if not format:
|
||||
file_type = get_file_type_from_extension(extension)
|
||||
mime_type: Optional[str] = None
|
||||
# For extension-less gs:// URIs, we cannot infer from path.
|
||||
# If callers pass `format`/`mime_type`, this branch is skipped.
|
||||
if extension:
|
||||
file_type = get_file_type_from_extension(extension)
|
||||
|
||||
# Validate the file type is supported by Gemini
|
||||
if not is_gemini_1_5_accepted_file_type(file_type):
|
||||
raise Exception(f"File type not supported by gemini - {file_type}")
|
||||
# Validate the file type is supported by Gemini
|
||||
if not is_gemini_1_5_accepted_file_type(file_type):
|
||||
raise litellm.BadRequestError(
|
||||
message=f"File type not supported by gemini - {file_type}",
|
||||
model=model,
|
||||
llm_provider="vertex_ai",
|
||||
)
|
||||
|
||||
mime_type = get_file_mime_type_for_file_type(file_type)
|
||||
mime_type = get_file_mime_type_for_file_type(file_type)
|
||||
else:
|
||||
mime_type = _get_gcs_object_content_type(
|
||||
image_url=image_url,
|
||||
vertex_project=vertex_project,
|
||||
vertex_credentials=vertex_credentials,
|
||||
)
|
||||
if mime_type is None:
|
||||
raise litellm.BadRequestError(
|
||||
message=(
|
||||
f"Unable to determine mime type for gs URI: {image_url}. "
|
||||
"This gs:// URI has no file extension and GCS metadata "
|
||||
"lookup failed. Set it explicitly using image_url.format "
|
||||
"(or image_url.mime_type/content_type) or "
|
||||
"message.content[].file.format."
|
||||
),
|
||||
model=model,
|
||||
llm_provider="vertex_ai",
|
||||
)
|
||||
else:
|
||||
mime_type = format
|
||||
if mime_type is None:
|
||||
raise litellm.BadRequestError(
|
||||
message=f"File type not supported by gemini - {image_url}",
|
||||
model=model,
|
||||
llm_provider="vertex_ai",
|
||||
)
|
||||
mime_type = _normalize_and_validate_gemini_mime_type(
|
||||
mime_type=mime_type,
|
||||
model=model,
|
||||
)
|
||||
file_data = FileDataType(mime_type=mime_type, file_uri=image_url)
|
||||
part: PartType = {"file_data": file_data}
|
||||
return _apply_gemini_metadata(
|
||||
|
|
@ -258,8 +415,6 @@ def _snake_to_camel(snake_str: str) -> str:
|
|||
|
||||
def _camel_to_snake(camel_str: str) -> str:
|
||||
"""Convert camelCase to snake_case"""
|
||||
import re
|
||||
|
||||
return re.sub(r"(?<!^)(?=[A-Z])", "_", camel_str).lower()
|
||||
|
||||
|
||||
|
|
@ -311,6 +466,7 @@ def check_if_part_exists_in_parts(
|
|||
def _gemini_convert_messages_with_history( # noqa: PLR0915
|
||||
messages: List[AllMessageValues],
|
||||
model: Optional[str] = None,
|
||||
litellm_params: Optional[dict] = None,
|
||||
) -> List[ContentType]:
|
||||
"""
|
||||
Converts given messages from OpenAI format to Gemini format
|
||||
|
|
@ -326,6 +482,16 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915
|
|||
|
||||
msg_i = 0
|
||||
tool_call_responses = []
|
||||
vertex_project = None
|
||||
vertex_credentials = None
|
||||
if litellm_params:
|
||||
vertex_project = litellm_params.get("vertex_project") or litellm_params.get(
|
||||
"vertex_ai_project"
|
||||
)
|
||||
vertex_credentials = litellm_params.get(
|
||||
"vertex_credentials"
|
||||
) or litellm_params.get("vertex_ai_credentials")
|
||||
|
||||
try:
|
||||
while msg_i < len(messages):
|
||||
user_content: List[PartType] = []
|
||||
|
|
@ -366,7 +532,11 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915
|
|||
model=model,
|
||||
llm_provider="vertex_ai",
|
||||
)
|
||||
format = raw_image_url.get("format")
|
||||
format = (
|
||||
raw_image_url.get("format")
|
||||
or raw_image_url.get("mime_type")
|
||||
or raw_image_url.get("content_type")
|
||||
)
|
||||
detail = raw_image_url.get("detail")
|
||||
media_resolution_enum = (
|
||||
_convert_detail_to_media_resolution_enum(detail)
|
||||
|
|
@ -378,6 +548,8 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915
|
|||
format=format,
|
||||
media_resolution_enum=media_resolution_enum,
|
||||
model=model,
|
||||
vertex_project=vertex_project,
|
||||
vertex_credentials=vertex_credentials,
|
||||
)
|
||||
_parts.append(_part)
|
||||
elif element["type"] == "input_audio":
|
||||
|
|
@ -403,6 +575,8 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915
|
|||
image_url=openai_image_str,
|
||||
format=audio_format_modified,
|
||||
model=model,
|
||||
vertex_project=vertex_project,
|
||||
vertex_credentials=vertex_credentials,
|
||||
)
|
||||
_parts.append(_part)
|
||||
elif element["type"] == "file":
|
||||
|
|
@ -414,11 +588,16 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915
|
|||
model=model,
|
||||
llm_provider="vertex_ai",
|
||||
)
|
||||
file_id = _file_field.get("file_id")
|
||||
format = _file_field.get("format")
|
||||
file_data = _file_field.get("file_data")
|
||||
detail = _file_field.get("detail")
|
||||
video_metadata = _file_field.get("video_metadata")
|
||||
file_dict = cast(Dict[str, Any], _file_field)
|
||||
file_id = file_dict.get("file_id")
|
||||
format = (
|
||||
file_dict.get("format")
|
||||
or file_dict.get("mime_type")
|
||||
or file_dict.get("content_type")
|
||||
)
|
||||
file_data = file_dict.get("file_data")
|
||||
detail = file_dict.get("detail")
|
||||
video_metadata = file_dict.get("video_metadata")
|
||||
passed_file = file_id or file_data
|
||||
if passed_file is None:
|
||||
raise Exception(
|
||||
|
|
@ -437,13 +616,23 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915
|
|||
model=model,
|
||||
media_resolution_enum=media_resolution_enum,
|
||||
video_metadata=video_metadata,
|
||||
vertex_project=vertex_project,
|
||||
vertex_credentials=vertex_credentials,
|
||||
)
|
||||
_parts.append(_part)
|
||||
except Exception:
|
||||
raise Exception(
|
||||
"Unable to determine mime type for file_id: {}, set this explicitly using message[{}].content[{}].file.format".format(
|
||||
file_id, msg_i, element_idx
|
||||
)
|
||||
except litellm.BadRequestError:
|
||||
raise
|
||||
except Exception as e:
|
||||
raise litellm.BadRequestError(
|
||||
message=(
|
||||
"Unable to determine mime type for file: "
|
||||
f"{file_id or 'provided data'}, set this explicitly "
|
||||
f"using message[{msg_i}].content[{element_idx}]."
|
||||
"file.format (or file.mime_type/content_type). "
|
||||
f"Original error: {str(e)}"
|
||||
),
|
||||
model=model,
|
||||
llm_provider="vertex_ai",
|
||||
)
|
||||
user_content.extend(_parts)
|
||||
elif _message_content is not None and isinstance(_message_content, str):
|
||||
|
|
@ -559,6 +748,8 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915
|
|||
format=format,
|
||||
media_resolution_enum=media_resolution_enum,
|
||||
model=model,
|
||||
vertex_project=vertex_project,
|
||||
vertex_credentials=vertex_credentials,
|
||||
)
|
||||
assistant_content.append(_part)
|
||||
|
||||
|
|
@ -733,11 +924,11 @@ def _transform_request_body( # noqa: PLR0915
|
|||
try:
|
||||
if custom_llm_provider == "gemini":
|
||||
content = litellm.GoogleAIStudioGeminiConfig()._transform_messages(
|
||||
messages=messages, model=model
|
||||
messages=messages, model=model, litellm_params=litellm_params
|
||||
)
|
||||
else:
|
||||
content = litellm.VertexGeminiConfig()._transform_messages(
|
||||
messages=messages, model=model
|
||||
messages=messages, model=model, litellm_params=litellm_params
|
||||
)
|
||||
tools: Optional[Tools] = optional_params.pop("tools", None)
|
||||
tool_choice: Optional[ToolConfig] = optional_params.pop("tool_choice", None)
|
||||
|
|
|
|||
|
|
@ -2533,9 +2533,14 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
|
|||
return model_response
|
||||
|
||||
def _transform_messages(
|
||||
self, messages: List[AllMessageValues], model: Optional[str] = None
|
||||
self,
|
||||
messages: List[AllMessageValues],
|
||||
model: Optional[str] = None,
|
||||
litellm_params: Optional[dict] = None,
|
||||
) -> List[ContentType]:
|
||||
return _gemini_convert_messages_with_history(messages=messages, model=model)
|
||||
return _gemini_convert_messages_with_history(
|
||||
messages=messages, model=model, litellm_params=litellm_params
|
||||
)
|
||||
|
||||
def get_error_class(
|
||||
self, error_message: str, status_code: int, headers: Union[Dict, httpx.Headers]
|
||||
|
|
|
|||
|
|
@ -1219,6 +1219,32 @@ def test_process_gemini_media():
|
|||
mime_type="image/jpeg", file_uri="gs://bucket/image"
|
||||
)
|
||||
|
||||
# Test gs url without extension using mime_type from image_url object
|
||||
image_message = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "image_url",
|
||||
"image_url": {
|
||||
"url": "gs://bucket/image-without-extension",
|
||||
"mime_type": "image/png",
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
]
|
||||
from litellm.llms.vertex_ai.gemini.transformation import (
|
||||
_gemini_convert_messages_with_history,
|
||||
)
|
||||
|
||||
converted = _gemini_convert_messages_with_history(
|
||||
messages=image_message, model="gemini-2.5-flash"
|
||||
)
|
||||
assert converted[0]["parts"][0]["file_data"] == FileDataType(
|
||||
mime_type="image/png", file_uri="gs://bucket/image-without-extension"
|
||||
)
|
||||
|
||||
# Test HTTPS JPG URL
|
||||
https_result = _process_gemini_media("https://example.com/image.jpg")
|
||||
print("https_result JPG", https_result)
|
||||
|
|
@ -1256,6 +1282,168 @@ def test_process_gemini_media():
|
|||
assert base64_result["inline_data"]["data"] == "/9j/4AAQSkZJRg..."
|
||||
|
||||
|
||||
def test_process_gemini_media_gcs_without_extension_raises_clear_error():
|
||||
with pytest.raises(litellm.BadRequestError) as exc_info:
|
||||
_process_gemini_media("gs://bucket/image-without-extension")
|
||||
|
||||
assert "Unable to determine mime type for gs URI" in str(exc_info.value)
|
||||
|
||||
|
||||
def test_process_gemini_media_gcs_without_extension_uses_gcs_metadata():
|
||||
from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_media
|
||||
from litellm.types.llms.vertex_ai import FileDataType
|
||||
|
||||
with patch(
|
||||
"litellm.llms.vertex_ai.gemini.transformation._get_gcs_object_content_type",
|
||||
return_value="image/jpeg",
|
||||
):
|
||||
result = _process_gemini_media("gs://bucket/image-without-extension")
|
||||
|
||||
assert result["file_data"] == FileDataType(
|
||||
mime_type="image/jpeg", file_uri="gs://bucket/image-without-extension"
|
||||
)
|
||||
|
||||
|
||||
def test_file_block_uses_mime_type_alias_for_extensionless_gcs():
|
||||
from litellm.llms.vertex_ai.gemini.transformation import (
|
||||
_gemini_convert_messages_with_history,
|
||||
)
|
||||
from litellm.types.llms.vertex_ai import FileDataType
|
||||
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "file",
|
||||
"file": {
|
||||
"file_id": "gs://bucket/no-extension-object",
|
||||
"mime_type": "application/pdf",
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
converted = _gemini_convert_messages_with_history(
|
||||
messages=messages, model="gemini-2.5-flash"
|
||||
)
|
||||
assert converted[0]["parts"][0]["file_data"] == FileDataType(
|
||||
mime_type="application/pdf", file_uri="gs://bucket/no-extension-object"
|
||||
)
|
||||
|
||||
|
||||
def test_real_file_id_without_extension_resolves_to_metadata_mime_type():
|
||||
from litellm.llms.vertex_ai.gemini.transformation import (
|
||||
_gemini_convert_messages_with_history,
|
||||
)
|
||||
from litellm.types.llms.vertex_ai import FileDataType
|
||||
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "file",
|
||||
"file": {
|
||||
"file_id": "gs://cdn.pagepop.top/gcs/input/source/18/1851d730e543cb91643835612508d7dd0e8d8c52015fe08b7429a7307c62076c",
|
||||
},
|
||||
}
|
||||
],
|
||||
}
|
||||
]
|
||||
|
||||
with patch(
|
||||
"litellm.llms.vertex_ai.gemini.transformation._get_gcs_object_content_type",
|
||||
return_value="image/jpeg",
|
||||
):
|
||||
converted = _gemini_convert_messages_with_history(
|
||||
messages=messages, model="gemini-2.5-flash"
|
||||
)
|
||||
|
||||
assert converted[0]["parts"][0]["file_data"] == FileDataType(
|
||||
mime_type="image/jpeg",
|
||||
file_uri="gs://cdn.pagepop.top/gcs/input/source/18/1851d730e543cb91643835612508d7dd0e8d8c52015fe08b7429a7307c62076c",
|
||||
)
|
||||
|
||||
|
||||
def test_dotted_bucket_name_up_to_222_chars_is_accepted():
|
||||
from litellm.llms.vertex_ai.gemini.transformation import _is_valid_gcs_bucket_name
|
||||
|
||||
dotted_bucket = ("a." * 110) + "aa" # total length = 222
|
||||
assert len(dotted_bucket) == 222
|
||||
assert _is_valid_gcs_bucket_name(dotted_bucket) is True
|
||||
|
||||
|
||||
def test_get_gcs_object_content_type_uses_shared_vertex_base_instance():
|
||||
from litellm.llms.vertex_ai.gemini import transformation as gemini_transformation
|
||||
|
||||
mock_vertex_base = MagicMock()
|
||||
mock_vertex_base.get_access_token.return_value = ("test-token", "test-project")
|
||||
mock_http_response = MagicMock()
|
||||
mock_http_response.json.return_value = {"contentType": "image/png"}
|
||||
mock_http_response.raise_for_status.return_value = None
|
||||
|
||||
with patch.object(
|
||||
gemini_transformation, "_GCS_METADATA_VERTEX_BASE", mock_vertex_base
|
||||
), patch("litellm.llms.vertex_ai.gemini.transformation.httpx.get") as mock_http_get:
|
||||
mock_http_get.return_value = mock_http_response
|
||||
content_type = gemini_transformation._get_gcs_object_content_type(
|
||||
image_url="gs://my-bucket/path/to/image-without-extension",
|
||||
vertex_project="project-123",
|
||||
vertex_credentials="credential-json",
|
||||
)
|
||||
|
||||
assert content_type == "image/png"
|
||||
mock_vertex_base.get_access_token.assert_called_once_with(
|
||||
credentials="credential-json",
|
||||
project_id="project-123",
|
||||
)
|
||||
|
||||
|
||||
def test_process_gemini_media_normalizes_metadata_mime_alias():
|
||||
from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_media
|
||||
from litellm.types.llms.vertex_ai import FileDataType
|
||||
|
||||
with patch(
|
||||
"litellm.llms.vertex_ai.gemini.transformation._get_gcs_object_content_type",
|
||||
return_value="image/jpg",
|
||||
):
|
||||
result = _process_gemini_media("gs://bucket/image-without-extension")
|
||||
|
||||
assert result["file_data"] == FileDataType(
|
||||
mime_type="image/jpeg", file_uri="gs://bucket/image-without-extension"
|
||||
)
|
||||
|
||||
|
||||
def test_process_gemini_media_rejects_unsupported_metadata_mime_type():
|
||||
from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_media
|
||||
|
||||
with patch(
|
||||
"litellm.llms.vertex_ai.gemini.transformation._get_gcs_object_content_type",
|
||||
return_value="application/octet-stream",
|
||||
):
|
||||
with pytest.raises(litellm.BadRequestError, match="File type not supported by gemini"):
|
||||
_process_gemini_media("gs://bucket/image-without-extension")
|
||||
|
||||
|
||||
def test_get_gcs_object_content_type_fails_fast_with_explicit_credentials():
|
||||
from litellm.llms.vertex_ai.gemini import transformation as gemini_transformation
|
||||
|
||||
mock_vertex_base = MagicMock()
|
||||
mock_vertex_base.get_access_token.side_effect = Exception("token failure")
|
||||
with patch.object(gemini_transformation, "_GCS_METADATA_VERTEX_BASE", mock_vertex_base):
|
||||
with pytest.raises(
|
||||
litellm.BadRequestError,
|
||||
match="Unable to fetch GCS metadata with provided Vertex credentials/project",
|
||||
):
|
||||
gemini_transformation._get_gcs_object_content_type(
|
||||
image_url="gs://my-bucket/path/to/image-without-extension",
|
||||
vertex_project="project-123",
|
||||
vertex_credentials="credential-json",
|
||||
)
|
||||
|
||||
|
||||
def test_get_image_mime_type_from_url():
|
||||
"""Test the _get_image_mime_type_from_url function for different image URLs"""
|
||||
from litellm.llms.vertex_ai.gemini.transformation import (
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue