From ad22861996ad0818f554ca2834deaf301922ebc1 Mon Sep 17 00:00:00 2001 From: Kris Xia Date: Wed, 13 May 2026 22:35:36 +0800 Subject: [PATCH] Fix Gemini MIME detection for extensionless GCS URIs (#27278) Squash-merged by litellm-agent from krisxia0506's PR. --- litellm/llms/gemini/chat/transformation.py | 9 +- .../llms/vertex_ai/gemini/transformation.py | 474 +++++++++++++++++- .../vertex_and_google_ai_studio_gemini.py | 9 +- .../llms/vertex_ai/test_vertex.py | 27 + .../test_vertex_gemini_gcs_uri_mime.py | 466 +++++++++++++++++ 5 files changed, 957 insertions(+), 28 deletions(-) create mode 100644 tests/test_litellm/llms/vertex_ai/test_vertex_gemini_gcs_uri_mime.py diff --git a/litellm/llms/gemini/chat/transformation.py b/litellm/llms/gemini/chat/transformation.py index fb5239e61f6..16e17dcc876 100644 --- a/litellm/llms/gemini/chat/transformation.py +++ b/litellm/llms/gemini/chat/transformation.py @@ -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 + ) diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 5102ed6c483..f92e3ae4ca8 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -6,13 +6,16 @@ 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 import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.prompt_templates.common_utils import ( _get_image_mime_type_from_url, ) @@ -57,6 +60,42 @@ from ..common_utils import ( get_supports_system_message, ) +# Typed as Any to avoid introducing a module-load-time cyclic import to +# vertex_llm_base. The instance is lazily constructed by _get_vertex_base() +# the first time GCS metadata needs to be fetched. +_GCS_METADATA_VERTEX_BASE: Optional[Any] = None +# Shared sync client for GCS JSON API metadata reads so proxy/SSL settings +# from litellm's HTTP stack apply (see Greptile review on PR #27278). +_GCS_METADATA_HTTP_HANDLER: Optional[HTTPHandler] = None +_GEMINI_MIME_TYPE_ALIASES: Dict[str, str] = { + "image/jpg": "image/jpeg", +} + + +def _apply_gemini_mime_type_aliases(mime_type: str) -> str: + """Normalize known MIME aliases only; does not consult the file-type registry.""" + return _GEMINI_MIME_TYPE_ALIASES.get( + mime_type.strip().lower(), mime_type.strip().lower() + ) + + +def _get_vertex_base() -> Any: + """Lazily return the shared VertexBase instance to avoid a module-load-time cyclic import.""" + global _GCS_METADATA_VERTEX_BASE + if _GCS_METADATA_VERTEX_BASE is None: + from ..vertex_llm_base import VertexBase + + _GCS_METADATA_VERTEX_BASE = VertexBase() + return _GCS_METADATA_VERTEX_BASE + + +def _get_gcs_metadata_http_handler() -> HTTPHandler: + global _GCS_METADATA_HTTP_HANDLER + if _GCS_METADATA_HTTP_HANDLER is None: + _GCS_METADATA_HTTP_HANDLER = HTTPHandler(timeout=5.0) + return _GCS_METADATA_HTTP_HANDLER + + if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj @@ -171,12 +210,299 @@ 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 _gs_uri_requires_content_type_metadata(url: str) -> bool: + """ + True when _process_gemini_media would call _get_gcs_object_content_type + (extension-less gs:// and no explicit format passed into that helper). + """ + if "gs://" not in url: + return False + extension_with_dot = os.path.splitext(url)[-1] + extension = extension_with_dot[1:] if extension_with_dot else "" + return len(extension) == 0 + + +def _image_url_payload_may_need_sync_gcs_metadata_fetch( + raw_image_url: Any, +) -> bool: + """ + True when this image_url value (content-part image_url or assistant ``images[]`` + entry) can trigger a blocking GCS metadata read for MIME resolution. + """ + fmt: Optional[str] = None + url: Optional[str] = None + if isinstance(raw_image_url, dict): + url = raw_image_url.get("url") # type: ignore[assignment] + if not isinstance(url, str): + return False + fmt = ( + raw_image_url.get("format") + or raw_image_url.get("mime_type") + or raw_image_url.get("content_type") + ) + elif isinstance(raw_image_url, str): + url = raw_image_url + else: + return False + if "gs://" not in url or fmt: + return False + return _gs_uri_requires_content_type_metadata(url) + + +def _openai_messages_may_need_sync_gcs_metadata_fetch( + messages: List[AllMessageValues], +) -> bool: + """ + Heuristic: True if any message part can trigger a blocking GCS JSON + metadata read inside _transform_request_body (extension-less gs:// without + explicit MIME hints). Covers user/system ``content`` parts and assistant + ``images`` (same paths as ``_gemini_convert_messages_with_history``). Used + to decide whether ``async_transform_request_body`` should offload the sync + transform via ``asyncify``. + """ + for raw in messages: + msg: Any = raw + if not isinstance(msg, dict) and hasattr(msg, "model_dump"): + msg = msg.model_dump(exclude_none=False) + if not isinstance(msg, dict): + continue + images_field = msg.get("images") + if isinstance(images_field, list): + for image_item in images_field: + if not isinstance(image_item, dict): + continue + if _image_url_payload_may_need_sync_gcs_metadata_fetch( + image_item.get("image_url") + ): + return True + + content = msg.get("content") + if not isinstance(content, list): + continue + for item in content: + if not isinstance(item, dict): + continue + itype = item.get("type") + if itype == "image_url": + if _image_url_payload_may_need_sync_gcs_metadata_fetch( + item.get("image_url") + ): + return True + elif itype == "file": + file_obj = item.get("file") + if not isinstance(file_obj, dict): + continue + fmt = ( + file_obj.get("format") + or file_obj.get("mime_type") + or file_obj.get("content_type") + ) + passed = file_obj.get("file_id") or file_obj.get("file_data") + if ( + isinstance(passed, str) + and "gs://" in passed + and not fmt + and _gs_uri_requires_content_type_metadata(passed) + ): + return True + return False + + +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. + + Only attaches a Bearer token when the caller explicitly supplies Vertex + credentials, to avoid using the server's default Google credentials on + the Gemini API-key (Google AI Studio) path and being used as an oracle + for private GCS object metadata. Without explicit credentials we only + issue an anonymous request, which only succeeds for publicly-readable + objects. + """ + 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 + ) + if explicit_vertex_auth_provided: + try: + access_token, _ = _get_vertex_base().get_access_token( + credentials=vertex_credentials, + project_id=vertex_project, + ) + headers["Authorization"] = f"Bearer {access_token}" + except Exception as e: + 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", + ) + + # Build the URL via httpx.URL with a fixed scheme/host and URL-encode both + # bucket and object so CodeQL does not flag the interpolation as a + # potential SSRF that could resolve to an arbitrary host. + encoded_bucket = quote(bucket, safe="") + encoded_object = quote(object_name, safe="") + metadata_url = httpx.URL( + scheme="https", + host="storage.googleapis.com", + path=f"/storage/v1/b/{encoded_bucket}/o/{encoded_object}", + params={"fields": "contentType"}, + ) + try: + response = _get_gcs_metadata_http_handler().get( + url=str(metadata_url), + headers=headers or None, + ) + except httpx.RequestError as e: + if explicit_vertex_auth_provided: + raise litellm.BadRequestError( + message=( + "Unable to reach GCS JSON API for object metadata with provided " + f"Vertex credentials. {type(e).__name__}: {e}" + ), + model=None, + llm_provider="vertex_ai", + ) from e + return None + + if response.is_error: + if explicit_vertex_auth_provided: + preview = (response.text or "")[:1024] + raise litellm.BadRequestError( + message=( + "Unable to read GCS object metadata with provided Vertex credentials. " + f"HTTP {response.status_code}. Response body (truncated): {preview!r}" + ), + model=None, + llm_provider="vertex_ai", + ) + return None + + try: + payload = response.json() + except ValueError as e: + if explicit_vertex_auth_provided: + raise litellm.BadRequestError( + message=( + "GCS metadata response was not valid JSON when using provided " + f"Vertex credentials (HTTP {response.status_code}). Error: {e}" + ), + model=None, + llm_provider="vertex_ai", + ) from e + return None + + if not isinstance(payload, dict): + if explicit_vertex_auth_provided: + raise litellm.BadRequestError( + message=( + "GCS metadata response was not a JSON object when using provided " + f"Vertex credentials (HTTP {response.status_code})." + ), + model=None, + llm_provider="vertex_ai", + ) + return None + + content_type = payload.get("contentType") + if isinstance(content_type, str) and len(content_type) > 0: + return content_type + + if explicit_vertex_auth_provided: + preview = (response.text or "")[:1024] + raise litellm.BadRequestError( + message=( + "GCS metadata JSON did not include a non-empty contentType field when " + f"using provided Vertex credentials (HTTP {response.status_code}). " + f"Body (truncated): {preview!r}" + ), + model=None, + llm_provider="vertex_ai", + ) + return None + + +def _normalize_and_validate_gemini_mime_type( + mime_type: str, model: Optional[str] +) -> str: + # Import lazily to avoid a module-level cyclic-import alert with + # litellm.types.files. + from litellm.types.files import get_file_extension_from_mime_type + + normalized_mime_type = _apply_gemini_mime_type_aliases(mime_type) + 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 +519,63 @@ 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" + explicit_gcs_format = False 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 + explicit_gcs_format = True + if mime_type is None: + raise litellm.BadRequestError( + message=f"File type not supported by gemini - {image_url}", + model=model, + llm_provider="vertex_ai", + ) + if explicit_gcs_format: + # Callers who pass format/mime_type explicitly for gs:// URIs + # rely on pass-through to Gemini (pre-PR behavior). Only apply + # known MIME aliases; skip litellm's file-type registry. + mime_type = _apply_gemini_mime_type_aliases(mime_type) + else: + 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 +627,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"(? List[ContentType]: """ Converts given messages from OpenAI format to Gemini format @@ -326,6 +694,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,8 +744,15 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915 model=model, llm_provider="vertex_ai", ) - format = raw_image_url.get("format") - detail = raw_image_url.get("detail") + # TypedDict does not declare mime_type/content_type; + # read via Dict[str, Any] for caller-provided MIME fields. + image_url_dict = cast(Dict[str, Any], raw_image_url) + format = ( + image_url_dict.get("format") + or image_url_dict.get("mime_type") + or image_url_dict.get("content_type") + ) + detail = image_url_dict.get("detail") media_resolution_enum = ( _convert_detail_to_media_resolution_enum(detail) ) @@ -378,6 +763,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 +790,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 +803,18 @@ 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") + # TypedDict does not declare mime_type/content_type; + # read via Dict[str, Any] for caller-provided MIME fields. + 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 +833,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): @@ -548,7 +954,11 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915 image_url_obj = image_item.get("image_url") if isinstance(image_url_obj, dict): assistant_image_url = image_url_obj.get("url") - format = image_url_obj.get("format") + format = ( + image_url_obj.get("format") + or image_url_obj.get("mime_type") + or image_url_obj.get("content_type") + ) detail = image_url_obj.get("detail") media_resolution_enum = ( _convert_detail_to_media_resolution_enum(detail) @@ -559,6 +969,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 +1145,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) @@ -913,6 +1325,20 @@ async def async_transform_request_body( vertex_auth_header=vertex_auth_header, ) + if _openai_messages_may_need_sync_gcs_metadata_fetch(messages): + # _transform_request_body may issue a sync httpx.get (up to 5s timeout) + # via _get_gcs_object_content_type to fetch GCS object metadata. Run the + # whole sync transformation on a worker thread so it does not block the + # async event loop. + return await asyncify(_transform_request_body)( + messages=messages, + model=model, + custom_llm_provider=custom_llm_provider, + litellm_params=litellm_params, + cached_content=cached_content, + optional_params=optional_params, + ) + return _transform_request_body( messages=messages, model=model, diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index f9899197854..49c1c335467 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -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] diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex.py b/tests/test_litellm/llms/vertex_ai/test_vertex.py index 2e9629f95de..be0e59e8b7d 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex.py @@ -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,7 @@ def test_process_gemini_media(): assert base64_result["inline_data"]["data"] == "/9j/4AAQSkZJRg..." + 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 ( diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_gemini_gcs_uri_mime.py b/tests/test_litellm/llms/vertex_ai/test_vertex_gemini_gcs_uri_mime.py new file mode 100644 index 00000000000..e0eccad80e2 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_gemini_gcs_uri_mime.py @@ -0,0 +1,466 @@ +"""Vertex Gemini: extensionless gs:// MIME + GCS metadata tests. + +Split from test_vertex.py to satisfy CI per-file size limits. +""" +import asyncio +import os +import sys +import time + +from dotenv import load_dotenv + +load_dotenv() + +import pytest + +import litellm +from unittest.mock import MagicMock, patch + +sys.path.insert(0, os.path.abspath("../..")) + +from litellm.llms.vertex_ai.gemini.transformation import _process_gemini_media + + +def test_process_gemini_media_gcs_explicit_format_octet_stream_and_alias(): + """Explicit format bypasses registry; image/jpg alias still applies.""" + from litellm.types.llms.vertex_ai import FileDataType + + r1 = _process_gemini_media( + "gs://bucket/object-no-ext", + format="application/octet-stream", + ) + assert r1["file_data"] == FileDataType( + mime_type="application/octet-stream", + file_uri="gs://bucket/object-no-ext", + ) + r2 = _process_gemini_media("gs://bucket/object-no-ext", format="image/jpg") + assert r2["file_data"] == FileDataType( + mime_type="image/jpeg", + file_uri="gs://bucket/object-no-ext", + ) + + +def test_process_gemini_media_gcs_without_extension_errors_and_metadata_mock(): + with patch( + "litellm.llms.vertex_ai.gemini.transformation._get_gcs_object_content_type", + return_value=None, + ): + with pytest.raises(litellm.BadRequestError) as exc: + _process_gemini_media("gs://bucket/image-without-extension") + assert "Unable to determine mime type for gs URI" in str(exc.value) + + 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", + ) as m: + r = _process_gemini_media("gs://bucket/image-without-extension") + assert r["file_data"] == FileDataType( + mime_type="image/jpeg", file_uri="gs://bucket/image-without-extension" + ) + m.assert_called() + + with patch( + "litellm.llms.vertex_ai.gemini.transformation._get_gcs_object_content_type", + return_value="image/jpg", + ): + r_alias = _process_gemini_media("gs://bucket/image-without-extension") + assert r_alias["file_data"]["mime_type"] == "image/jpeg" + + +def test_process_gemini_media_rejects_gcs_metadata_mime_not_supported_by_gemini(): + """Non-empty GCS contentType that fails _normalize_and_validate_gemini_mime_type.""" + with patch( + "litellm.llms.vertex_ai.gemini.transformation._get_gcs_object_content_type", + return_value="application/x-litellm-unit-test-unknown-mime", + ): + with pytest.raises( + litellm.BadRequestError, + match="File type not supported by gemini", + ): + _process_gemini_media("gs://bucket/object-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" + ) + + +@pytest.mark.parametrize( + "bucket,expected", + [ + (("a." * 110) + "aa", True), + ("ab", False), + ("a" * 64, False), + ("ab..cd", False), + ("1.2.3.4", False), + ("192.168.0.1", False), + ("Bucket-Upper", False), + ("bucket@name", False), + ("bucket name", False), + ("-mybucket", False), + ("mybucket-", False), + (".mybucket", False), + ("mybucket.", False), + ], +) +def test_is_valid_gcs_bucket_name_matrix(bucket, expected): + from litellm.llms.vertex_ai.gemini.transformation import _is_valid_gcs_bucket_name + + assert _is_valid_gcs_bucket_name(bucket) is expected + + +def test_get_gcs_object_content_type_explicit_vertex_success_and_token_failure(): + from litellm.llms.vertex_ai.gemini import transformation as gt + + mock_v = MagicMock() + mock_v.get_access_token.return_value = ("test-token", "test-project") + resp = MagicMock() + resp.is_error = False + resp.status_code = 200 + resp.json.return_value = {"contentType": "image/png"} + http = MagicMock() + http.get.return_value = resp + + with ( + patch.object(gt, "_GCS_METADATA_VERTEX_BASE", mock_v), + patch( + "litellm.llms.vertex_ai.gemini.transformation._get_gcs_metadata_http_handler", + return_value=http, + ), + ): + assert ( + gt._get_gcs_object_content_type( + image_url="gs://my-bucket/path/to/image-without-extension", + vertex_project="project-123", + vertex_credentials="credential-json", + ) + == "image/png" + ) + mock_v.get_access_token.assert_called_once_with( + credentials="credential-json", + project_id="project-123", + ) + + mock_v2 = MagicMock() + mock_v2.get_access_token.side_effect = Exception("token failure") + with patch.object(gt, "_GCS_METADATA_VERTEX_BASE", mock_v2): + with pytest.raises( + litellm.BadRequestError, + match="Unable to fetch GCS metadata with provided Vertex credentials/project", + ): + gt._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_gcs_object_content_type_http_error_explicit_vs_anonymous(): + from litellm.llms.vertex_ai.gemini import transformation as gt + + mock_v = MagicMock() + mock_v.get_access_token.return_value = ("t", "p") + err_resp = MagicMock() + err_resp.is_error = True + err_resp.status_code = 403 + err_resp.text = '{"error":{"message":"Permission denied"}}' + http = MagicMock() + http.get.return_value = err_resp + + with ( + patch.object(gt, "_GCS_METADATA_VERTEX_BASE", mock_v), + patch( + "litellm.llms.vertex_ai.gemini.transformation._get_gcs_metadata_http_handler", + return_value=http, + ), + ): + with pytest.raises(litellm.BadRequestError, match="HTTP 403") as ei: + gt._get_gcs_object_content_type( + image_url="gs://my-bucket/path/to/obj", + vertex_project="project-123", + vertex_credentials="credential-json", + ) + assert "Permission denied" in str(ei.value) + + mock_v2 = MagicMock() + anon_err = MagicMock() + anon_err.is_error = True + anon_err.status_code = 403 + anon_err.text = "Forbidden" + http2 = MagicMock() + http2.get.return_value = anon_err + with ( + patch.object(gt, "_GCS_METADATA_VERTEX_BASE", mock_v2), + patch( + "litellm.llms.vertex_ai.gemini.transformation._get_gcs_metadata_http_handler", + return_value=http2, + ), + ): + assert ( + gt._get_gcs_object_content_type(image_url="gs://public-bucket/public-object") + is None + ) + mock_v2.get_access_token.assert_not_called() + + +def test_get_gcs_object_content_type_anonymous_success_no_auth_header(): + from litellm.llms.vertex_ai.gemini import transformation as gt + + mock_v = MagicMock() + ok = MagicMock() + ok.is_error = False + ok.status_code = 200 + ok.json.return_value = {"contentType": "image/jpeg"} + http = MagicMock() + http.get.return_value = ok + + with ( + patch.object(gt, "_GCS_METADATA_VERTEX_BASE", mock_v), + patch( + "litellm.llms.vertex_ai.gemini.transformation._get_gcs_metadata_http_handler", + return_value=http, + ), + ): + assert ( + gt._get_gcs_object_content_type(image_url="gs://public-bucket/public-object") + == "image/jpeg" + ) + mock_v.get_access_token.assert_not_called() + hdrs = http.get.call_args.kwargs.get("headers") + assert hdrs is None or "Authorization" not in hdrs + + +def test_async_transform_request_body_offloads_extensionless_gs_not_plain_text(): + from litellm.llms.vertex_ai.gemini import transformation as gemini_transformation + + messages = [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": "gs://bucket/image-without-extension"}, + } + ], + } + ] + + def slow_http_get(*args, **kwargs): + time.sleep(0.5) + response = MagicMock() + response.is_error = False + response.status_code = 200 + response.raise_for_status.return_value = None + response.json.return_value = {"contentType": "image/png"} + return response + + async def fake_check_and_create_cache(self, **kwargs): + return kwargs["messages"], kwargs["optional_params"], None + + mock_v = MagicMock() + mock_v.get_access_token.return_value = ("token", "project") + mock_http = MagicMock() + mock_http.get.side_effect = slow_http_get + + async def run_scenario() -> float: + async def concurrent_sleep() -> float: + start = time.monotonic() + await asyncio.sleep(0.05) + return time.monotonic() - start + + task = asyncio.create_task( + gemini_transformation.async_transform_request_body( + gemini_api_key=None, + messages=messages, + api_base=None, + model="gemini-2.5-flash", + client=None, + timeout=None, + extra_headers=None, + optional_params={}, + logging_obj=MagicMock(), + custom_llm_provider="vertex_ai", + litellm_params={}, + vertex_project=None, + vertex_location=None, + vertex_auth_header=None, + ) + ) + elapsed = await concurrent_sleep() + await task + return elapsed + + with ( + patch.object(gemini_transformation, "_GCS_METADATA_VERTEX_BASE", mock_v), + patch( + "litellm.llms.vertex_ai.gemini.transformation._get_gcs_metadata_http_handler", + return_value=mock_http, + ), + patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching." + "ContextCachingEndpoints.async_check_and_create_cache", + new=fake_check_and_create_cache, + ), + ): + sleep_elapsed = asyncio.run(run_scenario()) + + assert sleep_elapsed < 0.4, ( + f"Event loop blocked for {sleep_elapsed:.3f}s; " + "async_transform_request_body did not offload sync GCS metadata" + ) + + async def fake_cache2(self, **kwargs): + return kwargs["messages"], kwargs["optional_params"], None + + async def run_plain(): + with patch( + "litellm.llms.vertex_ai.gemini.transformation.asyncify", + side_effect=AssertionError("asyncify must not run without extensionless gs://"), + ): + return await gemini_transformation.async_transform_request_body( + gemini_api_key=None, + messages=[{"role": "user", "content": "hello"}], + api_base=None, + model="gemini-2.5-flash", + client=None, + timeout=None, + extra_headers=None, + optional_params={}, + logging_obj=MagicMock(), + custom_llm_provider="vertex_ai", + litellm_params={}, + vertex_project=None, + vertex_location=None, + vertex_auth_header=None, + ) + + with patch( + "litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching." + "ContextCachingEndpoints.async_check_and_create_cache", + new=fake_cache2, + ): + body = asyncio.run(run_plain()) + assert body is not None and "contents" in body + + +@pytest.mark.parametrize( + "messages,expected", + [ + ([{"role": "user", "content": "hello"}], False), + ( + [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": "gs://bucket/image-without-extension"}, + } + ], + } + ], + True, + ), + ( + [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": {"url": "gs://bucket/image.png"}, + } + ], + } + ], + False, + ), + ( + [ + { + "role": "user", + "content": [ + { + "type": "image_url", + "image_url": { + "url": "gs://bucket/image-without-extension", + "mime_type": "image/png", + }, + } + ], + } + ], + False, + ), + ( + [ + { + "role": "assistant", + "content": [], + "images": [ + {"image_url": {"url": "gs://bucket/gen-without-extension"}}, + ], + } + ], + True, + ), + ( + [ + { + "role": "assistant", + "content": [], + "images": [{"image_url": {"url": "gs://bucket/gen.png"}}], + } + ], + False, + ), + ( + [ + { + "role": "assistant", + "content": [], + "images": [ + { + "image_url": { + "url": "gs://bucket/gen-no-ext", + "mime_type": "image/png", + }, + } + ], + } + ], + False, + ), + ], +) +def test_openai_messages_may_need_sync_gcs_metadata_fetch_matrix(messages, expected): + from litellm.llms.vertex_ai.gemini.transformation import ( + _openai_messages_may_need_sync_gcs_metadata_fetch, + ) + + assert _openai_messages_may_need_sync_gcs_metadata_fetch(messages) is expected