mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
Merge 65c3be3ada into b4a8bc715f
This commit is contained in:
Mateo Wang
GitHub
commit
baacb68f2e
6 changed files with 50 additions and 548 deletions
|
|
@ -35,7 +35,6 @@ FileDeleteProvider = Literal[
|
|||
]
|
||||
FileListProvider = Literal["openai", "azure", "litellm_proxy", "manus", "anthropic", "mistral", "xai"]
|
||||
import litellm
|
||||
from litellm import get_secret_str
|
||||
from litellm.files.streaming import FileContentStreamingResponse
|
||||
from litellm.files.types import FileContentProvider, FileContentStreamingResult, FileRetrieveProvider
|
||||
from litellm.litellm_core_utils.get_litellm_params import add_trusted_model_credentials_to_litellm_params
|
||||
|
|
@ -48,7 +47,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.llms.xai.batches.handler import XAIBatchesHandler
|
||||
from litellm.llms.xai.batches.transformation import is_xai_batch_results_id
|
||||
from litellm.types.llms.openai import (
|
||||
|
|
@ -103,7 +101,6 @@ def _file_content_logging_obj(kwargs: dict[str, object], _is_async: bool) -> Lit
|
|||
|
||||
openai_files_instance: Final = OpenAIFilesAPI()
|
||||
azure_files_instance: Final = AzureOpenAIFilesAPI()
|
||||
vertex_ai_files_instance: Final = VertexAIFilesHandler()
|
||||
bedrock_files_instance: Final = BedrockFilesHandler()
|
||||
xai_batch_results_instance: Final = XAIBatchesHandler()
|
||||
#################################################
|
||||
|
|
@ -1003,27 +1000,6 @@ def file_content(
|
|||
client=client,
|
||||
litellm_params=litellm_params_dict,
|
||||
)
|
||||
elif custom_llm_provider == "vertex_ai":
|
||||
api_base: Final = optional_params.api_base or ""
|
||||
vertex_ai_project: Final = (
|
||||
optional_params.vertex_project or litellm.vertex_project or get_secret_str("VERTEXAI_PROJECT")
|
||||
)
|
||||
vertex_ai_location: Final = (
|
||||
optional_params.vertex_location or litellm.vertex_location or get_secret_str("VERTEXAI_LOCATION")
|
||||
)
|
||||
vertex_credentials: Final = 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,239 +0,0 @@
|
|||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
from collections.abc import Coroutine, Mapping
|
||||
from typing import Final
|
||||
from urllib.parse import unquote
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm import LlmProviders
|
||||
from litellm.integrations.gcs_bucket.gcs_bucket_base import (
|
||||
GCSBucketBase,
|
||||
GCSLoggingConfig,
|
||||
)
|
||||
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.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.llms.custom_httpx.http_handler import get_async_httpx_client
|
||||
from litellm.types.llms.openai import (
|
||||
FileContentRequest,
|
||||
HttpxBinaryResponseContent,
|
||||
)
|
||||
from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES
|
||||
from litellm.types.utils import StandardCallbackDynamicParams
|
||||
|
||||
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
|
||||
``GCS_BATCH_BUCKET_NAME`` then ``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: Final[Mapping[str, object]] = litellm_params or {}
|
||||
bucket_candidate: Final = params.get("gcs_bucket_name") or params.get("bucket_name")
|
||||
configured_bucket_name = (
|
||||
bucket_candidate
|
||||
if isinstance(bucket_candidate, str)
|
||||
else os.getenv("GCS_BATCH_BUCKET_NAME") or os.getenv("GCS_BUCKET_NAME")
|
||||
)
|
||||
|
||||
credentials: Final = 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: dict | None = 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: VERTEX_CREDENTIALS_TYPES | None,
|
||||
vertex_project: str | None,
|
||||
vertex_location: str | None,
|
||||
timeout: float | httpx.Timeout,
|
||||
max_retries: int | None,
|
||||
litellm_params: dict | None = 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: Final = 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: Final = StandardCallbackDynamicParams(
|
||||
gcs_bucket_name=configured_bucket_name,
|
||||
gcs_path_service_account=path_service_account,
|
||||
)
|
||||
gcs_logging_config: Final[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: Final = {
|
||||
"standard_callback_dynamic_params": {
|
||||
"gcs_bucket_name": bucket_name,
|
||||
"gcs_path_service_account": gcs_logging_config["path_service_account"],
|
||||
}
|
||||
}
|
||||
|
||||
file_content: Final = await self.download_gcs_object(object_name=object_path, **download_kwargs)
|
||||
decoded_file_id: Final = unquote(file_id)
|
||||
|
||||
if file_content is None:
|
||||
raise ValueError(f"Failed to download file from GCS: {decoded_file_id}")
|
||||
|
||||
mock_response: Final = 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: Final = VertexAIFilesConfig()
|
||||
|
||||
# Create a logging object for transformation
|
||||
logging_obj: Final = 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: str | None,
|
||||
vertex_credentials: VERTEX_CREDENTIALS_TYPES | None,
|
||||
vertex_project: str | None,
|
||||
vertex_location: str | None,
|
||||
timeout: float | httpx.Timeout,
|
||||
max_retries: int | None,
|
||||
litellm_params: dict | None = None,
|
||||
) -> HttpxBinaryResponseContent | Coroutine[object, object, 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,
|
||||
)
|
||||
)
|
||||
|
|
@ -9,7 +9,7 @@ import respx
|
|||
from pydantic import ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
|
||||
from litellm.types.llms.openai import OpenAIFileObject
|
||||
|
||||
NATIVE_VERTEX_ROWS: Final = (
|
||||
|
|
@ -243,3 +243,40 @@ async def test_openai_file_operations_roundtrip(respx_mock: respx.MockRouter, mo
|
|||
assert delete_route.call_count == 1
|
||||
assert deleted.id == "file-abc123"
|
||||
assert deleted.deleted is True
|
||||
|
||||
|
||||
@pytest.mark.parametrize("is_async", [False, True], ids=["sync", "async"])
|
||||
@pytest.mark.asyncio
|
||||
async def test_vertex_file_content_reads_from_the_configured_bucket(is_async: bool):
|
||||
content: Final = b"batch output\n"
|
||||
|
||||
def download(request: httpx.Request) -> httpx.Response:
|
||||
assert str(request.url) == (
|
||||
"https://storage.googleapis.com/storage/v1/b/model-bucket/"
|
||||
"o/litellm-vertex-files%2Fuploads%2Foutput.jsonl?alt=media"
|
||||
)
|
||||
return httpx.Response(200, content=content)
|
||||
|
||||
transport: Final = httpx.MockTransport(download)
|
||||
if is_async:
|
||||
handler: Final = AsyncHTTPHandler(transport=transport)
|
||||
async with handler.client:
|
||||
result: Final = await litellm.afile_content(
|
||||
file_id="gs://model-bucket/litellm-vertex-files/uploads/output.jsonl",
|
||||
custom_llm_provider="vertex_ai",
|
||||
gcs_bucket_name="model-bucket",
|
||||
api_key="test-token",
|
||||
client=handler,
|
||||
)
|
||||
assert result.content == content
|
||||
return
|
||||
|
||||
with httpx.Client(transport=transport) as sync_client:
|
||||
sync_result: Final = litellm.file_content(
|
||||
file_id="gs://model-bucket/litellm-vertex-files/uploads/output.jsonl",
|
||||
custom_llm_provider="vertex_ai",
|
||||
gcs_bucket_name="model-bucket",
|
||||
api_key="test-token",
|
||||
client=HTTPHandler(client=sync_client),
|
||||
)
|
||||
assert sync_result.content == content
|
||||
|
|
|
|||
|
|
@ -1,262 +0,0 @@
|
|||
"""
|
||||
Test Vertex AI files handler functionality
|
||||
"""
|
||||
|
||||
import re
|
||||
from types import MappingProxyType
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
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=re.escape(
|
||||
"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.delenv("GCS_BATCH_BUCKET_NAME", raising=False)
|
||||
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_prefers_batch_env_over_logging_env(self, monkeypatch):
|
||||
monkeypatch.setenv("GCS_BATCH_BUCKET_NAME", "batch-bucket")
|
||||
monkeypatch.setenv("GCS_BUCKET_NAME", "logging-bucket")
|
||||
|
||||
bucket, _ = self.handler._resolve_read_gcs_config(litellm_params={}, vertex_credentials=None)
|
||||
|
||||
assert bucket == "batch-bucket"
|
||||
|
||||
def test_resolve_read_gcs_config_prefers_per_model_bucket_over_batch_env(self, monkeypatch):
|
||||
monkeypatch.setenv("GCS_BATCH_BUCKET_NAME", "batch-bucket")
|
||||
|
||||
bucket, _ = self.handler._resolve_read_gcs_config(
|
||||
litellm_params={"gcs_bucket_name": "my-model-bucket"},
|
||||
vertex_credentials=None,
|
||||
)
|
||||
|
||||
assert bucket == "my-model-bucket"
|
||||
|
||||
def test_resolve_read_gcs_config_prefers_gcs_bucket_name_over_legacy(self):
|
||||
bucket, _ = self.handler._resolve_read_gcs_config(
|
||||
litellm_params={"gcs_bucket_name": "my-model-bucket", "bucket_name": "legacy-bucket"},
|
||||
vertex_credentials=None,
|
||||
)
|
||||
|
||||
assert bucket == "my-model-bucket"
|
||||
|
||||
def test_resolve_read_gcs_config_accepts_legacy_bucket_name_alone(self):
|
||||
bucket, _ = self.handler._resolve_read_gcs_config(
|
||||
litellm_params={"bucket_name": "legacy-bucket"},
|
||||
vertex_credentials=None,
|
||||
)
|
||||
|
||||
assert bucket == "legacy-bucket"
|
||||
|
||||
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"}'
|
||||
|
|
@ -2,8 +2,9 @@
|
|||
Test Vertex AI files integration with main files API
|
||||
"""
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import litellm
|
||||
from litellm.types.llms.openai import HttpxBinaryResponseContent
|
||||
|
|
@ -12,26 +13,15 @@ from litellm.types.llms.openai import HttpxBinaryResponseContent
|
|||
class TestVertexAIFilesIntegration:
|
||||
"""Test integration of Vertex AI files with main litellm API"""
|
||||
|
||||
|
||||
|
||||
|
||||
def test_litellm_file_content_vertex_ai_error_cases(self):
|
||||
"""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",
|
||||
)
|
||||
@pytest.mark.parametrize("file_id", ["", "not-a-gcs-file"])
|
||||
def test_litellm_file_content_vertex_ai_error_cases(self, file_id: str):
|
||||
with pytest.raises(ValueError, match="gs://"):
|
||||
litellm.file_content(
|
||||
file_id=file_id,
|
||||
custom_llm_provider="vertex_ai",
|
||||
gcs_bucket_name="test-bucket",
|
||||
api_key="test-token",
|
||||
)
|
||||
|
||||
def test_vertex_ai_provider_in_supported_providers_list(self):
|
||||
"""Test that vertex_ai is included in supported providers for file_content"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue