diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 903d79ffa46..7878c5467d2 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/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( 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 e957f32aca1..120d5b48bcd 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 @@ -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"