mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
Merge pull request #25627 from BerriAI/litellm_vertex-batch-output-transformation
feat(vertex-ai): transform batch prediction outputs to OpenAI format
This commit is contained in:
Mateo Wang
GitHub
commit
6dd04357f6
4 changed files with 1193 additions and 51 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
import asyncio
|
||||
import time
|
||||
import urllib.parse
|
||||
from typing import Any, Coroutine, Optional, Tuple, Union
|
||||
|
||||
|
|
@ -16,9 +17,10 @@ from litellm.types.llms.openai import (
|
|||
HttpxBinaryResponseContent,
|
||||
OpenAIFileObject,
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES
|
||||
|
||||
from .transformation import VertexAIJsonlFilesTransformation
|
||||
from .transformation import VertexAIFilesConfig, VertexAIJsonlFilesTransformation
|
||||
|
||||
vertex_ai_files_transformation = VertexAIJsonlFilesTransformation()
|
||||
|
||||
|
|
@ -188,11 +190,30 @@ 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
|
||||
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,
|
||||
|
|
|
|||
|
|
@ -1,13 +1,18 @@
|
|||
import base64
|
||||
import json
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from typing import Any, Dict, List, Optional, Tuple, Union
|
||||
from typing import Any, Callable, Dict, List, Optional, Tuple, Union
|
||||
|
||||
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
|
||||
from litellm.litellm_core_utils.prompt_templates.common_utils import extract_file_data
|
||||
from litellm.llms.base_llm.chat.transformation import BaseLLMException
|
||||
from litellm.llms.base_llm.files.transformation import (
|
||||
|
|
@ -31,11 +36,135 @@ from litellm.types.llms.openai import (
|
|||
PathLike,
|
||||
)
|
||||
from litellm.types.llms.vertex_ai import GcsBucketResponse
|
||||
from litellm.types.utils import ExtractedFileData, LlmProviders
|
||||
from litellm.types.utils import ExtractedFileData, LlmProviders, ModelResponse
|
||||
|
||||
from ..common_utils import VertexAIError
|
||||
from ..vertex_llm_base import VertexBase
|
||||
|
||||
_GCP_LABEL_VALUE_MAX_LEN = 63
|
||||
_CUSTOM_ID_RAW_LABEL_PREFIX = "b32_"
|
||||
|
||||
|
||||
def _sanitize_gcp_label_value(value: str) -> str:
|
||||
"""
|
||||
Sanitize a string to meet GCP label value constraints.
|
||||
|
||||
GCP label values must:
|
||||
- Be lowercase
|
||||
- Contain only letters, numbers, underscores, and hyphens
|
||||
- Be max 63 characters
|
||||
|
||||
Args:
|
||||
value: The string to sanitize
|
||||
|
||||
Returns:
|
||||
A sanitized string that meets GCP label constraints
|
||||
"""
|
||||
sanitized = re.sub(r"[^a-z0-9_-]", "_", value.lower())
|
||||
return sanitized[:_GCP_LABEL_VALUE_MAX_LEN]
|
||||
|
||||
|
||||
def _encode_gcp_label_value_chunks(value: str) -> List[str]:
|
||||
"""Encode arbitrary text across one or more GCP-label-safe values."""
|
||||
max_encoded_len = _GCP_LABEL_VALUE_MAX_LEN - len(_CUSTOM_ID_RAW_LABEL_PREFIX)
|
||||
encoded = (
|
||||
base64.b32encode(value.encode("utf-8")).decode("ascii").rstrip("=").lower()
|
||||
)
|
||||
return [
|
||||
f"{_CUSTOM_ID_RAW_LABEL_PREFIX}{encoded[i : i + max_encoded_len]}"
|
||||
for i in range(0, len(encoded), max_encoded_len)
|
||||
] or [_CUSTOM_ID_RAW_LABEL_PREFIX]
|
||||
|
||||
|
||||
def _decode_gcp_label_value_chunks(values: List[str]) -> Optional[str]:
|
||||
"""Decode values produced by _encode_gcp_label_value_chunks."""
|
||||
encoded_parts = []
|
||||
for value in values:
|
||||
if not value.startswith(_CUSTOM_ID_RAW_LABEL_PREFIX):
|
||||
return None
|
||||
encoded_parts.append(value[len(_CUSTOM_ID_RAW_LABEL_PREFIX) :])
|
||||
encoded = "".join(encoded_parts).upper()
|
||||
padding = "=" * (-len(encoded) % 8)
|
||||
try:
|
||||
return base64.b32decode(encoded + padding).decode("utf-8")
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
def _set_litellm_batch_custom_id_labels(labels: Dict[str, str], custom_id: Any) -> None:
|
||||
"""
|
||||
Store OpenAI batch custom_id for Vertex batch correlation.
|
||||
|
||||
``litellm_custom_id`` is GCP-label-safe (may alter casing and characters).
|
||||
``litellm_custom_id_raw`` encodes the original string for
|
||||
round-trip correlation in batch output transforms.
|
||||
"""
|
||||
custom_id_str = str(custom_id)
|
||||
labels["litellm_custom_id"] = _sanitize_gcp_label_value(custom_id_str)
|
||||
raw_label_chunks = _encode_gcp_label_value_chunks(custom_id_str)
|
||||
labels["litellm_custom_id_raw"] = raw_label_chunks[0]
|
||||
for index, raw_label_chunk in enumerate(raw_label_chunks[1:], start=1):
|
||||
labels[f"litellm_custom_id_raw_{index}"] = raw_label_chunk
|
||||
|
||||
|
||||
def _get_litellm_batch_custom_id_from_labels(labels: Dict[str, Any]) -> str:
|
||||
"""Prefer encoded custom_id when present (see _set_litellm_batch_custom_id_labels)."""
|
||||
raw = labels.get("litellm_custom_id_raw")
|
||||
if raw:
|
||||
raw_chunks = [str(raw)]
|
||||
chunk_prefix = "litellm_custom_id_raw_"
|
||||
indexed_chunks = []
|
||||
for key, value in labels.items():
|
||||
if key.startswith(chunk_prefix) and key[len(chunk_prefix) :].isdigit():
|
||||
indexed_chunks.append((int(key[len(chunk_prefix) :]), str(value)))
|
||||
raw_chunks.extend(
|
||||
raw_label_chunk
|
||||
for _, raw_label_chunk in sorted(indexed_chunks, key=lambda item: item[0])
|
||||
)
|
||||
decoded = _decode_gcp_label_value_chunks(raw_chunks)
|
||||
if decoded is not None:
|
||||
return decoded
|
||||
return str(raw)
|
||||
return str(labels.get("litellm_custom_id", "unknown"))
|
||||
|
||||
|
||||
def _openai_batch_jsonl_entries_to_vertex_wrapped_requests(
|
||||
openai_jsonl_content: List[Dict[str, Any]],
|
||||
map_openai_to_vertex_params: Callable[[Dict[str, Any]], Dict[str, Any]],
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Transforms OpenAI JSONL batch entries to Vertex AI JSONL lines.
|
||||
|
||||
jsonl body for vertex is {"request": <request_body>}
|
||||
Example Vertex jsonl
|
||||
{"request":{"contents": [{"role": "user", "parts": [{"text": "What is the relation between the following video and image samples?"}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/video/animals.mp4", "mimeType": "video/mp4"}}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/image/cricket.jpeg", "mimeType": "image/jpeg"}}]}]}}
|
||||
{"request":{"contents": [{"role": "user", "parts": [{"text": "Describe what is happening in this video."}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/video/another_video.mov", "mimeType": "video/mov"}}]}]}}
|
||||
"""
|
||||
|
||||
vertex_jsonl_content = []
|
||||
for _openai_jsonl_content in openai_jsonl_content:
|
||||
openai_request_body = _openai_jsonl_content.get("body") or {}
|
||||
vertex_request_body = _transform_request_body(
|
||||
messages=openai_request_body.get("messages", []),
|
||||
model=openai_request_body.get("model", ""),
|
||||
optional_params=map_openai_to_vertex_params(openai_request_body),
|
||||
custom_llm_provider="vertex_ai",
|
||||
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 is not None:
|
||||
if "labels" not in vertex_request_body:
|
||||
vertex_request_body["labels"] = {}
|
||||
_set_litellm_batch_custom_id_labels(
|
||||
vertex_request_body["labels"], custom_id
|
||||
)
|
||||
|
||||
vertex_jsonl_content.append({"request": vertex_request_body})
|
||||
return vertex_jsonl_content
|
||||
|
||||
|
||||
class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
|
||||
"""
|
||||
|
|
@ -227,28 +356,10 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
|
|||
def _transform_openai_jsonl_content_to_vertex_ai_jsonl_content(
|
||||
self, openai_jsonl_content: List[Dict[str, Any]]
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Transforms OpenAI JSONL content to VertexAI JSONL content
|
||||
|
||||
jsonl body for vertex is {"request": <request_body>}
|
||||
Example Vertex jsonl
|
||||
{"request":{"contents": [{"role": "user", "parts": [{"text": "What is the relation between the following video and image samples?"}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/video/animals.mp4", "mimeType": "video/mp4"}}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/image/cricket.jpeg", "mimeType": "image/jpeg"}}]}]}}
|
||||
{"request":{"contents": [{"role": "user", "parts": [{"text": "Describe what is happening in this video."}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/video/another_video.mov", "mimeType": "video/mov"}}]}]}}
|
||||
"""
|
||||
|
||||
vertex_jsonl_content = []
|
||||
for _openai_jsonl_content in openai_jsonl_content:
|
||||
openai_request_body = _openai_jsonl_content.get("body") or {}
|
||||
vertex_request_body = _transform_request_body(
|
||||
messages=openai_request_body.get("messages", []),
|
||||
model=openai_request_body.get("model", ""),
|
||||
optional_params=self._map_openai_to_vertex_params(openai_request_body),
|
||||
custom_llm_provider="vertex_ai",
|
||||
litellm_params={},
|
||||
cached_content=None,
|
||||
)
|
||||
vertex_jsonl_content.append({"request": vertex_request_body})
|
||||
return vertex_jsonl_content
|
||||
return _openai_batch_jsonl_entries_to_vertex_wrapped_requests(
|
||||
openai_jsonl_content=openai_jsonl_content,
|
||||
map_openai_to_vertex_params=self._map_openai_to_vertex_params,
|
||||
)
|
||||
|
||||
def transform_create_file_request(
|
||||
self,
|
||||
|
|
@ -453,8 +564,253 @@ 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:
|
||||
# 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:
|
||||
transformed_content = self._try_transform_vertex_batch_output_to_openai(
|
||||
content=content,
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
if transformed_content != content:
|
||||
# Create a new response with transformed content and updated Content-Length
|
||||
# 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, logging_obj: Optional[LiteLLMLoggingObj] = None
|
||||
) -> 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", "litellm_custom_id_raw": "..."}},
|
||||
"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 with discriminating fields
|
||||
# Must have request, response, and processed_time
|
||||
# Plus either candidates (success) or status (error)
|
||||
has_base_structure = (
|
||||
"response" in first_line
|
||||
and "request" in first_line
|
||||
and "processed_time" in first_line
|
||||
)
|
||||
has_success_or_error = (
|
||||
"candidates" in first_line.get("response", {})
|
||||
or "promptFeedback" 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()
|
||||
# Always use a fresh local Logging object for the per-line transformation
|
||||
# so we never mutate the caller's logging_obj (which already went through
|
||||
# pre_call and has its own model/start_time/optional_params set).
|
||||
batch_transform_logging_obj = Logging(
|
||||
model="",
|
||||
messages=[],
|
||||
stream=False,
|
||||
call_type="batch_transform",
|
||||
start_time=time.time(),
|
||||
litellm_call_id="",
|
||||
function_id="",
|
||||
)
|
||||
batch_transform_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:
|
||||
if not line.strip():
|
||||
continue
|
||||
|
||||
try:
|
||||
vertex_output = json.loads(line)
|
||||
openai_output = (
|
||||
self._transform_single_vertex_batch_output_to_openai(
|
||||
vertex_output=vertex_output,
|
||||
vertex_gemini_config=vertex_gemini_config,
|
||||
logging_obj=batch_transform_logging_obj,
|
||||
mock_httpx_response=mock_httpx_response,
|
||||
)
|
||||
)
|
||||
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],
|
||||
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.
|
||||
Uses the existing VertexGeminiConfig transformation for the response.
|
||||
"""
|
||||
# Extract custom_id from request labels (prefer raw for OpenAI round-trip)
|
||||
request_data = vertex_output.get("request", {})
|
||||
labels = request_data.get("labels", {}) or {}
|
||||
custom_id = _get_litellm_batch_custom_id_from_labels(labels)
|
||||
|
||||
# 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]
|
||||
|
||||
try:
|
||||
# Use existing VertexGeminiConfig transformation
|
||||
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):
|
||||
"""
|
||||
|
|
@ -492,29 +848,11 @@ class VertexAIJsonlFilesTransformation(VertexGeminiConfig):
|
|||
|
||||
def _transform_openai_jsonl_content_to_vertex_ai_jsonl_content(
|
||||
self, openai_jsonl_content: List[Dict[str, Any]]
|
||||
):
|
||||
"""
|
||||
Transforms OpenAI JSONL content to VertexAI JSONL content
|
||||
|
||||
jsonl body for vertex is {"request": <request_body>}
|
||||
Example Vertex jsonl
|
||||
{"request":{"contents": [{"role": "user", "parts": [{"text": "What is the relation between the following video and image samples?"}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/video/animals.mp4", "mimeType": "video/mp4"}}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/image/cricket.jpeg", "mimeType": "image/jpeg"}}]}]}}
|
||||
{"request":{"contents": [{"role": "user", "parts": [{"text": "Describe what is happening in this video."}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/video/another_video.mov", "mimeType": "video/mov"}}]}]}}
|
||||
"""
|
||||
|
||||
vertex_jsonl_content = []
|
||||
for _openai_jsonl_content in openai_jsonl_content:
|
||||
openai_request_body = _openai_jsonl_content.get("body") or {}
|
||||
vertex_request_body = _transform_request_body(
|
||||
messages=openai_request_body.get("messages", []),
|
||||
model=openai_request_body.get("model", ""),
|
||||
optional_params=self._map_openai_to_vertex_params(openai_request_body),
|
||||
custom_llm_provider="vertex_ai",
|
||||
litellm_params={},
|
||||
cached_content=None,
|
||||
)
|
||||
vertex_jsonl_content.append({"request": vertex_request_body})
|
||||
return vertex_jsonl_content
|
||||
) -> List[Dict[str, Any]]:
|
||||
return _openai_batch_jsonl_entries_to_vertex_wrapped_requests(
|
||||
openai_jsonl_content=openai_jsonl_content,
|
||||
map_openai_to_vertex_params=self._map_openai_to_vertex_params,
|
||||
)
|
||||
|
||||
def _get_gcs_object_name(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -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,12 @@ 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,
|
||||
_get_litellm_batch_custom_id_from_labels,
|
||||
_sanitize_gcp_label_value,
|
||||
)
|
||||
from litellm.types.llms.openai import OpenAIFileObject, HttpxBinaryResponseContent
|
||||
from openai.types.file_deleted import FileDeleted
|
||||
|
||||
|
|
@ -143,6 +149,108 @@ class TestTransformFileContent:
|
|||
assert isinstance(result, HttpxBinaryResponseContent)
|
||||
assert result.response.content == b'{"line": 1}\n{"line": 2}\n'
|
||||
|
||||
def test_should_not_mutate_caller_logging_obj_for_batch_output_transform(
|
||||
self, config, monkeypatch
|
||||
):
|
||||
original_model = "vertex_ai/original-model"
|
||||
original_start_time = 123.456
|
||||
original_optional_params = {"temperature": 0.1}
|
||||
raw_response = httpx.Response(
|
||||
status_code=200,
|
||||
content=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"),
|
||||
headers={"content-type": "application/octet-stream"},
|
||||
request=httpx.Request("GET", "https://example.com"),
|
||||
)
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.model = original_model
|
||||
logging_obj.start_time = original_start_time
|
||||
logging_obj.optional_params = original_optional_params
|
||||
captured = {}
|
||||
|
||||
def mock_transform_single(
|
||||
vertex_output,
|
||||
vertex_gemini_config,
|
||||
logging_obj,
|
||||
mock_httpx_response,
|
||||
):
|
||||
captured["logging_obj"] = logging_obj
|
||||
logging_obj.model = "gemini-2.0-flash-001"
|
||||
logging_obj.start_time = 789.0
|
||||
return {
|
||||
"custom_id": vertex_output["request"]["labels"]["litellm_custom_id"]
|
||||
}
|
||||
|
||||
monkeypatch.setattr(
|
||||
config,
|
||||
"_transform_single_vertex_batch_output_to_openai",
|
||||
mock_transform_single,
|
||||
)
|
||||
|
||||
result = config.transform_file_content_response(
|
||||
raw_response=raw_response,
|
||||
logging_obj=logging_obj,
|
||||
litellm_params={},
|
||||
)
|
||||
|
||||
assert captured["logging_obj"] is not logging_obj
|
||||
assert logging_obj.model == original_model
|
||||
assert logging_obj.start_time == original_start_time
|
||||
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):
|
||||
|
|
@ -239,3 +347,677 @@ 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_vertex_batch_output_legacy_labels_only_sanitized(self, config):
|
||||
"""Older LiteLLM batches only stored litellm_custom_id (sanitized); read path still works."""
|
||||
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": "myrequest-1"},
|
||||
},
|
||||
"response": {
|
||||
"candidates": [
|
||||
{
|
||||
"content": {
|
||||
"parts": [{"text": "Hello!"}],
|
||||
"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"))
|
||||
|
||||
assert result["custom_id"] == "myrequest-1"
|
||||
|
||||
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_transform_vertex_batch_output_with_first_line_prompt_feedback(
|
||||
self, config, monkeypatch
|
||||
):
|
||||
"""Test that promptFeedback-only first lines are detected as Vertex batch output."""
|
||||
vertex_outputs = [
|
||||
{
|
||||
"status": "",
|
||||
"processed_time": "2024-11-01T18:13:16.826+00:00",
|
||||
"request": {"labels": {"litellm_custom_id": "blocked-request"}},
|
||||
"response": {
|
||||
"promptFeedback": {"blockReason": "SAFETY"},
|
||||
"modelVersion": "gemini-2.0-flash-001@default",
|
||||
},
|
||||
},
|
||||
{
|
||||
"status": "",
|
||||
"processed_time": "2024-11-01T18:13:17.826+00:00",
|
||||
"request": {"labels": {"litellm_custom_id": "request-2"}},
|
||||
"response": {"candidates": [{"content": {"parts": [{"text": "ok"}]}}]},
|
||||
},
|
||||
]
|
||||
|
||||
def mock_transform_single(
|
||||
vertex_output,
|
||||
vertex_gemini_config,
|
||||
logging_obj,
|
||||
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
|
||||
)
|
||||
results = [
|
||||
json.loads(line) for line in transformed_content.decode("utf-8").split("\n")
|
||||
]
|
||||
|
||||
assert [result["custom_id"] for result in results] == [
|
||||
"blocked-request",
|
||||
"request-2",
|
||||
]
|
||||
|
||||
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"
|
||||
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 TestTryTransformDoesNotMutateCallerLoggingObj:
|
||||
"""Regression tests: _try_transform_vertex_batch_output_to_openai must not mutate
|
||||
the caller's logging_obj (model, start_time, optional_params)."""
|
||||
|
||||
def _make_vertex_batch_line(self) -> bytes:
|
||||
return json.dumps(
|
||||
{
|
||||
"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": "Hi!"}],
|
||||
"role": "model",
|
||||
},
|
||||
"finishReason": "STOP",
|
||||
}
|
||||
],
|
||||
"modelVersion": "gemini-2.0-flash-001@default",
|
||||
"usageMetadata": {
|
||||
"promptTokenCount": 5,
|
||||
"candidatesTokenCount": 3,
|
||||
"totalTokenCount": 8,
|
||||
},
|
||||
},
|
||||
}
|
||||
).encode("utf-8")
|
||||
|
||||
def test_should_not_overwrite_model_on_caller_logging_obj(self, config):
|
||||
sentinel_model = "original-caller-model"
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.model = sentinel_model
|
||||
logging_obj.optional_params = {"temperature": 0.9}
|
||||
|
||||
config._try_transform_vertex_batch_output_to_openai(
|
||||
content=self._make_vertex_batch_line(),
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
|
||||
assert (
|
||||
logging_obj.model == sentinel_model
|
||||
), "logging_obj.model was mutated by _try_transform_vertex_batch_output_to_openai"
|
||||
|
||||
def test_should_not_overwrite_start_time_on_caller_logging_obj(self, config):
|
||||
sentinel_start = 1234567890.0
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.start_time = sentinel_start
|
||||
logging_obj.optional_params = {}
|
||||
|
||||
config._try_transform_vertex_batch_output_to_openai(
|
||||
content=self._make_vertex_batch_line(),
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
|
||||
assert (
|
||||
logging_obj.start_time == sentinel_start
|
||||
), "logging_obj.start_time was mutated by _try_transform_vertex_batch_output_to_openai"
|
||||
|
||||
def test_should_not_overwrite_optional_params_on_caller_logging_obj(self, config):
|
||||
sentinel_params = {"temperature": 0.5, "top_p": 0.9}
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.optional_params = sentinel_params
|
||||
|
||||
config._try_transform_vertex_batch_output_to_openai(
|
||||
content=self._make_vertex_batch_line(),
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
|
||||
assert (
|
||||
logging_obj.optional_params is sentinel_params
|
||||
), "logging_obj.optional_params was replaced by _try_transform_vertex_batch_output_to_openai"
|
||||
assert logging_obj.optional_params == {
|
||||
"temperature": 0.5,
|
||||
"top_p": 0.9,
|
||||
}, "logging_obj.optional_params contents were mutated"
|
||||
|
||||
def test_should_still_transform_content_correctly(self, config):
|
||||
logging_obj = MagicMock()
|
||||
logging_obj.model = "original-model"
|
||||
logging_obj.start_time = 9999.0
|
||||
logging_obj.optional_params = {"max_tokens": 100}
|
||||
|
||||
result = config._try_transform_vertex_batch_output_to_openai(
|
||||
content=self._make_vertex_batch_line(),
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
|
||||
# Transformation should still succeed
|
||||
transformed = json.loads(result.decode("utf-8"))
|
||||
assert transformed["custom_id"] == "request-1"
|
||||
assert transformed["response"]["status_code"] == 200
|
||||
|
||||
|
||||
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"
|
||||
raw_label = vertex_request["request"]["labels"]["litellm_custom_id_raw"]
|
||||
assert raw_label != "request-1"
|
||||
assert _sanitize_gcp_label_value(raw_label) == raw_label
|
||||
|
||||
def test_long_custom_id_round_trips_across_raw_label_chunks(self):
|
||||
"""Test that long custom_ids are not truncated in raw labels."""
|
||||
transformation = VertexAIJsonlFilesTransformation()
|
||||
custom_id_a = "shared-prefix-that-is-longer-than-thirty-six-bytes-A"
|
||||
custom_id_b = "shared-prefix-that-is-longer-than-thirty-six-bytes-B"
|
||||
|
||||
openai_jsonl_content = [
|
||||
{
|
||||
"custom_id": custom_id,
|
||||
"method": "POST",
|
||||
"url": "/v1/chat/completions",
|
||||
"body": {
|
||||
"model": "gemini-1.5-flash-001",
|
||||
"messages": [{"role": "user", "content": "Question"}],
|
||||
},
|
||||
}
|
||||
for custom_id in (custom_id_a, custom_id_b)
|
||||
]
|
||||
|
||||
vertex_jsonl_content = (
|
||||
transformation._transform_openai_jsonl_content_to_vertex_ai_jsonl_content(
|
||||
openai_jsonl_content
|
||||
)
|
||||
)
|
||||
labels_a = vertex_jsonl_content[0]["request"]["labels"]
|
||||
labels_b = vertex_jsonl_content[1]["request"]["labels"]
|
||||
|
||||
assert "litellm_custom_id_raw_1" in labels_a
|
||||
assert "litellm_custom_id_raw_1" in labels_b
|
||||
assert labels_a["litellm_custom_id_raw"] == labels_b["litellm_custom_id_raw"]
|
||||
assert (
|
||||
labels_a["litellm_custom_id_raw_1"] != labels_b["litellm_custom_id_raw_1"]
|
||||
)
|
||||
assert _get_litellm_batch_custom_id_from_labels(labels_a) == custom_id_a
|
||||
assert _get_litellm_batch_custom_id_from_labels(labels_b) == custom_id_b
|
||||
|
||||
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
|
||||
)
|
||||
raw_label = vertex_request["request"]["labels"]["litellm_custom_id_raw"]
|
||||
assert raw_label != expected_custom_id
|
||||
assert _sanitize_gcp_label_value(raw_label) == raw_label
|
||||
|
||||
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 (mixed case exercises raw label)
|
||||
openai_input = [
|
||||
{
|
||||
"custom_id": "MyRequest-1",
|
||||
"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 both labels are GCP-safe and encoded raw preserves round-trip.
|
||||
assert (
|
||||
vertex_input[0]["request"]["labels"]["litellm_custom_id"] == "myrequest-1"
|
||||
)
|
||||
raw_label = vertex_input[0]["request"]["labels"]["litellm_custom_id_raw"]
|
||||
assert raw_label != "MyRequest-1"
|
||||
assert _sanitize_gcp_label_value(raw_label) == raw_label
|
||||
|
||||
# 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 (original casing, not sanitized label)
|
||||
assert openai_output["custom_id"] == "MyRequest-1"
|
||||
assert openai_output["response"]["status_code"] == 200
|
||||
|
||||
def test_custom_id_label_sanitization(self):
|
||||
"""Test that custom_id values are sanitized to meet GCP label constraints"""
|
||||
transformation = VertexAIJsonlFilesTransformation()
|
||||
|
||||
# Test sanitization function
|
||||
assert _sanitize_gcp_label_value("MyRequest-1") == "myrequest-1"
|
||||
assert _sanitize_gcp_label_value("Request.With.Dots") == "request_with_dots"
|
||||
assert _sanitize_gcp_label_value("Request With Spaces") == "request_with_spaces"
|
||||
assert _sanitize_gcp_label_value("Request@#$%Special") == "request____special"
|
||||
|
||||
# Test max length (63 chars)
|
||||
long_id = "a" * 100
|
||||
assert len(_sanitize_gcp_label_value(long_id)) == 63
|
||||
|
||||
# Test in actual transformation
|
||||
openai_input = [
|
||||
{
|
||||
"custom_id": "MyRequest-1",
|
||||
"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 both labels are safe for GCP labels.
|
||||
assert (
|
||||
vertex_input[0]["request"]["labels"]["litellm_custom_id"] == "myrequest-1"
|
||||
)
|
||||
raw_label = vertex_input[0]["request"]["labels"]["litellm_custom_id_raw"]
|
||||
assert raw_label != "MyRequest-1"
|
||||
assert _sanitize_gcp_label_value(raw_label) == raw_label
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue