thread trusted params through vertex file content

This commit is contained in:
user 2026-05-01 18:24:22 -07:00
parent e7c23fc041
commit bb8f6fd5e2
4 changed files with 31 additions and 4 deletions

View file

@ -1016,6 +1016,7 @@ def file_content(
vertex_location=vertex_ai_location,
timeout=timeout,
max_retries=optional_params.max_retries,
litellm_params=litellm_params_dict,
)
elif custom_llm_provider == "bedrock":
response = bedrock_files_instance.file_content(

View file

@ -37,8 +37,6 @@ _supported_callback_params = [
"langfuse_secret_key",
"langfuse_host",
"langfuse_prompt_version",
"gcs_bucket_name",
"gcs_path_service_account",
"langsmith_api_key",
"langsmith_project",
"langsmith_base_url",
@ -69,7 +67,7 @@ def initialize_standard_callback_dynamic_params(
"""
Initialize the standard callback dynamic params from the kwargs
checks if langfuse_secret_key, gcs_bucket_name in kwargs and sets the corresponding attributes in StandardCallbackDynamicParams
checks supported request callback params in kwargs and sets the corresponding attributes in StandardCallbackDynamicParams
"""
standard_callback_dynamic_params = StandardCallbackDynamicParams()

View file

@ -11,6 +11,7 @@ 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
@ -117,7 +118,10 @@ class VertexAIFilesHandler(GCSBucketBase):
)
def _extract_bucket_and_object_from_file_id(
self, file_id: str, configured_bucket_name: str
self,
file_id: str,
configured_bucket_name: str,
litellm_params: Optional[dict] = None,
) -> Tuple[str, str]:
"""
Validate and extract bucket name and object path from file_id.
@ -134,6 +138,9 @@ 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(
litellm_params
),
)
async def afile_content(
@ -144,6 +151,7 @@ class VertexAIFilesHandler(GCSBucketBase):
vertex_location: Optional[str],
timeout: Union[float, httpx.Timeout],
max_retries: Optional[int],
litellm_params: Optional[dict] = None,
) -> HttpxBinaryResponseContent:
"""
Download file content from GCS bucket for VertexAI files.
@ -169,6 +177,7 @@ class VertexAIFilesHandler(GCSBucketBase):
bucket_name, object_path = self._extract_bucket_and_object_from_file_id(
file_id=file_id,
configured_bucket_name=gcs_logging_config["bucket_name"],
litellm_params=litellm_params,
)
download_kwargs = {
@ -205,6 +214,7 @@ class VertexAIFilesHandler(GCSBucketBase):
vertex_location: Optional[str],
timeout: Union[float, httpx.Timeout],
max_retries: Optional[int],
litellm_params: Optional[dict] = None,
) -> Union[
HttpxBinaryResponseContent, Coroutine[Any, Any, HttpxBinaryResponseContent]
]:
@ -233,6 +243,7 @@ class VertexAIFilesHandler(GCSBucketBase):
vertex_location=vertex_location,
timeout=timeout,
max_retries=max_retries,
litellm_params=litellm_params,
)
else:
return asyncio.run(
@ -243,5 +254,6 @@ class VertexAIFilesHandler(GCSBucketBase):
vertex_location=vertex_location,
timeout=timeout,
max_retries=max_retries,
litellm_params=litellm_params,
)
)

View file

@ -3,6 +3,7 @@ Test Vertex AI files handler functionality
"""
import asyncio
from types import MappingProxyType
import pytest
from unittest.mock import AsyncMock, patch
@ -66,6 +67,21 @@ class TestVertexAIFilesHandler:
configured_bucket_name="test-bucket",
)
def test_extract_bucket_and_object_from_file_id_allows_trusted_legacy_flag(self):
file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt"
trusted_credentials = MappingProxyType({"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",
litellm_params={
"_litellm_internal_model_credentials": trusted_credentials,
},
)
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"