Fix Vertex batch output transformation detection

litellm/llms/vertex_ai/files/transformation.py

@@ -633,13 +633,31 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
                 and "processed_time" in first_line
             )
             has_success_or_error = (
-                "candidates" in first_line.get("response", {}) or "status" in first_line
+                "candidates" in first_line.get("response", {})
+                or bool(first_line.get("status"))
             )
 
             if not (has_base_structure and has_success_or_error):
                 # Not a Vertex AI batch output, return as-is
                 return content
 
+            vertex_gemini_config = VertexGeminiConfig()
+            logging_obj = Logging(
+                model="",
+                messages=[],
+                stream=False,
+                call_type="batch_transform",
+                start_time=time.time(),
+                litellm_call_id="",
+                function_id="",
+            )
+            logging_obj.optional_params = {}
+            mock_httpx_response = httpx.Response(
+                status_code=200,
+                headers={"content-type": "application/json"},
+                request=httpx.Request(method="POST", url="https://example.com"),
+            )
+
             # Transform all lines
             transformed_lines = []
             for line in lines:
@@ -650,7 +668,10 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
                     vertex_output = json.loads(line)
                     openai_output = (
                         self._transform_single_vertex_batch_output_to_openai(
-                            vertex_output
+                            vertex_output=vertex_output,
+                            vertex_gemini_config=vertex_gemini_config,
+                            logging_obj=logging_obj,
+                            mock_httpx_response=mock_httpx_response,
                         )
                     )
                     transformed_lines.append(json.dumps(openai_output))
@@ -666,7 +687,11 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
             return content
 
     def _transform_single_vertex_batch_output_to_openai(
-        self, vertex_output: Dict[str, Any]
+        self,
+        vertex_output: Dict[str, Any],
+        vertex_gemini_config: VertexGeminiConfig,
+        logging_obj: Logging,
+        mock_httpx_response: httpx.Response,
     ) -> Dict[str, Any]:
         """
         Transform a single Vertex AI batch output line to OpenAI format.
@@ -712,29 +737,11 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
         if "@" in model:
             model = model.split("@")[0]
 
-        # Create logging object for transformation
-        logging_obj = Logging(
-            model=model,
-            messages=[],
-            stream=False,
-            call_type="batch_transform",
-            start_time=time.time(),
-            litellm_call_id="",
-            function_id="",
-        )
-        logging_obj.optional_params = {}
-
-        # Create mock httpx response for transformation
-        mock_httpx_response = httpx.Response(
-            status_code=200,
-            content=json.dumps(vertex_response).encode("utf-8"),
-            headers={"content-type": "application/json"},
-            request=httpx.Request(method="POST", url="https://example.com"),
-        )
+        logging_obj.model = model
+        logging_obj.start_time = time.time()
 
         try:
             # Use existing VertexGeminiConfig transformation
-            vertex_gemini_config = VertexGeminiConfig()
             model_response = ModelResponse()
 
             transformed_response = vertex_gemini_config._transform_google_generate_content_to_openai_model_response(

tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py

@@ -449,6 +449,66 @@ class TestVertexBatchOutputTransformation:
             assert "choices" in body
             assert len(body["choices"]) > 0
 
+    def test_batch_detection_requires_candidates_or_non_empty_status(self, config):
+        """Test that JSONL with a blank status but no candidates is returned as-is."""
+        non_batch_output = {
+            "status": "",
+            "processed_time": "2024-11-01T18:13:16.826+00:00",
+            "request": {"metadata": "not a Vertex batch request"},
+            "response": {"metadata": "not a Gemini response"},
+        }
+
+        content = json.dumps(non_batch_output).encode("utf-8")
+        transformed_content = config._try_transform_vertex_batch_output_to_openai(
+            content
+        )
+
+        assert transformed_content == content
+
+    def test_reuses_batch_transform_helpers_per_jsonl_file(self, config, monkeypatch):
+        """Test that heavy helper objects are reused while transforming a JSONL file."""
+        vertex_outputs = [
+            {
+                "status": "",
+                "processed_time": "2024-11-01T18:13:16.826+00:00",
+                "request": {"labels": {"litellm_custom_id": f"request-{i}"}},
+                "response": {"candidates": [{"content": {"parts": [{"text": "ok"}]}}]},
+            }
+            for i in range(2)
+        ]
+        helper_ids = []
+
+        def mock_transform_single(
+            vertex_output,
+            vertex_gemini_config,
+            logging_obj,
+            mock_httpx_response,
+        ):
+            helper_ids.append(
+                (
+                    id(vertex_gemini_config),
+                    id(logging_obj),
+                    id(mock_httpx_response),
+                )
+            )
+            return {"custom_id": vertex_output["request"]["labels"]["litellm_custom_id"]}
+
+        monkeypatch.setattr(
+            config,
+            "_transform_single_vertex_batch_output_to_openai",
+            mock_transform_single,
+        )
+
+        content = "\n".join(json.dumps(output) for output in vertex_outputs).encode(
+            "utf-8"
+        )
+        transformed_content = config._try_transform_vertex_batch_output_to_openai(
+            content
+        )
+
+        assert len(transformed_content.decode("utf-8").strip().split("\n")) == 2
+        assert len(set(helper_ids)) == 1
+
     def test_non_batch_output_passthrough(self, config):
         """Test that non-batch output is returned as-is"""
         regular_content = b"This is just a regular file content"