From a696ffe4a69d3461930001f985ca5937e7ede33e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Felipe=20Gar=C3=A9?= <90070734+FelipeRodriguesGare@users.noreply.github.com> Date: Fri, 19 Sep 2025 19:22:52 -0300 Subject: [PATCH] Litellm gemini batch (#14733) * feat: add Vertex AI support for file content retrieval - Extended `custom_llm_provider` to include "vertex_ai" in `afile_content` function. - Implemented file content retrieval logic for Vertex AI in `VertexAIFilesHandler`. - Added helper method to extract bucket and object from URL-encoded file_id. - Created comprehensive unit and integration tests for Vertex AI file handling. - Updated transformation logic to ensure compatibility with Vertex AI file responses. * fix: update Vertex AI file transformation logic - Modified the transformation logic in `VertexAIFilesConfig` to return a newline-separated JSON string for batch JSONL files instead of a array if JSON strings. * fix: enhance Vertex AI output handling in transformation logic - Updated the transformation logic in `VertexAIBatchTransformation` to utilize the new `OutputInfo` TypedDict for retrieving the GCS output directory. - Added `OutputInfo` class to type definitions for better structure and clarity in Vertex AI responses. --- litellm/files/main.py | 28 +- .../llms/vertex_ai/batches/transformation.py | 9 +- litellm/llms/vertex_ai/files/handler.py | 143 +++++++++- .../llms/vertex_ai/files/transformation.py | 6 +- litellm/types/llms/vertex_ai.py | 5 + .../files/test_vertex_ai_files_handler.py | 270 ++++++++++++++++++ .../files/test_vertex_ai_files_integration.py | 225 +++++++++++++++ 7 files changed, 679 insertions(+), 7 deletions(-) create mode 100644 tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py create mode 100644 tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_integration.py diff --git a/litellm/files/main.py b/litellm/files/main.py index 299e52895b..18be2c702b 100644 --- a/litellm/files/main.py +++ b/litellm/files/main.py @@ -731,7 +731,7 @@ def file_list( async def afile_content( file_id: str, - custom_llm_provider: Literal["openai", "azure"] = "openai", + custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai", extra_headers: Optional[Dict[str, str]] = None, extra_body: Optional[Dict[str, str]] = None, **kwargs, @@ -887,6 +887,32 @@ 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, + ) else: raise litellm.exceptions.BadRequestError( message="LiteLLM doesn't support {} for 'custom_llm_provider'. Supported providers are 'openai', 'azure', 'vertex_ai'.".format( diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py index a97f312d48..5b6d21b594 100644 --- a/litellm/llms/vertex_ai/batches/transformation.py +++ b/litellm/llms/vertex_ai/batches/transformation.py @@ -114,7 +114,14 @@ class VertexAIBatchTransformation: """ Gets the output file id from the Vertex AI Batch response """ - output_file_id: str = "" + + output_file_id: str = ( + response.get("outputInfo", OutputInfo()).get("gcsOutputDirectory", "") + + "/predictions.jsonl" + ) + if output_file_id != "/predictions.jsonl": + return output_file_id + output_config = response.get("outputConfig") if output_config is None: return output_file_id diff --git a/litellm/llms/vertex_ai/files/handler.py b/litellm/llms/vertex_ai/files/handler.py index a666a2c37f..6636bccd6a 100644 --- a/litellm/llms/vertex_ai/files/handler.py +++ b/litellm/llms/vertex_ai/files/handler.py @@ -1,5 +1,6 @@ import asyncio -from typing import Any, Coroutine, Optional, Union +import urllib.parse +from typing import Any, Coroutine, Optional, Tuple, Union import httpx @@ -9,7 +10,12 @@ from litellm.integrations.gcs_bucket.gcs_bucket_base import ( GCSLoggingConfig, ) from litellm.llms.custom_httpx.http_handler import get_async_httpx_client -from litellm.types.llms.openai import CreateFileRequest, OpenAIFileObject +from litellm.types.llms.openai import ( + CreateFileRequest, + FileContentRequest, + HttpxBinaryResponseContent, + OpenAIFileObject, +) from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES from .transformation import VertexAIJsonlFilesTransformation @@ -105,3 +111,136 @@ class VertexAIFilesHandler(GCSBucketBase): max_retries=max_retries, ) ) + + def _extract_bucket_and_object_from_file_id(self, file_id: str) -> Tuple[str, str]: + """ + Extract bucket name and object path from URL-encoded file_id. + + Expected format: gs%3A%2F%2Fbucket-name%2Fpath%2Fto%2Ffile + Which decodes to: gs://bucket-name/path/to/file + + Returns: + tuple: (bucket_name, url_encoded_object_path) + - bucket_name: "bucket-name" + - url_encoded_object_path: "path%2Fto%2Ffile" + """ + decoded_path = urllib.parse.unquote(file_id) + + if decoded_path.startswith("gs://"): + full_path = decoded_path[5:] # Remove 'gs://' prefix + else: + full_path = decoded_path + + if "/" in full_path: + bucket_name, object_path = full_path.split("/", 1) + else: + bucket_name = full_path + object_path = "" + + encoded_object_path = urllib.parse.quote(object_path, safe="") + + return bucket_name, encoded_object_path + + 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], + ) -> 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") + + bucket_name, encoded_object_path = self._extract_bucket_and_object_from_file_id( + file_id + ) + + download_kwargs = { + "standard_callback_dynamic_params": {"gcs_bucket_name": bucket_name} + } + + file_content = await self.download_gcs_object( + object_name=encoded_object_path, **download_kwargs + ) + + if file_content is None: + decoded_path = urllib.parse.unquote(file_id) + raise ValueError(f"Failed to download file from GCS: {decoded_path}") + + decoded_path = urllib.parse.unquote(file_id) + mock_response = httpx.Response( + status_code=200, + content=file_content, + headers={"content-type": "application/octet-stream"}, + request=httpx.Request(method="GET", url=decoded_path), + ) + + return HttpxBinaryResponseContent(response=mock_response) + + 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], + ) -> 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, + ) + 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, + ) + ) diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 1066a8f036..f2e5a5b5d2 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -261,10 +261,10 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): raise ValueError("file is required") extracted_file_data = extract_file_data(file_data) extracted_file_data_content = extracted_file_data.get("content") - + if extracted_file_data_content is None: raise ValueError("file content is required") - + if FilesAPIUtils.is_batch_jsonl_file( create_file_data=create_file_data, extracted_file_data=extracted_file_data, @@ -283,7 +283,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): openai_jsonl_content ) ) - return json.dumps(vertex_jsonl_content) + return "\n".join(json.dumps(item) for item in vertex_jsonl_content) elif isinstance(extracted_file_data_content, bytes): return extracted_file_data_content else: diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index f17a284ddf..854d37522c 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -553,6 +553,10 @@ class OutputConfig(TypedDict, total=False): gcsDestination: GcsDestination +class OutputInfo(TypedDict, total=False): + gcsOutputDirectory: str + + class GcsBucketResponse(TypedDict): """ TypedDict for GCS bucket upload response @@ -611,6 +615,7 @@ class VertexBatchPredictionResponse(TypedDict, total=False): model: str inputConfig: InputConfig outputConfig: OutputConfig + outputInfo: OutputInfo state: str createTime: str updateTime: str diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py new file mode 100644 index 0000000000..ea056a2fc8 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_handler.py @@ -0,0 +1,270 @@ +""" +Test Vertex AI files handler functionality +""" + +import asyncio +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 + + +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%2Ftest-folder" "%2Fsub-folder%2Ftest-file.txt" + + bucket_name, encoded_object_path = ( + self.handler._extract_bucket_and_object_from_file_id(file_id) + ) + + # Verify bucket name extraction + assert bucket_name == "test-bucket" + + # Verify object path encoding + expected_encoded_object = "test-folder%2Fsub-folder%2Ftest-file.txt" + assert encoded_object_path == expected_encoded_object + + def test_extract_bucket_and_object_from_file_id_bucket_only(self): + """Test extraction when only bucket name is provided""" + file_id = "gs%3A%2F%2Ftest-bucket" + + bucket_name, encoded_object_path = ( + self.handler._extract_bucket_and_object_from_file_id(file_id) + ) + + assert bucket_name == "test-bucket" + assert encoded_object_path == "" + + def test_extract_bucket_and_object_from_file_id_simple_path(self): + """Test extraction with simple path""" + file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt" + + bucket_name, encoded_object_path = ( + self.handler._extract_bucket_and_object_from_file_id(file_id) + ) + + assert bucket_name == "test-bucket" + assert encoded_object_path == "test-file.txt" + + def test_extract_bucket_and_object_from_file_id_no_gs_prefix(self): + """Test extraction when gs:// prefix is missing""" + file_id = "test-bucket%2Ftest-file.txt" + + bucket_name, encoded_object_path = ( + self.handler._extract_bucket_and_object_from_file_id(file_id) + ) + + assert bucket_name == "test-bucket" + assert encoded_object_path == "test-file.txt" + + @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%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 download_gcs_object method + with patch.object( + self.handler, "download_gcs_object", new_callable=AsyncMock + ) as mock_download: + 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"] == "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%2Ftest-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: + 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/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_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") diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_integration.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_integration.py new file mode 100644 index 0000000000..723594dc39 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_integration.py @@ -0,0 +1,225 @@ +""" +Test Vertex AI files integration with main files API +""" + +import pytest +from unittest.mock import AsyncMock, patch + +import litellm +from litellm.types.llms.openai import HttpxBinaryResponseContent + + +class TestVertexAIFilesIntegration: + """Test integration of Vertex AI files with main litellm API""" + + @pytest.mark.asyncio + async def test_litellm_afile_content_vertex_ai_provider(self): + """Test litellm.afile_content with vertex_ai provider""" + file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt" + expected_content = b"test file content" + + # Mock the vertex_ai_files_instance.file_content method + with patch( + "litellm.files.main.vertex_ai_files_instance.file_content", + new_callable=AsyncMock, + ) as mock_file_content: + # Create a mock HttpxBinaryResponseContent response + import httpx + + 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_file_content.return_value = HttpxBinaryResponseContent( + response=mock_response + ) + + # Call litellm.afile_content + result = await litellm.afile_content( + file_id=file_id, + custom_llm_provider="vertex_ai", + vertex_project="test-project", + vertex_location="us-central1", + vertex_credentials=None, + ) + + # Verify the result + assert isinstance(result, HttpxBinaryResponseContent) + assert result.response.content == expected_content + assert result.response.status_code == 200 + + # Verify the mock was called with correct parameters + mock_file_content.assert_called_once() + call_kwargs = mock_file_content.call_args.kwargs + assert call_kwargs["_is_async"] is True + assert call_kwargs["file_content_request"]["file_id"] == file_id + assert call_kwargs["vertex_project"] == "test-project" + assert call_kwargs["vertex_location"] == "us-central1" + + def test_litellm_file_content_vertex_ai_provider(self): + """Test litellm.file_content with vertex_ai provider (sync)""" + file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt" + expected_content = b"test file content" + + # Mock the vertex_ai_files_instance.file_content method + with patch( + "litellm.files.main.vertex_ai_files_instance.file_content" + ) as mock_file_content: + # Create a mock HttpxBinaryResponseContent response + import httpx + + 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_file_content.return_value = HttpxBinaryResponseContent( + response=mock_response + ) + + # Call litellm.file_content + result = litellm.file_content( + file_id=file_id, + custom_llm_provider="vertex_ai", + vertex_project="test-project", + vertex_location="us-central1", + vertex_credentials=None, + ) + + # Verify the result + assert isinstance(result, HttpxBinaryResponseContent) + assert result.response.content == expected_content + assert result.response.status_code == 200 + + # Verify the mock was called with correct parameters + mock_file_content.assert_called_once() + call_kwargs = mock_file_content.call_args.kwargs + assert call_kwargs["_is_async"] is False + assert call_kwargs["file_content_request"]["file_id"] == file_id + assert call_kwargs["vertex_project"] == "test-project" + assert call_kwargs["vertex_location"] == "us-central1" + + def test_litellm_file_content_vertex_ai_with_model_provider_detection(self): + """Test litellm.file_content with model parameter for provider detection""" + file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt" + expected_content = b"test file content" + + # Mock the vertex_ai_files_instance.file_content method + with patch( + "litellm.files.main.vertex_ai_files_instance.file_content" + ) as mock_file_content: + # Mock get_llm_provider to return vertex_ai + with patch("litellm.files.main.get_llm_provider") as mock_get_provider: + mock_get_provider.return_value = ( + "vertex_ai/gemini-pro", + "vertex_ai", + None, + None, + ) + + # Create a mock HttpxBinaryResponseContent response + import httpx + + 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_file_content.return_value = HttpxBinaryResponseContent( + response=mock_response + ) + + # Call litellm.file_content with model to trigger provider detection + result = litellm.file_content( + file_id=file_id, + model="vertex_ai/gemini-pro", # This should trigger provider detection + vertex_project="test-project", + vertex_location="us-central1", + ) + + # Verify the result + assert isinstance(result, HttpxBinaryResponseContent) + assert result.response.content == expected_content + + # Verify provider detection was called + mock_get_provider.assert_called_once() + + def test_litellm_file_content_vertex_ai_error_cases(self): + """Test error handling in vertex_ai file_content""" + # Test missing file_id + 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", + ) + + def test_vertex_ai_provider_in_supported_providers_list(self): + """Test that vertex_ai is included in supported providers for file_content""" + # This test ensures the type annotations and error messages include vertex_ai + + # Test that calling with unsupported provider raises appropriate error + with pytest.raises(Exception) as exc_info: + litellm.file_content( + file_id="test-file-id", + custom_llm_provider="unsupported_provider", # This should fail + ) + + # The error message should mention supported providers including vertex_ai + error_message = str(exc_info.value) + assert "vertex_ai" in error_message or "supported" in error_message.lower() + + @pytest.mark.asyncio + async def test_vertex_ai_file_content_with_timeout_and_retries(self): + """Test vertex_ai file_content with timeout and retry configuration""" + file_id = "gs%3A%2F%2Ftest-bucket%2Ftest-file.txt" + expected_content = b"test file content" + + # Mock the vertex_ai_files_instance.file_content method + with patch( + "litellm.files.main.vertex_ai_files_instance.file_content", + new_callable=AsyncMock, + ) as mock_file_content: + # Create a mock HttpxBinaryResponseContent response + import httpx + + 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_file_content.return_value = HttpxBinaryResponseContent( + response=mock_response + ) + + # Call with custom timeout and max_retries + result = await litellm.afile_content( + file_id=file_id, + custom_llm_provider="vertex_ai", + vertex_project="test-project", + vertex_location="us-central1", + timeout=120, + max_retries=5, + ) + + # Verify the result + assert isinstance(result, HttpxBinaryResponseContent) + assert result.response.content == expected_content + + # Verify the timeout and max_retries were passed through + call_kwargs = mock_file_content.call_args.kwargs + assert call_kwargs["timeout"] == 120 + assert call_kwargs["max_retries"] == 5