mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-11 22:51:28 +00:00
Fix Vertex batch custom id labels
This commit is contained in:
Cursor Agent
parent
8b85deb1df
commit
7deac91a92
2 changed files with 230 additions and 160 deletions
|
|
@ -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": <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:
|
||||
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": <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,
|
||||
)
|
||||
|
||||
# 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": <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,
|
||||
)
|
||||
|
||||
# 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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue