mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
refactor(files): remove dead vertex_ai file_content branch and handler
This commit is contained in:
mateo-berri
parent
47f1fb394e
commit
866863dd64
8 changed files with 32 additions and 701 deletions
|
|
@ -24,7 +24,7 @@
|
|||
"limit": 42
|
||||
},
|
||||
"reportExplicitAny": {
|
||||
"limit": 10228
|
||||
"limit": 10226
|
||||
},
|
||||
"reportFunctionMemberAccess": {
|
||||
"limit": 11
|
||||
|
|
@ -57,7 +57,7 @@
|
|||
"limit": 5893
|
||||
},
|
||||
"reportMissingTypeArgument": {
|
||||
"limit": 15886
|
||||
"limit": 15883
|
||||
},
|
||||
"reportMissingTypeStubs": {
|
||||
"limit": 41
|
||||
|
|
@ -99,19 +99,19 @@
|
|||
"limit": 0
|
||||
},
|
||||
"reportUnknownArgumentType": {
|
||||
"limit": 45567
|
||||
"limit": 45564
|
||||
},
|
||||
"reportUnknownLambdaType": {
|
||||
"limit": 113
|
||||
},
|
||||
"reportUnknownMemberType": {
|
||||
"limit": 40525
|
||||
"limit": 40517
|
||||
},
|
||||
"reportUnknownParameterType": {
|
||||
"limit": 20384
|
||||
"limit": 20381
|
||||
},
|
||||
"reportUnknownVariableType": {
|
||||
"limit": 32099
|
||||
"limit": 32097
|
||||
},
|
||||
"reportUnnecessaryCast": {
|
||||
"limit": 177
|
||||
|
|
|
|||
|
|
@ -30,7 +30,6 @@ FileRetrieveProvider = Literal["openai", "azure", "gemini", "vertex_ai", "hosted
|
|||
FileDeleteProvider = Literal["openai", "azure", "gemini", "manus", "anthropic"]
|
||||
FileListProvider = Literal["openai", "azure", "manus", "anthropic"]
|
||||
import litellm
|
||||
from litellm import get_secret_str
|
||||
from litellm.files.streaming import FileContentStreamingResponse
|
||||
from litellm.files.types import FileContentProvider, FileContentStreamingResult
|
||||
from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
|
||||
|
|
@ -42,7 +41,6 @@ from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
|||
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
|
||||
from litellm.llms.openai.common_utils import get_openai_credentials
|
||||
from litellm.llms.openai.openai import FileDeleted, FileObject, OpenAIFilesAPI
|
||||
from litellm.llms.vertex_ai.files.handler import VertexAIFilesHandler
|
||||
from litellm.types.llms.openai import (
|
||||
CreateFileRequest,
|
||||
FileContentRequest,
|
||||
|
|
@ -79,7 +77,6 @@ def _should_sdk_support_streaming(
|
|||
|
||||
openai_files_instance = OpenAIFilesAPI()
|
||||
azure_files_instance = AzureOpenAIFilesAPI()
|
||||
vertex_ai_files_instance = VertexAIFilesHandler()
|
||||
bedrock_files_instance = BedrockFilesHandler()
|
||||
#################################################
|
||||
|
||||
|
|
@ -956,27 +953,6 @@ def file_content(
|
|||
client=client,
|
||||
litellm_params=litellm_params_dict,
|
||||
)
|
||||
elif custom_llm_provider == "vertex_ai":
|
||||
api_base = optional_params.api_base or ""
|
||||
vertex_ai_project = (
|
||||
optional_params.vertex_project or litellm.vertex_project or get_secret_str("VERTEXAI_PROJECT")
|
||||
)
|
||||
vertex_ai_location = (
|
||||
optional_params.vertex_location or litellm.vertex_location or get_secret_str("VERTEXAI_LOCATION")
|
||||
)
|
||||
vertex_credentials = optional_params.vertex_credentials or get_secret_str("VERTEXAI_CREDENTIALS")
|
||||
|
||||
response = vertex_ai_files_instance.file_content(
|
||||
_is_async=_is_async,
|
||||
file_content_request=_file_content_request,
|
||||
api_base=api_base,
|
||||
vertex_credentials=vertex_credentials,
|
||||
vertex_project=vertex_ai_project,
|
||||
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(
|
||||
_is_async=_is_async,
|
||||
|
|
|
|||
|
|
@ -1,7 +1,7 @@
|
|||
"""
|
||||
Supports writing files to Google AI Studio Files API.
|
||||
|
||||
For vertex ai, check out the vertex_ai/files/handler.py file.
|
||||
For vertex ai, check out the vertex_ai/files/transformation.py file.
|
||||
"""
|
||||
|
||||
import time
|
||||
|
|
|
|||
|
|
@ -1,234 +0,0 @@
|
|||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
from urllib.parse import unquote
|
||||
from typing import Any, Coroutine, Mapping, Optional, Tuple, Union
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm import LlmProviders
|
||||
from litellm.integrations.gcs_bucket.gcs_bucket_base import (
|
||||
GCSBucketBase,
|
||||
GCSLoggingConfig,
|
||||
)
|
||||
from litellm.types.utils import StandardCallbackDynamicParams
|
||||
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
|
||||
from litellm.types.llms.openai import (
|
||||
FileContentRequest,
|
||||
HttpxBinaryResponseContent,
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES
|
||||
|
||||
from .transformation import VertexAIFilesConfig
|
||||
|
||||
|
||||
class VertexAIFilesHandler(GCSBucketBase):
|
||||
"""
|
||||
Handles Calling VertexAI in OpenAI Files API format v1/files/*
|
||||
|
||||
This implementation uploads files on GCS Buckets
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.async_httpx_client = get_async_httpx_client(
|
||||
llm_provider=LlmProviders.VERTEX_AI,
|
||||
)
|
||||
|
||||
def _resolve_read_gcs_config(
|
||||
self,
|
||||
litellm_params: Mapping[str, object] | None,
|
||||
vertex_credentials: VERTEX_CREDENTIALS_TYPES | None,
|
||||
) -> tuple[str | None, str | None]:
|
||||
"""
|
||||
Resolve the GCS bucket and service-account credentials for the read/content path.
|
||||
|
||||
Sources them from the deployment's ``litellm_params`` (``gcs_bucket_name`` /
|
||||
``bucket_name`` and ``vertex_credentials``), mirroring the write path in
|
||||
``VertexAIFilesConfig._get_configured_bucket_name``, and falls back to the global
|
||||
``GCS_BUCKET_NAME`` / ``GCS_PATH_SERVICE_ACCOUNT`` env vars. This lets Vertex batch
|
||||
run entirely at the model-group level, so output written to a per-model bucket is
|
||||
readable without setting the global env vars.
|
||||
"""
|
||||
params: Mapping[str, object] = litellm_params or {}
|
||||
bucket_candidate = params.get("gcs_bucket_name") or params.get("bucket_name")
|
||||
configured_bucket_name = bucket_candidate if isinstance(bucket_candidate, str) else os.getenv("GCS_BUCKET_NAME")
|
||||
|
||||
credentials = params.get("vertex_credentials") or vertex_credentials
|
||||
if isinstance(credentials, dict):
|
||||
path_service_account: str | None = json.dumps(credentials)
|
||||
elif isinstance(credentials, str):
|
||||
path_service_account = credentials
|
||||
else:
|
||||
path_service_account = os.getenv("GCS_PATH_SERVICE_ACCOUNT")
|
||||
|
||||
return configured_bucket_name, path_service_account
|
||||
|
||||
def _extract_bucket_and_object_from_file_id(
|
||||
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.
|
||||
|
||||
Expected format: gs://bucket-name/litellm-vertex-files/path/to/file
|
||||
|
||||
Returns:
|
||||
tuple: (bucket_name, object_path)
|
||||
- bucket_name: "bucket-name"
|
||||
- object_path: "litellm-vertex-files/path/to/file"
|
||||
"""
|
||||
return validate_managed_cloud_file_id(
|
||||
file_id=file_id,
|
||||
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(
|
||||
self,
|
||||
file_content_request: FileContentRequest,
|
||||
vertex_credentials: Optional[VERTEX_CREDENTIALS_TYPES],
|
||||
vertex_project: Optional[str],
|
||||
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.
|
||||
|
||||
Args:
|
||||
file_content_request: Contains file_id (URL-encoded GCS path)
|
||||
vertex_credentials: VertexAI credentials
|
||||
vertex_project: VertexAI project ID
|
||||
vertex_location: VertexAI location
|
||||
timeout: Request timeout
|
||||
max_retries: Max retry attempts
|
||||
|
||||
Returns:
|
||||
HttpxBinaryResponseContent: Binary content wrapped in compatible response format
|
||||
"""
|
||||
file_id = file_content_request.get("file_id")
|
||||
if not file_id:
|
||||
raise ValueError("file_id is required in file_content_request")
|
||||
|
||||
configured_bucket_name, path_service_account = self._resolve_read_gcs_config(
|
||||
litellm_params=litellm_params,
|
||||
vertex_credentials=vertex_credentials,
|
||||
)
|
||||
dynamic_params = StandardCallbackDynamicParams(
|
||||
gcs_bucket_name=configured_bucket_name,
|
||||
gcs_path_service_account=path_service_account,
|
||||
)
|
||||
gcs_logging_config: GCSLoggingConfig = await self.get_gcs_logging_config(
|
||||
kwargs={"standard_callback_dynamic_params": dynamic_params}
|
||||
)
|
||||
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 = {
|
||||
"standard_callback_dynamic_params": {
|
||||
"gcs_bucket_name": bucket_name,
|
||||
"gcs_path_service_account": gcs_logging_config["path_service_account"],
|
||||
}
|
||||
}
|
||||
|
||||
file_content = await self.download_gcs_object(object_name=object_path, **download_kwargs)
|
||||
decoded_file_id = unquote(file_id)
|
||||
|
||||
if file_content is None:
|
||||
raise ValueError(f"Failed to download file from GCS: {decoded_file_id}")
|
||||
|
||||
mock_response = httpx.Response(
|
||||
status_code=200,
|
||||
content=file_content,
|
||||
headers={
|
||||
"content-type": "application/octet-stream",
|
||||
"content-length": str(len(file_content)),
|
||||
},
|
||||
request=httpx.Request(method="GET", url=decoded_file_id),
|
||||
)
|
||||
|
||||
# Apply transformation to convert Vertex AI batch outputs to OpenAI format
|
||||
config = VertexAIFilesConfig()
|
||||
|
||||
# Create a logging object for transformation
|
||||
logging_obj = Logging(
|
||||
model="",
|
||||
messages=[],
|
||||
stream=False,
|
||||
call_type="afile_content",
|
||||
start_time=time.time(),
|
||||
litellm_call_id="",
|
||||
function_id="",
|
||||
)
|
||||
|
||||
return config.transform_file_content_response(
|
||||
raw_response=mock_response, logging_obj=logging_obj, litellm_params={}
|
||||
)
|
||||
|
||||
def file_content(
|
||||
self,
|
||||
_is_async: bool,
|
||||
file_content_request: FileContentRequest,
|
||||
api_base: Optional[str],
|
||||
vertex_credentials: Optional[VERTEX_CREDENTIALS_TYPES],
|
||||
vertex_project: Optional[str],
|
||||
vertex_location: Optional[str],
|
||||
timeout: Union[float, httpx.Timeout],
|
||||
max_retries: Optional[int],
|
||||
litellm_params: Optional[dict] = None,
|
||||
) -> Union[HttpxBinaryResponseContent, Coroutine[Any, Any, HttpxBinaryResponseContent]]:
|
||||
"""
|
||||
Download file content from GCS bucket for VertexAI files.
|
||||
Supports both sync and async operations.
|
||||
|
||||
Args:
|
||||
_is_async: Whether to run asynchronously
|
||||
file_content_request: Contains file_id (URL-encoded GCS path)
|
||||
api_base: API base (unused for GCS operations)
|
||||
vertex_credentials: VertexAI credentials
|
||||
vertex_project: VertexAI project ID
|
||||
vertex_location: VertexAI location
|
||||
timeout: Request timeout
|
||||
max_retries: Max retry attempts
|
||||
|
||||
Returns:
|
||||
HttpxBinaryResponseContent or Coroutine: Binary content wrapped in compatible response format
|
||||
"""
|
||||
if _is_async:
|
||||
return self.afile_content(
|
||||
file_content_request=file_content_request,
|
||||
vertex_credentials=vertex_credentials,
|
||||
vertex_project=vertex_project,
|
||||
vertex_location=vertex_location,
|
||||
timeout=timeout,
|
||||
max_retries=max_retries,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
else:
|
||||
return asyncio.run(
|
||||
self.afile_content(
|
||||
file_content_request=file_content_request,
|
||||
vertex_credentials=vertex_credentials,
|
||||
vertex_project=vertex_project,
|
||||
vertex_location=vertex_location,
|
||||
timeout=timeout,
|
||||
max_retries=max_retries,
|
||||
litellm_params=litellm_params,
|
||||
)
|
||||
)
|
||||
|
|
@ -15,7 +15,7 @@
|
|||
"limit": 944
|
||||
},
|
||||
"ANN204": {
|
||||
"limit": 724
|
||||
"limit": 723
|
||||
},
|
||||
"ANN205": {
|
||||
"limit": 127
|
||||
|
|
@ -123,7 +123,7 @@
|
|||
"limit": 52
|
||||
},
|
||||
"I001": {
|
||||
"limit": 270
|
||||
"limit": 269
|
||||
},
|
||||
"LOG015": {
|
||||
"limit": 8
|
||||
|
|
@ -306,7 +306,7 @@
|
|||
"limit": 9
|
||||
},
|
||||
"TID251": {
|
||||
"limit": 2652
|
||||
"limit": 2651
|
||||
},
|
||||
"TRY002": {
|
||||
"limit": 548
|
||||
|
|
@ -324,10 +324,10 @@
|
|||
"limit": 883
|
||||
},
|
||||
"UP006": {
|
||||
"limit": 12147
|
||||
"limit": 12146
|
||||
},
|
||||
"UP007": {
|
||||
"limit": 2526
|
||||
"limit": 2523
|
||||
},
|
||||
"UP008": {
|
||||
"limit": 5
|
||||
|
|
@ -354,7 +354,7 @@
|
|||
"limit": 4
|
||||
},
|
||||
"UP035": {
|
||||
"limit": 2232
|
||||
"limit": 2230
|
||||
},
|
||||
"UP036": {
|
||||
"limit": 4
|
||||
|
|
@ -363,6 +363,6 @@
|
|||
"limit": 105
|
||||
},
|
||||
"UP045": {
|
||||
"limit": 17824
|
||||
"limit": 17812
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1,411 +0,0 @@
|
|||
"""
|
||||
Test Vertex AI files handler functionality
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from types import MappingProxyType
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm.llms.vertex_ai.files.handler import VertexAIFilesHandler
|
||||
from litellm.types.llms.openai import FileContentRequest, HttpxBinaryResponseContent
|
||||
|
||||
|
||||
def _mock_gcs_logging_config(bucket_name: str = "test-bucket"):
|
||||
return {
|
||||
"bucket_name": bucket_name,
|
||||
"path_service_account": None,
|
||||
"vertex_instance": None,
|
||||
}
|
||||
|
||||
|
||||
class TestVertexAIFilesHandler:
|
||||
"""Test Vertex AI files handler"""
|
||||
|
||||
def setup_method(self):
|
||||
"""Setup test method"""
|
||||
self.handler = VertexAIFilesHandler()
|
||||
|
||||
def test_extract_bucket_and_object_from_file_id_standard_path(self):
|
||||
"""Test extraction of bucket and object from URL-encoded file_id with standard path"""
|
||||
# Sample file_id with nested folder structure
|
||||
file_id = "gs%3A%2F%2Ftest-bucket%2Flitellm-vertex-files%2Ftest-folder%2Fsub-folder%2Ftest-file.txt"
|
||||
|
||||
bucket_name, object_path = self.handler._extract_bucket_and_object_from_file_id(
|
||||
file_id=file_id,
|
||||
configured_bucket_name="test-bucket",
|
||||
)
|
||||
|
||||
# Verify bucket name extraction
|
||||
assert bucket_name == "test-bucket"
|
||||
|
||||
expected_object = "litellm-vertex-files/test-folder/sub-folder/test-file.txt"
|
||||
assert object_path == expected_object
|
||||
|
||||
def test_extract_bucket_and_object_from_file_id_rejects_bucket_only(self):
|
||||
"""Test extraction when only bucket name is provided"""
|
||||
file_id = "gs%3A%2F%2Ftest-bucket"
|
||||
|
||||
with pytest.raises(ValueError, match="object name"):
|
||||
self.handler._extract_bucket_and_object_from_file_id(
|
||||
file_id=file_id,
|
||||
configured_bucket_name="test-bucket",
|
||||
)
|
||||
|
||||
def test_extract_bucket_and_object_from_file_id_rejects_unmanaged_path(self):
|
||||
"""Test extraction with simple path"""
|
||||
file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt"
|
||||
|
||||
with pytest.raises(ValueError, match="LiteLLM-managed"):
|
||||
self.handler._extract_bucket_and_object_from_file_id(
|
||||
file_id=file_id,
|
||||
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"
|
||||
|
||||
with pytest.raises(ValueError, match="gs://"):
|
||||
self.handler._extract_bucket_and_object_from_file_id(
|
||||
file_id=file_id,
|
||||
configured_bucket_name="test-bucket",
|
||||
)
|
||||
|
||||
def test_extract_bucket_and_object_from_file_id_rejects_wrong_bucket(self):
|
||||
file_id = "gs%3A%2F%2Fother-bucket%2Flitellm-vertex-files%2Ftest-file.txt"
|
||||
|
||||
with pytest.raises(ValueError, match="configured storage bucket"):
|
||||
self.handler._extract_bucket_and_object_from_file_id(
|
||||
file_id=file_id,
|
||||
configured_bucket_name="test-bucket",
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_afile_content_success(self):
|
||||
"""Test successful async file content retrieval"""
|
||||
# Setup test data
|
||||
file_id = "gs%3A%2F%2Ftest-bucket%2Flitellm-vertex-files%2Fuploads%2Fabc-test-file.txt"
|
||||
expected_content = b"test file content"
|
||||
|
||||
file_content_request = FileContentRequest(file_id=file_id, extra_headers=None, extra_body=None)
|
||||
|
||||
# Mock the download_gcs_object method
|
||||
with (
|
||||
patch.object(self.handler, "download_gcs_object", new_callable=AsyncMock) as mock_download,
|
||||
patch.object(
|
||||
self.handler,
|
||||
"get_gcs_logging_config",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_mock_gcs_logging_config(),
|
||||
),
|
||||
):
|
||||
mock_download.return_value = expected_content
|
||||
|
||||
# Call the method
|
||||
result = await self.handler.afile_content(
|
||||
file_content_request=file_content_request,
|
||||
vertex_credentials=None,
|
||||
vertex_project="test-project",
|
||||
vertex_location="us-central1",
|
||||
timeout=60.0,
|
||||
max_retries=3,
|
||||
)
|
||||
|
||||
# Verify the result
|
||||
assert isinstance(result, HttpxBinaryResponseContent)
|
||||
assert hasattr(result, "response")
|
||||
assert result.response.content == expected_content
|
||||
assert result.response.status_code == 200
|
||||
|
||||
# Verify the download was called with correct parameters
|
||||
mock_download.assert_called_once()
|
||||
call_args = mock_download.call_args
|
||||
assert call_args.kwargs["object_name"] == "litellm-vertex-files/uploads/abc-test-file.txt"
|
||||
assert "standard_callback_dynamic_params" in call_args.kwargs
|
||||
assert call_args.kwargs["standard_callback_dynamic_params"]["gcs_bucket_name"] == "test-bucket"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_afile_content_missing_file_id(self):
|
||||
"""Test async file content retrieval with missing file_id"""
|
||||
file_content_request = FileContentRequest(extra_headers=None, extra_body=None)
|
||||
|
||||
# Should raise ValueError for missing file_id
|
||||
with pytest.raises(ValueError, match="file_id is required in file_content_request"):
|
||||
await self.handler.afile_content(
|
||||
file_content_request=file_content_request,
|
||||
vertex_credentials=None,
|
||||
vertex_project="test-project",
|
||||
vertex_location="us-central1",
|
||||
timeout=60.0,
|
||||
max_retries=3,
|
||||
)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_afile_content_download_failure(self):
|
||||
"""Test async file content retrieval when download fails"""
|
||||
file_id = "gs%3A%2F%2Ftest-bucket%2Flitellm-vertex-files%2Fuploads%2Fabc-test-file.txt"
|
||||
|
||||
file_content_request = FileContentRequest(file_id=file_id, extra_headers=None, extra_body=None)
|
||||
|
||||
# Mock download to return None (failure)
|
||||
with (
|
||||
patch.object(self.handler, "download_gcs_object", new_callable=AsyncMock) as mock_download,
|
||||
patch.object(
|
||||
self.handler,
|
||||
"get_gcs_logging_config",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_mock_gcs_logging_config(),
|
||||
),
|
||||
):
|
||||
mock_download.return_value = None
|
||||
|
||||
# Should raise ValueError for failed download
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="Failed to download file from GCS: gs://test-bucket/litellm-vertex-files/uploads/abc-test-file.txt",
|
||||
):
|
||||
await self.handler.afile_content(
|
||||
file_content_request=file_content_request,
|
||||
vertex_credentials=None,
|
||||
vertex_project="test-project",
|
||||
vertex_location="us-central1",
|
||||
timeout=60.0,
|
||||
max_retries=3,
|
||||
)
|
||||
|
||||
def test_resolve_read_gcs_config_prefers_per_model_bucket(self, monkeypatch):
|
||||
monkeypatch.setenv("GCS_BUCKET_NAME", "env-default-bucket")
|
||||
monkeypatch.setenv("GCS_PATH_SERVICE_ACCOUNT", "/env/sa.json")
|
||||
|
||||
bucket, service_account = self.handler._resolve_read_gcs_config(
|
||||
litellm_params={
|
||||
"gcs_bucket_name": "my-model-bucket",
|
||||
"vertex_credentials": "/model/sa.json",
|
||||
},
|
||||
vertex_credentials=None,
|
||||
)
|
||||
|
||||
assert bucket == "my-model-bucket"
|
||||
assert service_account == "/model/sa.json"
|
||||
|
||||
def test_resolve_read_gcs_config_falls_back_to_env(self, monkeypatch):
|
||||
monkeypatch.setenv("GCS_BUCKET_NAME", "env-default-bucket")
|
||||
monkeypatch.setenv("GCS_PATH_SERVICE_ACCOUNT", "/env/sa.json")
|
||||
|
||||
bucket, service_account = self.handler._resolve_read_gcs_config(litellm_params={}, vertex_credentials=None)
|
||||
|
||||
assert bucket == "env-default-bucket"
|
||||
assert service_account == "/env/sa.json"
|
||||
|
||||
def test_resolve_read_gcs_config_serializes_dict_credentials(self, monkeypatch):
|
||||
monkeypatch.delenv("GCS_PATH_SERVICE_ACCOUNT", raising=False)
|
||||
|
||||
_, service_account = self.handler._resolve_read_gcs_config(
|
||||
litellm_params={"gcs_bucket_name": "my-model-bucket"},
|
||||
vertex_credentials={"type": "service_account", "project_id": "p"},
|
||||
)
|
||||
|
||||
assert service_account == '{"type": "service_account", "project_id": "p"}'
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_afile_content_honors_per_model_bucket_over_env(self, monkeypatch):
|
||||
"""
|
||||
Regression for #32640: a batch output written to a per-model gcs_bucket_name must be
|
||||
readable even when the global GCS_BUCKET_NAME points at a different bucket. Before the
|
||||
fix the read path resolved the bucket from env only and raised
|
||||
"file_id bucket does not match the configured storage bucket".
|
||||
"""
|
||||
monkeypatch.setenv("GCS_BUCKET_NAME", "env-default-bucket")
|
||||
monkeypatch.delenv("GCS_PATH_SERVICE_ACCOUNT", raising=False)
|
||||
|
||||
file_id = "gs%3A%2F%2Fmy-model-bucket%2Flitellm-vertex-files%2Fuploads%2Fabc-batch-output.jsonl"
|
||||
file_content_request = FileContentRequest(file_id=file_id, extra_headers=None, extra_body=None)
|
||||
|
||||
with (
|
||||
patch.object(self.handler, "download_gcs_object", new_callable=AsyncMock) as mock_download,
|
||||
patch.object(
|
||||
self.handler,
|
||||
"get_or_create_vertex_instance",
|
||||
new_callable=AsyncMock,
|
||||
return_value=object(),
|
||||
),
|
||||
):
|
||||
mock_download.return_value = b"batch output"
|
||||
|
||||
result = await self.handler.afile_content(
|
||||
file_content_request=file_content_request,
|
||||
vertex_credentials="/model/sa.json",
|
||||
vertex_project="test-project",
|
||||
vertex_location="us-central1",
|
||||
timeout=60.0,
|
||||
max_retries=0,
|
||||
litellm_params={
|
||||
"gcs_bucket_name": "my-model-bucket",
|
||||
"vertex_credentials": "/model/sa.json",
|
||||
},
|
||||
)
|
||||
|
||||
assert isinstance(result, HttpxBinaryResponseContent)
|
||||
assert result.response.content == b"batch output"
|
||||
|
||||
dynamic_params = mock_download.call_args.kwargs["standard_callback_dynamic_params"]
|
||||
assert dynamic_params["gcs_bucket_name"] == "my-model-bucket"
|
||||
assert dynamic_params["gcs_path_service_account"] == "/model/sa.json"
|
||||
assert mock_download.call_args.kwargs["object_name"] == "litellm-vertex-files/uploads/abc-batch-output.jsonl"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_afile_content_reads_without_global_env_bucket(self, monkeypatch):
|
||||
"""
|
||||
Regression for #32640: with no global GCS_BUCKET_NAME set, a model-group-level
|
||||
deployment (per-model gcs_bucket_name) must still be readable. Before the fix the read
|
||||
path raised "GCS_BUCKET_NAME is not set in the environment".
|
||||
"""
|
||||
monkeypatch.delenv("GCS_BUCKET_NAME", raising=False)
|
||||
monkeypatch.delenv("GCS_PATH_SERVICE_ACCOUNT", raising=False)
|
||||
|
||||
file_id = "gs%3A%2F%2Fmy-model-bucket%2Flitellm-vertex-files%2Fuploads%2Fabc-batch-output.jsonl"
|
||||
file_content_request = FileContentRequest(file_id=file_id, extra_headers=None, extra_body=None)
|
||||
|
||||
with (
|
||||
patch.object(self.handler, "download_gcs_object", new_callable=AsyncMock) as mock_download,
|
||||
patch.object(
|
||||
self.handler,
|
||||
"get_or_create_vertex_instance",
|
||||
new_callable=AsyncMock,
|
||||
return_value=object(),
|
||||
),
|
||||
):
|
||||
mock_download.return_value = b"batch output"
|
||||
|
||||
result = await self.handler.afile_content(
|
||||
file_content_request=file_content_request,
|
||||
vertex_credentials="/model/sa.json",
|
||||
vertex_project="test-project",
|
||||
vertex_location="us-central1",
|
||||
timeout=60.0,
|
||||
max_retries=0,
|
||||
litellm_params={"gcs_bucket_name": "my-model-bucket"},
|
||||
)
|
||||
|
||||
assert isinstance(result, HttpxBinaryResponseContent)
|
||||
dynamic_params = mock_download.call_args.kwargs["standard_callback_dynamic_params"]
|
||||
assert dynamic_params["gcs_bucket_name"] == "my-model-bucket"
|
||||
|
||||
def test_file_content_sync_success(self):
|
||||
"""Test successful sync file content retrieval"""
|
||||
file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt"
|
||||
expected_content = b"test file content"
|
||||
|
||||
file_content_request = FileContentRequest(file_id=file_id, extra_headers=None, extra_body=None)
|
||||
|
||||
# Create expected response
|
||||
mock_response = httpx.Response(
|
||||
status_code=200,
|
||||
content=expected_content,
|
||||
headers={"content-type": "application/octet-stream"},
|
||||
request=httpx.Request(method="GET", url="gs://test-bucket/test-file.txt"),
|
||||
)
|
||||
expected_result = HttpxBinaryResponseContent(response=mock_response)
|
||||
|
||||
# Mock asyncio.run to return our expected result
|
||||
with patch("asyncio.run") as mock_run:
|
||||
mock_run.return_value = expected_result
|
||||
|
||||
result = self.handler.file_content(
|
||||
_is_async=False,
|
||||
file_content_request=file_content_request,
|
||||
api_base="",
|
||||
vertex_credentials=None,
|
||||
vertex_project="test-project",
|
||||
vertex_location="us-central1",
|
||||
timeout=60.0,
|
||||
max_retries=3,
|
||||
)
|
||||
|
||||
# Verify the result
|
||||
assert result == expected_result
|
||||
|
||||
# Verify asyncio.run was called (indicating sync execution)
|
||||
mock_run.assert_called_once()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_file_content_async_mode(self):
|
||||
"""Test async file content retrieval when _is_async=True"""
|
||||
file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt"
|
||||
expected_content = b"test file content"
|
||||
|
||||
file_content_request = FileContentRequest(file_id=file_id, extra_headers=None, extra_body=None)
|
||||
|
||||
# Mock the afile_content method
|
||||
with patch.object(self.handler, "afile_content", new_callable=AsyncMock) as mock_afile_content:
|
||||
mock_response = httpx.Response(
|
||||
status_code=200,
|
||||
content=expected_content,
|
||||
headers={"content-type": "application/octet-stream"},
|
||||
request=httpx.Request(method="GET", url="gs://test-bucket/test-file.txt"),
|
||||
)
|
||||
mock_afile_content.return_value = HttpxBinaryResponseContent(response=mock_response)
|
||||
|
||||
# Call the method with _is_async=True
|
||||
result = self.handler.file_content(
|
||||
_is_async=True,
|
||||
file_content_request=file_content_request,
|
||||
api_base="",
|
||||
vertex_credentials=None,
|
||||
vertex_project="test-project",
|
||||
vertex_location="us-central1",
|
||||
timeout=60.0,
|
||||
max_retries=3,
|
||||
)
|
||||
|
||||
# Should return a coroutine since _is_async=True
|
||||
assert asyncio.iscoroutine(result)
|
||||
|
||||
# Await the result
|
||||
final_result = await result
|
||||
assert isinstance(final_result, HttpxBinaryResponseContent)
|
||||
assert final_result.response.content == expected_content
|
||||
|
||||
def test_httpx_response_compatibility(self):
|
||||
"""Test that the created HttpxBinaryResponseContent is compatible with expected interface"""
|
||||
# Test the mock response creation logic
|
||||
expected_content = b"test file content"
|
||||
decoded_path = "gs://test-bucket/test-file.txt"
|
||||
|
||||
mock_response = httpx.Response(
|
||||
status_code=200,
|
||||
content=expected_content,
|
||||
headers={"content-type": "application/octet-stream"},
|
||||
request=httpx.Request(method="GET", url=decoded_path),
|
||||
)
|
||||
|
||||
result = HttpxBinaryResponseContent(response=mock_response)
|
||||
|
||||
# Verify the response properties
|
||||
assert result.response.status_code == 200
|
||||
assert result.response.content == expected_content
|
||||
assert result.response.headers["content-type"] == "application/octet-stream"
|
||||
|
||||
# Verify it has the expected interface (matching OpenAI file content response)
|
||||
assert hasattr(result, "response")
|
||||
assert hasattr(result.response, "content")
|
||||
assert hasattr(result.response, "status_code")
|
||||
assert hasattr(result.response, "headers")
|
||||
|
|
@ -136,23 +136,23 @@ class TestVertexAIFilesIntegration:
|
|||
# Verify provider detection was called
|
||||
mock_get_provider.assert_called_once()
|
||||
|
||||
def test_litellm_file_content_vertex_ai_error_cases(self):
|
||||
def test_litellm_file_content_vertex_ai_error_cases(self, monkeypatch):
|
||||
"""Test error handling in vertex_ai file_content"""
|
||||
# Test missing file_id - the VertexAI provider config's
|
||||
# transform_file_content_request should handle empty file_id.
|
||||
# Since the code now goes through base_llm_http_handler, we mock
|
||||
# ProviderConfigManager to return None so it falls through to the
|
||||
# old vertex_ai code path that validates file_id.
|
||||
with patch(
|
||||
"litellm.files.main.ProviderConfigManager.get_provider_files_config",
|
||||
return_value=None,
|
||||
):
|
||||
with pytest.raises(ValueError, match="file_id is required"):
|
||||
litellm.file_content(
|
||||
file_id="", # Empty file_id should cause error
|
||||
custom_llm_provider="vertex_ai",
|
||||
vertex_project="test-project",
|
||||
)
|
||||
monkeypatch.delenv("GCS_BUCKET_NAME", raising=False)
|
||||
with pytest.raises(ValueError, match="bucket_name is required"):
|
||||
litellm.file_content(
|
||||
file_id="",
|
||||
custom_llm_provider="vertex_ai",
|
||||
vertex_project="test-project",
|
||||
)
|
||||
|
||||
monkeypatch.setenv("GCS_BUCKET_NAME", "test-bucket")
|
||||
with pytest.raises(ValueError, match="gs://"):
|
||||
litellm.file_content(
|
||||
file_id="",
|
||||
custom_llm_provider="vertex_ai",
|
||||
vertex_project="test-project",
|
||||
)
|
||||
|
||||
def test_vertex_ai_provider_in_supported_providers_list(self):
|
||||
"""Test that vertex_ai is included in supported providers for file_content"""
|
||||
|
|
|
|||
|
|
@ -1,9 +1,9 @@
|
|||
{
|
||||
"LIT001": {
|
||||
"limit": 23287
|
||||
"limit": 23284
|
||||
},
|
||||
"LIT002": {
|
||||
"limit": 27473
|
||||
"limit": 27466
|
||||
},
|
||||
"LIT003": {
|
||||
"limit": 292
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue