From 65c3be3adadc490738bd0697f7f15f2a959a8590 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 28 Oct 2026 23:36:34 -0700 Subject: [PATCH] refactor(files): remove unreachable Vertex AI file handler --- litellm/files/main.py | 24 -- litellm/llms/gemini/files/transformation.py | 2 +- litellm/llms/vertex_ai/files/handler.py | 239 ---------------- tests/unit/files/test_main.py | 39 ++- .../files/test_vertex_ai_files_handler.py | 262 ------------------ .../files/test_vertex_ai_files_integration.py | 32 +-- 6 files changed, 50 insertions(+), 548 deletions(-) delete mode 100644 litellm/llms/vertex_ai/files/handler.py delete mode 100644 tests/unit/llms/vertex_ai/files/test_vertex_ai_files_handler.py diff --git a/litellm/files/main.py b/litellm/files/main.py index c5172bc01ae..0b6e0a3ee76 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -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, diff --git a/litellm/llms/gemini/files/transformation.py b/litellm/llms/gemini/files/transformation.py index f9a7e230ea9..9b5e476afa0 100644 --- a/litellm/llms/gemini/files/transformation.py +++ b/litellm/llms/gemini/files/transformation.py @@ -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 diff --git a/litellm/llms/vertex_ai/files/handler.py b/litellm/llms/vertex_ai/files/handler.py deleted file mode 100644 index 54c0fa97550..00000000000 --- a/litellm/llms/vertex_ai/files/handler.py +++ /dev/null @@ -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, - ) - ) diff --git a/tests/unit/files/test_main.py b/tests/unit/files/test_main.py index e91951fb5ba..8faf03a3c76 100644 --- a/tests/unit/files/test_main.py +++ b/tests/unit/files/test_main.py @@ -8,7 +8,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 = ( @@ -187,3 +187,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 diff --git a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_handler.py b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_handler.py deleted file mode 100644 index 39310836bf2..00000000000 --- a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_handler.py +++ /dev/null @@ -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"}' diff --git a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_integration.py b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_integration.py index 402c463d134..ee3e79ef1ab 100644 --- a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_integration.py +++ b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_integration.py @@ -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"""