mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
harden cloud file compatibility path
This commit is contained in:
parent
2776e0162c
commit
e7c23fc041
8 changed files with 94 additions and 32 deletions
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue