scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 124

feat(vertex-ai): transform batch prediction outputs to OpenAI format

Browse source 
- 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
This commit is contained in:
Sameer Kankute 2026-04-13 16:42:25 +05:30
parent 5544803b35
commit f6d5502faa
No known key found for this signature in database
3 changed files with 571 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

29
litellm/llms/vertex_ai/files/handler.py
View file

@ -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,

245
litellm/llms/vertex_ai/files/transformation.py
View file

@ -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": {<OpenAI chat completion response>}
},
"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

300
tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_files_transformation.py
View file

@ -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