diff --git a/litellm/__init__.py b/litellm/__init__.py index d02d39cae15..28112f5c12a 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -288,6 +288,7 @@ disable_token_counter: bool = False disable_add_transform_inline_image_block: bool = False disable_add_user_agent_to_request_tags: bool = False disable_anthropic_gemini_context_caching_transform: bool = False +disable_vertex_batch_output_transformation: bool = False extra_spend_tag_headers: Optional[List[str]] = None in_memory_llm_clients_cache: "LLMClientCache" safe_memory_mode: bool = False diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 02ef2ab6575..7df16723870 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -9,6 +9,7 @@ import httpx from httpx import Headers, Response from openai.types.file_deleted import FileDeleted +import litellm from litellm._uuid import uuid from litellm.files.utils import FilesAPIUtils from litellm.litellm_core_utils.litellm_logging import Logging @@ -573,6 +574,11 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): is returned as-is to maintain backward compatibility. """ try: + # Allow users to opt out of automatic Vertex batch output -> OpenAI + # transformation, e.g. if they consume raw `predictions.jsonl` directly. + if getattr(litellm, "disable_vertex_batch_output_transformation", False): + return HttpxBinaryResponseContent(response=raw_response) + # Try to transform batch output if it's a JSONL file content = raw_response.content if content: diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py index 6537fb6a946..4adad53a98f 100644 --- a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -210,6 +210,47 @@ class TestTransformFileContent: assert logging_obj.optional_params == original_optional_params assert result.response is not raw_response + def test_should_skip_batch_output_transformation_when_opt_out_flag_set( + self, config, monkeypatch + ): + """When `litellm.disable_vertex_batch_output_transformation` is True the + Vertex predictions.jsonl content must be returned untouched, so callers + that parse raw `candidates`/`modelVersion` keep working.""" + import litellm + + raw_jsonl = json.dumps( + { + "status": "", + "processed_time": "2024-11-01T18:13:16.826+00:00", + "request": {"labels": {"litellm_custom_id": "request-1"}}, + "response": { + "candidates": [ + {"content": {"parts": [{"text": "ok"}], "role": "model"}} + ], + "modelVersion": "gemini-2.0-flash-001@default", + }, + } + ).encode("utf-8") + raw_response = httpx.Response( + status_code=200, + content=raw_jsonl, + headers={"content-type": "application/octet-stream"}, + request=httpx.Request("GET", "https://example.com"), + ) + + monkeypatch.setattr( + litellm, "disable_vertex_batch_output_transformation", True, raising=False + ) + + result = config.transform_file_content_response( + raw_response=raw_response, + logging_obj=MagicMock(), + litellm_params={}, + ) + + assert isinstance(result, HttpxBinaryResponseContent) + assert result.response.content == raw_jsonl + class TestTransformDeleteFile: def test_should_build_correct_gcs_delete_url(self, config):