mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-10 03:28:53 +00:00
thread trusted params through vertex file content
This commit is contained in:
parent
e7c23fc041
commit
bb8f6fd5e2
4 changed files with 31 additions and 4 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue