harden cloud file compatibility path

This commit is contained in:
user 2026-05-01 18:14:36 -07:00
parent 2776e0162c
commit e7c23fc041
8 changed files with 94 additions and 32 deletions

View file

@ -86,6 +86,16 @@ bedrock_files_instance = BedrockFilesHandler()
#################################################
def _add_trusted_model_credentials_to_litellm_params(
litellm_params_dict: Dict[str, Any], kwargs: Dict[str, Any]
) -> None:
trusted_model_credentials = kwargs.get("_litellm_internal_model_credentials")
if isinstance(trusted_model_credentials, type(MappingProxyType({}))):
litellm_params_dict["_litellm_internal_model_credentials"] = (
trusted_model_credentials
)
@client
async def acreate_file(
file: FileTypes,
@ -374,6 +384,10 @@ def file_retrieve(
)
if provider_config is not None:
litellm_params_dict = get_litellm_params(**kwargs)
_add_trusted_model_credentials_to_litellm_params(
litellm_params_dict=litellm_params_dict,
kwargs=kwargs,
)
litellm_params_dict["api_key"] = optional_params.api_key
litellm_params_dict["api_base"] = optional_params.api_base
@ -498,6 +512,10 @@ def file_delete(
pass
optional_params = GenericLiteLLMParams(**kwargs)
litellm_params_dict = get_litellm_params(**kwargs)
_add_trusted_model_credentials_to_litellm_params(
litellm_params_dict=litellm_params_dict,
kwargs=kwargs,
)
### TIMEOUT LOGIC ###
timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600
# set timeout for 10 minutes by default
@ -847,11 +865,10 @@ def file_content(
try:
optional_params = GenericLiteLLMParams(**kwargs)
litellm_params_dict = get_litellm_params(**kwargs)
trusted_model_credentials = kwargs.get("_litellm_internal_model_credentials")
if isinstance(trusted_model_credentials, type(MappingProxyType({}))):
litellm_params_dict["_litellm_internal_model_credentials"] = (
trusted_model_credentials
)
_add_trusted_model_credentials_to_litellm_params(
litellm_params_dict=litellm_params_dict,
kwargs=kwargs,
)
### TIMEOUT LOGIC ###
timeout = optional_params.timeout or kwargs.get("request_timeout", 600) or 600
client = kwargs.get("client")

View file

@ -1,7 +1,7 @@
import os
import posixpath
import re
from typing import Optional, Sequence, Tuple
from types import MappingProxyType
from typing import Any, Mapping, Optional, Sequence, Tuple, cast
from urllib.parse import quote, unquote
from litellm._uuid import uuid
@ -15,6 +15,7 @@ BEDROCK_MANAGED_S3_PREFIXES = (
BEDROCK_MANAGED_S3_UPLOAD_PREFIX,
BEDROCK_MANAGED_S3_OUTPUT_PREFIX,
)
_MAPPING_PROXY_TYPE: type = type(MappingProxyType({}))
_SAFE_OBJECT_COMPONENT_PATTERN = re.compile(r"[^A-Za-z0-9._-]+")
@ -110,12 +111,20 @@ def encode_s3_object_key_for_url(object_key: str) -> str:
return quote(unquote(object_key), safe="/")
def should_allow_legacy_cloud_file_ids(litellm_params: Optional[dict] = None) -> bool:
def should_allow_legacy_cloud_file_ids(
litellm_params: Optional[Mapping[str, Any]] = None,
) -> bool:
value = None
if isinstance(litellm_params, dict):
if isinstance(litellm_params, Mapping):
value = litellm_params.get("allow_legacy_cloud_file_ids")
if value is None:
value = os.getenv("LITELLM_ALLOW_LEGACY_CLOUD_FILE_IDS")
if value is None:
trusted_model_credentials = litellm_params.get(
"_litellm_internal_model_credentials"
)
if isinstance(trusted_model_credentials, _MAPPING_PROXY_TYPE):
value = cast(Mapping[str, Any], trusted_model_credentials).get(
"allow_legacy_cloud_file_ids"
)
if isinstance(value, bool):
return value

View file

@ -2,7 +2,7 @@ import asyncio
import base64
import os
from types import MappingProxyType
from typing import Any, Coroutine, Optional, Tuple, Union
from typing import Any, Coroutine, Mapping, Optional, Tuple, Union, cast
import httpx
@ -101,7 +101,14 @@ class BedrockFilesHandler(BaseAWSLLM):
)
bucket_name = None
if isinstance(trusted_model_credentials, type(MappingProxyType({}))):
bucket_name = trusted_model_credentials.get("s3_bucket_name")
trusted_model_credentials_mapping = cast(
Mapping[str, Any], trusted_model_credentials
)
candidate_bucket_name = trusted_model_credentials_mapping.get(
"s3_bucket_name"
)
if isinstance(candidate_bucket_name, str):
bucket_name = candidate_bucket_name
bucket_name = bucket_name or os.getenv("AWS_S3_BUCKET_NAME")
if not bucket_name:
raise ValueError(

View file

@ -11,7 +11,6 @@ from litellm.integrations.gcs_bucket.gcs_bucket_base import (
)
from litellm.litellm_core_utils.cloud_storage_security import (
VERTEX_AI_MANAGED_GCS_PREFIX,
should_allow_legacy_cloud_file_ids,
validate_managed_cloud_file_id,
)
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
@ -135,7 +134,6 @@ class VertexAIFilesHandler(GCSBucketBase):
scheme="gs://",
configured_bucket_name=configured_bucket_name,
allowed_object_prefixes=(VERTEX_AI_MANAGED_GCS_PREFIX,),
allow_legacy_cloud_file_ids=should_allow_legacy_cloud_file_ids(),
)
async def afile_content(

View file

@ -950,6 +950,7 @@ async def get_file(
data=data,
credentials=credentials, # type: ignore
file_id=original_file_id,
include_internal_credentials=True,
)
response = await litellm.afile_retrieve(**data) # type: ignore
@ -1150,6 +1151,7 @@ async def delete_file(
data=data,
credentials=credentials, # type: ignore
file_id=original_file_id,
include_internal_credentials=True,
)
response = await litellm.afile_delete(

View file

@ -184,3 +184,23 @@ def test_should_forward_trusted_model_credentials_to_bedrock_provider_config():
litellm_params = mock_retrieve_file_content.call_args.kwargs["litellm_params"]
assert litellm_params["_litellm_internal_model_credentials"] is trusted_credentials
assert "s3_bucket_name" not in litellm_params
def test_should_forward_trusted_model_credentials_to_retrieve_provider_config():
trusted_credentials = MappingProxyType({"allow_legacy_cloud_file_ids": True})
mock_response = MagicMock()
with patch.object(
files_main.base_llm_http_handler,
"retrieve_file",
return_value=mock_response,
) as mock_retrieve_file:
response = files_main.file_retrieve(
file_id="gs://safe-bucket/private/file.jsonl",
custom_llm_provider="vertex_ai",
_litellm_internal_model_credentials=trusted_credentials,
)
assert response is mock_response
litellm_params = mock_retrieve_file.call_args.kwargs["litellm_params"]
assert litellm_params["_litellm_internal_model_credentials"] is trusted_credentials

View file

@ -3,7 +3,6 @@ Test Vertex AI files handler functionality
"""
import asyncio
import os
import pytest
from unittest.mock import AsyncMock, patch
@ -67,22 +66,6 @@ class TestVertexAIFilesHandler:
configured_bucket_name="test-bucket",
)
def test_extract_bucket_and_object_from_file_id_allows_legacy_path_with_env_flag(
self,
):
file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt"
with patch.dict(os.environ, {"LITELLM_ALLOW_LEGACY_CLOUD_FILE_IDS": "true"}):
bucket_name, object_path = (
self.handler._extract_bucket_and_object_from_file_id(
file_id=file_id,
configured_bucket_name="test-bucket",
)
)
assert bucket_name == "test-bucket"
assert object_path == "test-file.txt"
def test_extract_bucket_and_object_from_file_id_rejects_no_gs_prefix(self):
"""Test extraction when gs:// prefix is missing"""
file_id = "test-bucket%2Flitellm-vertex-files%2Ftest-file.txt"

View file

@ -3,6 +3,7 @@ Tests for VertexAIFilesConfig transformation methods (Issues 5-7).
"""
import urllib.parse
from types import MappingProxyType
from urllib.parse import parse_qs, urlparse
import httpx
@ -85,6 +86,31 @@ class TestParseGcsUri:
assert bucket == "my-bucket"
assert encoded == urllib.parse.quote("private/object.txt", safe="")
def test_should_allow_legacy_object_path_with_trusted_server_flag(self, config):
trusted_credentials = MappingProxyType({"allow_legacy_cloud_file_ids": True})
bucket, encoded = config._parse_gcs_uri(
"gs://my-bucket/private/object.txt",
litellm_params={
"bucket_name": "my-bucket",
"_litellm_internal_model_credentials": trusted_credentials,
},
)
assert bucket == "my-bucket"
assert encoded == urllib.parse.quote("private/object.txt", safe="")
def test_should_reject_user_supplied_legacy_flag_snapshot(self, config):
with pytest.raises(ValueError, match="LiteLLM-managed"):
config._parse_gcs_uri(
"gs://my-bucket/private/object.txt",
litellm_params={
"bucket_name": "my-bucket",
"_litellm_internal_model_credentials": {
"allow_legacy_cloud_file_ids": True
},
},
)
def test_should_keep_configured_prefix_for_legacy_object_path(self, config):
bucket, encoded = config._parse_gcs_uri(
"gs://my-bucket/team-a/private/object.txt",