diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 75d2c69bf1c..903d79ffa46 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -1,8 +1,9 @@ +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 @@ -39,8 +40,8 @@ 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: @@ -62,27 +63,95 @@ def _sanitize_gcp_label_value(value: str) -> str: return sanitized[:_GCP_LABEL_VALUE_MAX_LEN] +def _encode_gcp_label_value(value: str) -> str: + """Encode arbitrary text into a GCP-label-safe value.""" + max_encoded_len = _GCP_LABEL_VALUE_MAX_LEN - len(_CUSTOM_ID_RAW_LABEL_PREFIX) + max_raw_bytes = (max_encoded_len * 5) // 8 + raw_bytes = value.encode("utf-8")[:max_raw_bytes] + while raw_bytes: + try: + raw_bytes.decode("utf-8") + break + except UnicodeDecodeError: + raw_bytes = raw_bytes[:-1] + encoded = base64.b32encode(raw_bytes).decode("ascii").rstrip("=").lower() + return f"{_CUSTOM_ID_RAW_LABEL_PREFIX}{encoded}" + + +def _decode_gcp_label_value(value: str) -> Optional[str]: + """Decode values produced by _encode_gcp_label_value.""" + if not value.startswith(_CUSTOM_ID_RAW_LABEL_PREFIX): + return None + encoded = value[len(_CUSTOM_ID_RAW_LABEL_PREFIX) :].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`` preserves the original string (truncated) for + ``litellm_custom_id_raw`` encodes the original string (truncated) 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) - labels["litellm_custom_id_raw"] = custom_id_str[:_GCP_LABEL_VALUE_MAX_LEN] + labels["litellm_custom_id_raw"] = _encode_gcp_label_value(custom_id_str) def _get_litellm_batch_custom_id_from_labels(labels: Dict[str, Any]) -> str: - """Prefer unsanitized custom_id when present (see _set_litellm_batch_custom_id_labels).""" + """Prefer encoded custom_id when present (see _set_litellm_batch_custom_id_labels).""" raw = labels.get("litellm_custom_id_raw") if raw: + decoded = _decode_gcp_label_value(str(raw)) + if decoded is not None: + return decoded return str(raw) return str(labels.get("litellm_custom_id", "unknown")) +def _transform_openai_jsonl_content_to_vertex_ai_jsonl_content( + 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 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=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: + 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): """ Config for VertexAI Files @@ -273,38 +342,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, - ) - - # Add custom_id as a label for correlation in batch outputs - custom_id = _openai_jsonl_content.get("custom_id") - if custom_id: - if "labels" not in vertex_request_body: - vertex_request_body["labels"] = {} - _set_litellm_batch_custom_id_labels( - vertex_request_body["labels"], custom_id - ) - - vertex_jsonl_content.append({"request": vertex_request_body}) - return vertex_jsonl_content + return _transform_openai_jsonl_content_to_vertex_ai_jsonl_content( + openai_jsonl_content=openai_jsonl_content, + map_openai_to_vertex_params=self._map_openai_to_vertex_params, + ) def transform_create_file_request( self, @@ -779,39 +820,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, - ) - - # Add custom_id as a label for correlation in batch outputs - custom_id = _openai_jsonl_content.get("custom_id") - if custom_id: - if "labels" not in vertex_request_body: - vertex_request_body["labels"] = {} - _set_litellm_batch_custom_id_labels( - vertex_request_body["labels"], custom_id - ) - - vertex_jsonl_content.append({"request": vertex_request_body}) - return vertex_jsonl_content + ) -> List[Dict[str, Any]]: + return _transform_openai_jsonl_content_to_vertex_ai_jsonl_content( + 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 9c49cc9b06f..e957f32aca1 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 @@ -13,6 +13,7 @@ from unittest.mock import MagicMock from litellm.llms.vertex_ai.files.transformation import ( VertexAIFilesConfig, VertexAIJsonlFilesTransformation, + _sanitize_gcp_label_value, ) from litellm.types.llms.openai import OpenAIFileObject, HttpxBinaryResponseContent from openai.types.file_deleted import FileDeleted @@ -256,27 +257,31 @@ class TestVertexBatchOutputTransformation: "processed_time": "2024-11-01T18:13:16.826+00:00", "request": { "contents": [{"role": "user", "parts": [{"text": "Hello world!"}]}], - "labels": {"litellm_custom_id": "request-1"} + "labels": {"litellm_custom_id": "request-1"}, }, "response": { - "candidates": [{ - "content": { - "parts": [{"text": "Hello! How can I help you today?"}], - "role": "model" - }, - "finishReason": "STOP" - }], + "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 - } - } + "totalTokenCount": 30, + }, + }, } content = json.dumps(vertex_output).encode("utf-8") - transformed_content = config._try_transform_vertex_batch_output_to_openai(content) + transformed_content = config._try_transform_vertex_batch_output_to_openai( + content + ) result = json.loads(transformed_content.decode("utf-8")) # Verify OpenAI format @@ -284,20 +289,20 @@ class TestVertexBatchOutputTransformation: 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] @@ -312,13 +317,15 @@ class TestVertexBatchOutputTransformation: "processed_time": "2024-11-01T18:13:16.826+00:00", "request": { "contents": [{"role": "user", "parts": [{"text": "Hello world!"}]}], - "labels": {"litellm_custom_id": "request-error"} + "labels": {"litellm_custom_id": "request-error"}, }, - "response": {} + "response": {}, } content = json.dumps(vertex_output).encode("utf-8") - transformed_content = config._try_transform_vertex_batch_output_to_openai(content) + transformed_content = config._try_transform_vertex_batch_output_to_openai( + content + ) result = json.loads(transformed_content.decode("utf-8")) # Verify error format @@ -337,13 +344,15 @@ class TestVertexBatchOutputTransformation: "labels": {"litellm_custom_id": "myrequest-1"}, }, "response": { - "candidates": [{ - "content": { - "parts": [{"text": "Hello!"}], - "role": "model", - }, - "finishReason": "STOP", - }], + "candidates": [ + { + "content": { + "parts": [{"text": "Hello!"}], + "role": "model", + }, + "finishReason": "STOP", + } + ], "modelVersion": "gemini-2.0-flash-001@default", "usageMetadata": { "promptTokenCount": 10, @@ -354,7 +363,9 @@ class TestVertexBatchOutputTransformation: } content = json.dumps(vertex_output).encode("utf-8") - transformed_content = config._try_transform_vertex_batch_output_to_openai(content) + 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" @@ -366,48 +377,66 @@ class TestVertexBatchOutputTransformation: "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"} + "contents": [ + {"role": "user", "parts": [{"text": "First request"}]} + ], + "labels": {"litellm_custom_id": "request-1"}, }, "response": { - "candidates": [{ - "content": {"parts": [{"text": "First response"}], "role": "model"}, - "finishReason": "STOP" - }], + "candidates": [ + { + "content": { + "parts": [{"text": "First response"}], + "role": "model", + }, + "finishReason": "STOP", + } + ], "modelVersion": "gemini-2.0-flash-001@default", "usageMetadata": { "promptTokenCount": 5, "candidatesTokenCount": 10, - "totalTokenCount": 15 - } - } + "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"} + "contents": [ + {"role": "user", "parts": [{"text": "Second request"}]} + ], + "labels": {"litellm_custom_id": "request-2"}, }, "response": { - "candidates": [{ - "content": {"parts": [{"text": "Second response"}], "role": "model"}, - "finishReason": "STOP" - }], + "candidates": [ + { + "content": { + "parts": [{"text": "Second response"}], + "role": "model", + }, + "finishReason": "STOP", + } + ], "modelVersion": "gemini-2.0-flash-001@default", "usageMetadata": { "promptTokenCount": 6, "candidatesTokenCount": 11, - "totalTokenCount": 17 - } - } - } + "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) + 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): @@ -423,13 +452,17 @@ class TestVertexBatchOutputTransformation: 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) + 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) + transformed_content = config._try_transform_vertex_batch_output_to_openai( + invalid_content + ) assert transformed_content == invalid_content @@ -448,13 +481,15 @@ class TestVertexBatchCustomIdLabels: "body": { "model": "gemini-1.5-flash-001", "messages": [{"role": "user", "content": "What is 2+2?"}], - "max_tokens": 10 - } + "max_tokens": 10, + }, } ] - vertex_jsonl_content = transformation._transform_openai_jsonl_content_to_vertex_ai_jsonl_content( - openai_jsonl_content + vertex_jsonl_content = ( + transformation._transform_openai_jsonl_content_to_vertex_ai_jsonl_content( + openai_jsonl_content + ) ) assert len(vertex_jsonl_content) == 1 @@ -464,7 +499,9 @@ class TestVertexBatchCustomIdLabels: 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" - assert vertex_request["request"]["labels"]["litellm_custom_id_raw"] == "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_multiple_requests_each_get_their_own_label(self): """Test that multiple requests each get their own custom_id label""" @@ -478,21 +515,28 @@ class TestVertexBatchCustomIdLabels: "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 + 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 - assert vertex_request["request"]["labels"]["litellm_custom_id_raw"] == expected_custom_id + 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""" @@ -505,12 +549,14 @@ class TestVertexBatchCustomIdLabels: "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 + 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 @@ -533,17 +579,23 @@ class TestVertexBatchCustomIdLabels: "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 + vertex_input = ( + transformation._transform_openai_jsonl_content_to_vertex_ai_jsonl_content( + openai_input + ) ) - # Verify GCP-safe label and preserved raw for round-trip - assert vertex_input[0]["request"]["labels"]["litellm_custom_id"] == "myrequest-1" - assert vertex_input[0]["request"]["labels"]["litellm_custom_id_raw"] == "MyRequest-1" + # 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 = { @@ -551,22 +603,26 @@ class TestVertexBatchCustomIdLabels: "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" - }], + "candidates": [ + { + "content": {"parts": [{"text": "Hi there!"}], "role": "model"}, + "finishReason": "STOP", + } + ], "modelVersion": "gemini-2.0-flash-001@default", "usageMetadata": { "promptTokenCount": 5, "candidatesTokenCount": 10, - "totalTokenCount": 15 - } - } + "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) + 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) @@ -575,11 +631,6 @@ class TestVertexBatchCustomIdLabels: def test_custom_id_label_sanitization(self): """Test that custom_id values are sanitized to meet GCP label constraints""" - from litellm.llms.vertex_ai.files.transformation import ( - VertexAIJsonlFilesTransformation, - _sanitize_gcp_label_value, - ) - transformation = VertexAIJsonlFilesTransformation() # Test sanitization function @@ -587,7 +638,7 @@ class TestVertexBatchCustomIdLabels: 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 @@ -601,14 +652,20 @@ class TestVertexBatchCustomIdLabels: "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 + vertex_input = ( + transformation._transform_openai_jsonl_content_to_vertex_ai_jsonl_content( + openai_input + ) ) - # Verify label was sanitized and original retained for read-back - assert vertex_input[0]["request"]["labels"]["litellm_custom_id"] == "myrequest-1" - assert vertex_input[0]["request"]["labels"]["litellm_custom_id_raw"] == "MyRequest-1" + # 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