diff --git a/litellm/__init__.py b/litellm/__init__.py index d02d39cae15..28112f5c12a 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -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 diff --git a/litellm/llms/vertex_ai/files/handler.py b/litellm/llms/vertex_ai/files/handler.py index 6636bccd6a3..bd4b2ac8bbc 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 @@ -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, diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 070ec508283..7df16723870 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -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": } + 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": } - 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": {} + }, + "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": } - 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, 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 596726cdb4b..4adad53a98f 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,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