mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-03 02:22:24 +00:00
feat(vertex_ai/gemini): implement shared HTTP handler for GCS metadata
- Introduced a shared HTTP handler for GCS JSON API metadata reads to ensure proxy/SSL settings from litellm's HTTP stack are applied. - Added new function `_get_gcs_metadata_http_handler` to manage the HTTP handler's lifecycle. - Enhanced tests to validate the new handler's integration and behavior, including checks for bucket name validation and ensuring no authorization headers are sent for anonymous requests.
This commit is contained in:
S0ngRu1
parent
c71b6fffa3
commit
e6e2108f59
2 changed files with 105 additions and 18 deletions
|
|
@ -64,11 +64,21 @@ from ..common_utils import (
|
|||
# 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
|
||||
|
|
@ -79,6 +89,13 @@ def _get_vertex_base() -> Any:
|
|||
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
|
||||
|
||||
|
|
@ -279,7 +296,10 @@ def _get_gcs_object_content_type(
|
|||
params={"fields": "contentType"},
|
||||
)
|
||||
try:
|
||||
response = httpx.get(url=metadata_url, headers=headers, timeout=5.0)
|
||||
response = _get_gcs_metadata_http_handler().get(
|
||||
url=str(metadata_url),
|
||||
headers=headers or None,
|
||||
)
|
||||
response.raise_for_status()
|
||||
content_type = response.json().get("contentType")
|
||||
if isinstance(content_type, str) and len(content_type) > 0:
|
||||
|
|
@ -296,9 +316,7 @@ def _normalize_and_validate_gemini_mime_type(
|
|||
# litellm.types.files.
|
||||
from litellm.types.files import get_file_extension_from_mime_type
|
||||
|
||||
normalized_mime_type = _GEMINI_MIME_TYPE_ALIASES.get(
|
||||
mime_type.strip().lower(), mime_type.strip().lower()
|
||||
)
|
||||
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)
|
||||
|
|
@ -346,6 +364,7 @@ def _process_gemini_media(
|
|||
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:
|
||||
mime_type: Optional[str] = None
|
||||
# For extension-less gs:// URIs, we cannot infer from path.
|
||||
|
|
@ -382,16 +401,23 @@ def _process_gemini_media(
|
|||
)
|
||||
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",
|
||||
)
|
||||
mime_type = _normalize_and_validate_gemini_mime_type(
|
||||
mime_type=mime_type,
|
||||
model=model,
|
||||
)
|
||||
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(
|
||||
|
|
|
|||
|
|
@ -1282,6 +1282,34 @@ def test_process_gemini_media():
|
|||
assert base64_result["inline_data"]["data"] == "/9j/4AAQSkZJRg..."
|
||||
|
||||
|
||||
def test_process_gemini_media_gcs_explicit_format_skips_litellm_registry():
|
||||
"""Explicit format for gs:// must not require litellm's FILE_MIME_TYPES entry."""
|
||||
from litellm.types.llms.vertex_ai import FileDataType
|
||||
|
||||
result = _process_gemini_media(
|
||||
"gs://bucket/object-no-ext",
|
||||
format="application/octet-stream",
|
||||
)
|
||||
assert result["file_data"] == FileDataType(
|
||||
mime_type="application/octet-stream",
|
||||
file_uri="gs://bucket/object-no-ext",
|
||||
)
|
||||
|
||||
|
||||
def test_process_gemini_media_gcs_explicit_format_still_applies_mime_alias():
|
||||
"""Known MIME aliases still apply when format is explicit for gs:// URIs."""
|
||||
from litellm.types.llms.vertex_ai import FileDataType
|
||||
|
||||
result = _process_gemini_media(
|
||||
"gs://bucket/object-no-ext",
|
||||
format="image/jpg",
|
||||
)
|
||||
assert result["file_data"] == FileDataType(
|
||||
mime_type="image/jpeg",
|
||||
file_uri="gs://bucket/object-no-ext",
|
||||
)
|
||||
|
||||
|
||||
def test_process_gemini_media_gcs_without_extension_raises_clear_error():
|
||||
# Mock the GCS metadata lookup to avoid real outbound HTTP in tests.
|
||||
with patch(
|
||||
|
|
@ -1380,6 +1408,26 @@ def test_dotted_bucket_name_up_to_222_chars_is_accepted():
|
|||
assert _is_valid_gcs_bucket_name(dotted_bucket) is True
|
||||
|
||||
|
||||
def test_is_valid_gcs_bucket_name_rejects_too_short():
|
||||
from litellm.llms.vertex_ai.gemini.transformation import _is_valid_gcs_bucket_name
|
||||
|
||||
assert _is_valid_gcs_bucket_name("ab") is False
|
||||
|
||||
|
||||
def test_is_valid_gcs_bucket_name_rejects_too_long_plain():
|
||||
from litellm.llms.vertex_ai.gemini.transformation import _is_valid_gcs_bucket_name
|
||||
|
||||
bucket = "a" * 64
|
||||
assert len(bucket) == 64
|
||||
assert _is_valid_gcs_bucket_name(bucket) is False
|
||||
|
||||
|
||||
def test_is_valid_gcs_bucket_name_rejects_double_dot():
|
||||
from litellm.llms.vertex_ai.gemini.transformation import _is_valid_gcs_bucket_name
|
||||
|
||||
assert _is_valid_gcs_bucket_name("ab..cd") is False
|
||||
|
||||
|
||||
def test_get_gcs_object_content_type_uses_shared_vertex_base_instance():
|
||||
from litellm.llms.vertex_ai.gemini import transformation as gemini_transformation
|
||||
|
||||
|
|
@ -1389,15 +1437,18 @@ def test_get_gcs_object_content_type_uses_shared_vertex_base_instance():
|
|||
mock_http_response.json.return_value = {"contentType": "image/png"}
|
||||
mock_http_response.raise_for_status.return_value = None
|
||||
|
||||
mock_http_handler = MagicMock()
|
||||
mock_http_handler.get.return_value = mock_http_response
|
||||
|
||||
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,
|
||||
"litellm.llms.vertex_ai.gemini.transformation._get_gcs_metadata_http_handler",
|
||||
return_value=mock_http_handler,
|
||||
),
|
||||
):
|
||||
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",
|
||||
|
|
@ -1472,15 +1523,18 @@ def test_get_gcs_object_content_type_without_credentials_skips_auth():
|
|||
mock_http_response.json.return_value = {"contentType": "image/jpeg"}
|
||||
mock_http_response.raise_for_status.return_value = None
|
||||
|
||||
mock_http_handler = MagicMock()
|
||||
mock_http_handler.get.return_value = mock_http_response
|
||||
|
||||
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,
|
||||
"litellm.llms.vertex_ai.gemini.transformation._get_gcs_metadata_http_handler",
|
||||
return_value=mock_http_handler,
|
||||
),
|
||||
):
|
||||
mock_http_get.return_value = mock_http_response
|
||||
content_type = gemini_transformation._get_gcs_object_content_type(
|
||||
image_url="gs://public-bucket/public-object"
|
||||
)
|
||||
|
|
@ -1488,9 +1542,10 @@ def test_get_gcs_object_content_type_without_credentials_skips_auth():
|
|||
# Must not call get_access_token (so default server credentials are not used)
|
||||
mock_vertex_base.get_access_token.assert_not_called()
|
||||
# An anonymous request is still sent, covering publicly-readable objects.
|
||||
mock_http_get.assert_called_once()
|
||||
call_kwargs = mock_http_get.call_args.kwargs
|
||||
assert "Authorization" not in call_kwargs.get("headers", {})
|
||||
mock_http_handler.get.assert_called_once()
|
||||
call_kwargs = mock_http_handler.get.call_args.kwargs
|
||||
headers = call_kwargs.get("headers")
|
||||
assert headers is None or "Authorization" not in headers
|
||||
assert content_type == "image/jpeg"
|
||||
|
||||
|
||||
|
|
@ -1559,11 +1614,17 @@ def test_async_transform_request_body_does_not_block_event_loop():
|
|||
await transform_task
|
||||
return sleep_elapsed
|
||||
|
||||
mock_http_handler = MagicMock()
|
||||
mock_http_handler.get.side_effect = slow_http_get
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
gemini_transformation, "_GCS_METADATA_VERTEX_BASE", mock_vertex_base
|
||||
),
|
||||
patch.object(gemini_transformation.httpx, "get", side_effect=slow_http_get),
|
||||
patch(
|
||||
"litellm.llms.vertex_ai.gemini.transformation._get_gcs_metadata_http_handler",
|
||||
return_value=mock_http_handler,
|
||||
),
|
||||
patch(
|
||||
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching."
|
||||
"ContextCachingEndpoints.async_check_and_create_cache",
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue