From f6d5502faa59e56fb4bc38227153296fb9fb4c49 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Mon, 13 Apr 2026 16:42:25 +0530 Subject: [PATCH] feat(vertex-ai): transform batch prediction outputs to OpenAI format - Add automatic conversion of Vertex AI batch prediction JSONL to OpenAI format - Preserve custom_id via Vertex AI labels for request correlation - Fix Content-Length header mismatch in transformed responses - Add comprehensive tests for batch output transformation Made-with: Cursor --- litellm/llms/vertex_ai/files/handler.py | 29 +- .../llms/vertex_ai/files/transformation.py | 245 ++++++++++++++ .../test_vertex_ai_files_transformation.py | 300 +++++++++++++++++- 3 files changed, 571 insertions(+), 3 deletions(-) diff --git a/litellm/llms/vertex_ai/files/handler.py b/litellm/llms/vertex_ai/files/handler.py index 6636bccd6a3..81bf7852c82 100644 --- a/litellm/llms/vertex_ai/files/handler.py +++ b/litellm/llms/vertex_ai/files/handler.py @@ -1,4 +1,5 @@ import asyncio +import time import urllib.parse from typing import Any, Coroutine, Optional, Tuple, Union @@ -188,11 +189,35 @@ class VertexAIFilesHandler(GCSBucketBase): mock_response = httpx.Response( status_code=200, content=file_content, - headers={"content-type": "application/octet-stream"}, + headers={ + "content-type": "application/octet-stream", + "content-length": str(len(file_content)), + }, request=httpx.Request(method="GET", url=decoded_path), ) - return HttpxBinaryResponseContent(response=mock_response) + # Apply transformation to convert Vertex AI batch outputs to OpenAI format + from .transformation import VertexAIFilesConfig + from litellm.litellm_core_utils.litellm_logging import Logging + + config = VertexAIFilesConfig() + + # Create a logging object for transformation + logging_obj = 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, diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 070ec508283..a967e968e4a 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -3,6 +3,7 @@ import os import time from typing import Any, Dict, List, Optional, Tuple, Union +import httpx from httpx import Headers, Response from openai.types.file_deleted import FileDeleted @@ -247,6 +248,14 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): litellm_params={}, cached_content=None, ) + + # Add custom_id as a label for correlation in batch outputs + custom_id = _openai_jsonl_content.get("custom_id") + if custom_id: + if "labels" not in vertex_request_body: + vertex_request_body["labels"] = {} + vertex_request_body["labels"]["litellm_custom_id"] = str(custom_id) + vertex_jsonl_content.append({"request": vertex_request_body}) return vertex_jsonl_content @@ -453,7 +462,235 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): logging_obj: LiteLLMLoggingObj, litellm_params: dict, ) -> HttpxBinaryResponseContent: + """ + Transform file content response, converting Vertex AI batch output to OpenAI format if applicable. + + This method automatically detects and transforms Vertex AI batch prediction outputs + (predictions.jsonl files) into OpenAI-compatible batch response format. + + If the file is not a batch output or transformation fails, the original content + is returned as-is to maintain backward compatibility. + """ + try: + # Try to transform batch output if it's a JSONL file + content = raw_response.content + if content: + transformed_content = self._try_transform_vertex_batch_output_to_openai( + content + ) + if transformed_content != content: + # Create a new response with transformed content and updated Content-Length + import httpx + + # Update headers with correct Content-Length + new_headers = dict(raw_response.headers) + new_headers["content-length"] = str(len(transformed_content)) + + mock_response = httpx.Response( + status_code=raw_response.status_code, + content=transformed_content, + headers=new_headers, + request=raw_response.request, + ) + return HttpxBinaryResponseContent(response=mock_response) + except Exception: + # If transformation fails, return as-is + pass + return HttpxBinaryResponseContent(response=raw_response) + + def _try_transform_vertex_batch_output_to_openai( + self, content: bytes + ) -> bytes: + """ + Try to transform Vertex AI batch output to OpenAI format. + If conversion fails at any point, return the original content as-is. + + Vertex AI batch output format (predictions.jsonl): + { + "request": {"contents": [...], "labels": {"litellm_custom_id": "request-1"}}, + "status": "", + "response": {"candidates": [...], "modelVersion": "gemini-2.5-flash", ...}, + "processed_time": "2026-04-13T10:18:18.102004+00:00" + } + + OpenAI batch output format: + { + "id": "batch_req_...", + "custom_id": "request-1", + "response": { + "status_code": 200, + "request_id": "chatcmpl-...", + "body": {} + }, + "error": null + } + """ + try: + # Decode content + content_str = content.decode("utf-8") + + # Check if it's JSONL (multiple lines) + lines = content_str.strip().split("\n") + if not lines: + return content + + # Try to parse the first line to see if it's Vertex AI batch output + first_line = json.loads(lines[0]) + + # Check if it has Vertex AI batch output structure + if not ("response" in first_line and "request" in first_line): + # Not a Vertex AI batch output, return as-is + return content + + # Transform all lines + transformed_lines = [] + for line in lines: + if not line.strip(): + continue + + try: + vertex_output = json.loads(line) + openai_output = self._transform_single_vertex_batch_output_to_openai( + vertex_output + ) + transformed_lines.append(json.dumps(openai_output)) + except Exception: + # If any line fails, return original content + return content + + # Return transformed content + return "\n".join(transformed_lines).encode("utf-8") + + except Exception: + # If anything fails, return original content + return content + + def _transform_single_vertex_batch_output_to_openai( + self, vertex_output: Dict[str, Any] + ) -> Dict[str, Any]: + """ + Transform a single Vertex AI batch output line to OpenAI format. + Uses the existing VertexGeminiConfig transformation for the response. + """ + from litellm.types.utils import ModelResponse + import httpx + import time + + # Extract custom_id from request labels + custom_id = "unknown" + request_data = vertex_output.get("request", {}) + labels = request_data.get("labels", {}) + if "litellm_custom_id" in labels: + custom_id = labels["litellm_custom_id"] + + # Check if there's an error + status = vertex_output.get("status", "") + has_error = bool(status) + + if has_error: + # Return error response in OpenAI format + return { + "id": f"batch_req_{uuid.uuid4()}", + "custom_id": custom_id, + "response": { + "status_code": 400, + "request_id": "", + "body": { + "error": { + "message": status, + "type": "vertex_ai_error", + "code": "vertex_ai_error" + } + } + }, + "error": { + "message": status, + "type": "vertex_ai_error", + "code": "vertex_ai_error" + } + } + + # Transform successful response using existing transformation + vertex_response = vertex_output.get("response", {}) + + # Extract model from response + model = vertex_response.get("modelVersion", "gemini-1.5-flash-001") + if "@" in model: + model = model.split("@")[0] + + # Create logging object for transformation + from litellm.litellm_core_utils.litellm_logging import Logging + + 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"), + ) + + 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( + completion_response=vertex_response, + model_response=model_response, + model=model, + logging_obj=logging_obj, + raw_response=mock_httpx_response, + ) + + # Convert ModelResponse to dict + response_dict = transformed_response.model_dump() + + # Return in OpenAI batch format + return { + "id": f"batch_req_{uuid.uuid4()}", + "custom_id": custom_id, + "response": { + "status_code": 200, + "request_id": response_dict.get("id", ""), + "body": response_dict + }, + "error": None + } + + except Exception as e: + # If transformation fails, return error + return { + "id": f"batch_req_{uuid.uuid4()}", + "custom_id": custom_id, + "response": { + "status_code": 500, + "request_id": "", + "body": { + "error": { + "message": f"Failed to transform response: {str(e)}", + "type": "transformation_error", + "code": "transformation_error" + } + } + }, + "error": { + "message": f"Failed to transform response: {str(e)}", + "type": "transformation_error", + "code": "transformation_error" + } + } class VertexAIJsonlFilesTransformation(VertexGeminiConfig): @@ -513,6 +750,14 @@ class VertexAIJsonlFilesTransformation(VertexGeminiConfig): litellm_params={}, cached_content=None, ) + + # Add custom_id as a label for correlation in batch outputs + custom_id = _openai_jsonl_content.get("custom_id") + if custom_id: + if "labels" not in vertex_request_body: + vertex_request_body["labels"] = {} + vertex_request_body["labels"]["litellm_custom_id"] = str(custom_id) + vertex_jsonl_content.append({"request": vertex_request_body}) return vertex_jsonl_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 598ad255aca..7549eb9eaa6 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 @@ -1,5 +1,6 @@ """ Tests for VertexAIFilesConfig transformation methods (Issues 5-7). +Includes tests for Vertex AI batch output transformation to OpenAI format. """ import json @@ -9,7 +10,10 @@ import httpx import pytest from unittest.mock import MagicMock -from litellm.llms.vertex_ai.files.transformation import VertexAIFilesConfig +from litellm.llms.vertex_ai.files.transformation import ( + VertexAIFilesConfig, + VertexAIJsonlFilesTransformation, +) from litellm.types.llms.openai import OpenAIFileObject, HttpxBinaryResponseContent from openai.types.file_deleted import FileDeleted @@ -228,3 +232,297 @@ class TestTransformDeleteFile: "gs://prod-bucket/litellm-vertex-files/publishers/google/" "models/gemini-2.0-flash-001/abc-123" ) + + +class TestVertexBatchOutputTransformation: + """Test transformation of Vertex AI batch outputs to OpenAI format""" + + def test_transform_successful_vertex_batch_output(self, config): + """Test transformation of a successful Vertex AI batch output""" + # Sample Vertex AI batch output (based on actual format) + vertex_output = { + "status": "", + "processed_time": "2024-11-01T18:13:16.826+00:00", + "request": { + "contents": [{"role": "user", "parts": [{"text": "Hello world!"}]}], + "labels": {"litellm_custom_id": "request-1"} + }, + "response": { + "candidates": [{ + "content": { + "parts": [{"text": "Hello! How can I help you today?"}], + "role": "model" + }, + "finishReason": "STOP" + }], + "modelVersion": "gemini-2.0-flash-001@default", + "usageMetadata": { + "promptTokenCount": 10, + "candidatesTokenCount": 20, + "totalTokenCount": 30 + } + } + } + + content = json.dumps(vertex_output).encode("utf-8") + transformed_content = config._try_transform_vertex_batch_output_to_openai(content) + result = json.loads(transformed_content.decode("utf-8")) + + # Verify OpenAI format + assert "id" in result + assert "custom_id" in result + assert "response" in result + assert "error" in result + + # Verify custom_id was extracted from labels + assert result["custom_id"] == "request-1" + + # Verify response structure + assert result["response"]["status_code"] == 200 + assert "body" in result["response"] + + # Verify body has OpenAI format + body = result["response"]["body"] + assert "choices" in body + assert "usage" in body + assert "model" in body + + # Verify choices + assert len(body["choices"]) > 0 + choice = body["choices"][0] + assert "message" in choice + assert "content" in choice["message"] + assert "Hello! How can I help you today?" in choice["message"]["content"] + + def test_transform_error_vertex_batch_output(self, config): + """Test transformation of an error Vertex AI batch output""" + vertex_output = { + "status": "Error: Invalid request", + "processed_time": "2024-11-01T18:13:16.826+00:00", + "request": { + "contents": [{"role": "user", "parts": [{"text": "Hello world!"}]}], + "labels": {"litellm_custom_id": "request-error"} + }, + "response": {} + } + + content = json.dumps(vertex_output).encode("utf-8") + transformed_content = config._try_transform_vertex_batch_output_to_openai(content) + result = json.loads(transformed_content.decode("utf-8")) + + # Verify error format + assert result["response"]["status_code"] == 400 + assert result["error"] is not None + assert "Invalid request" in result["error"]["message"] + assert result["custom_id"] == "request-error" + + def test_transform_multiple_vertex_batch_outputs(self, config): + """Test transformation of multiple Vertex AI batch outputs (JSONL)""" + vertex_outputs = [ + { + "status": "", + "processed_time": "2024-11-01T18:13:16.826+00:00", + "request": { + "contents": [{"role": "user", "parts": [{"text": "First request"}]}], + "labels": {"litellm_custom_id": "request-1"} + }, + "response": { + "candidates": [{ + "content": {"parts": [{"text": "First response"}], "role": "model"}, + "finishReason": "STOP" + }], + "modelVersion": "gemini-2.0-flash-001@default", + "usageMetadata": { + "promptTokenCount": 5, + "candidatesTokenCount": 10, + "totalTokenCount": 15 + } + } + }, + { + "status": "", + "processed_time": "2024-11-01T18:13:17.826+00:00", + "request": { + "contents": [{"role": "user", "parts": [{"text": "Second request"}]}], + "labels": {"litellm_custom_id": "request-2"} + }, + "response": { + "candidates": [{ + "content": {"parts": [{"text": "Second response"}], "role": "model"}, + "finishReason": "STOP" + }], + "modelVersion": "gemini-2.0-flash-001@default", + "usageMetadata": { + "promptTokenCount": 6, + "candidatesTokenCount": 11, + "totalTokenCount": 17 + } + } + } + ] + + 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) + lines = transformed_content.decode("utf-8").strip().split("\n") + + assert len(lines) == 2 + + for i, line in enumerate(lines): + result = json.loads(line) + assert "id" in result + assert "response" in result + assert result["response"]["status_code"] == 200 + assert result["custom_id"] == f"request-{i+1}" + body = result["response"]["body"] + assert "choices" in body + assert len(body["choices"]) > 0 + + 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" + transformed_content = config._try_transform_vertex_batch_output_to_openai(regular_content) + assert transformed_content == regular_content + + def test_invalid_json_passthrough(self, config): + """Test that invalid JSON is returned as-is""" + invalid_content = b'{"invalid": json content}' + transformed_content = config._try_transform_vertex_batch_output_to_openai(invalid_content) + assert transformed_content == invalid_content + + +class TestVertexBatchCustomIdLabels: + """Test custom_id handling in batch transformations""" + + def test_custom_id_added_to_labels_in_vertex_request(self): + """Test that custom_id from OpenAI format is added as a label in Vertex AI format""" + transformation = VertexAIJsonlFilesTransformation() + + openai_jsonl_content = [ + { + "custom_id": "request-1", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "gemini-1.5-flash-001", + "messages": [{"role": "user", "content": "What is 2+2?"}], + "max_tokens": 10 + } + } + ] + + vertex_jsonl_content = transformation._transform_openai_jsonl_content_to_vertex_ai_jsonl_content( + openai_jsonl_content + ) + + assert len(vertex_jsonl_content) == 1 + vertex_request = vertex_jsonl_content[0] + + # Verify labels were added + assert "labels" in vertex_request["request"] + assert "litellm_custom_id" in vertex_request["request"]["labels"] + assert vertex_request["request"]["labels"]["litellm_custom_id"] == "request-1" + + def test_multiple_requests_each_get_their_own_label(self): + """Test that multiple requests each get their own custom_id label""" + transformation = VertexAIJsonlFilesTransformation() + + openai_jsonl_content = [ + { + "custom_id": f"request-{i+1}", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "gemini-1.5-flash-001", + "messages": [{"role": "user", "content": f"Question {i+1}"}], + } + } + for i in range(3) + ] + + vertex_jsonl_content = transformation._transform_openai_jsonl_content_to_vertex_ai_jsonl_content( + openai_jsonl_content + ) + + assert len(vertex_jsonl_content) == 3 + + for i, vertex_request in enumerate(vertex_jsonl_content): + expected_custom_id = f"request-{i+1}" + assert vertex_request["request"]["labels"]["litellm_custom_id"] == expected_custom_id + + def test_request_without_custom_id_has_no_label(self): + """Test that requests without custom_id don't get a label""" + transformation = VertexAIJsonlFilesTransformation() + + openai_jsonl_content = [ + { + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "gemini-1.5-flash-001", + "messages": [{"role": "user", "content": "Question"}], + } + } + ] + + vertex_jsonl_content = transformation._transform_openai_jsonl_content_to_vertex_ai_jsonl_content( + openai_jsonl_content + ) + + # Should not have labels if no custom_id was provided + assert "labels" not in vertex_jsonl_content[0]["request"] + + def test_end_to_end_custom_id_round_trip(self): + """ + Test the full round trip: OpenAI format -> Vertex AI format -> Vertex AI output -> OpenAI output + Verify that custom_id is preserved through the entire flow. + """ + transformation = VertexAIJsonlFilesTransformation() + config = VertexAIFilesConfig() + + # Step 1: Transform OpenAI input to Vertex AI format + openai_input = [ + { + "custom_id": "my-custom-request-id", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "gemini-1.5-flash-001", + "messages": [{"role": "user", "content": "Hello"}], + } + } + ] + + vertex_input = transformation._transform_openai_jsonl_content_to_vertex_ai_jsonl_content( + openai_input + ) + + # Verify label was added + assert vertex_input[0]["request"]["labels"]["litellm_custom_id"] == "my-custom-request-id" + + # Step 2: Simulate Vertex AI batch output (with the label echoed back) + vertex_output = { + "status": "", + "processed_time": "2024-11-01T18:13:16.826+00:00", + "request": vertex_input[0]["request"], + "response": { + "candidates": [{ + "content": {"parts": [{"text": "Hi there!"}], "role": "model"}, + "finishReason": "STOP" + }], + "modelVersion": "gemini-2.0-flash-001@default", + "usageMetadata": { + "promptTokenCount": 5, + "candidatesTokenCount": 10, + "totalTokenCount": 15 + } + } + } + + # Step 3: Transform Vertex AI output back to OpenAI format + content = json.dumps(vertex_output).encode("utf-8") + transformed_content = config._try_transform_vertex_batch_output_to_openai(content) + openai_output = json.loads(transformed_content.decode("utf-8")) + + # Step 4: Verify custom_id was preserved + assert openai_output["custom_id"] == "my-custom-request-id" + assert openai_output["response"]["status_code"] == 200