Fix Vertex batch output edge cases
This commit is contained in:
Cursor Agent
parent
0cbb35cdcc
commit
60ca4424a6
@ -66,23 +66,39 @@ def _sanitize_gcp_label_value(value: str) -> str:
|
||||
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()
|
||||
encoded = (
|
||||
base64.b32encode(value.encode("utf-8")).decode("ascii").rstrip("=").lower()
|
||||
)
|
||||
if len(encoded) > max_encoded_len:
|
||||
raise ValueError("Encoded label value exceeds GCP label length")
|
||||
return f"{_CUSTOM_ID_RAW_LABEL_PREFIX}{encoded}"
|
||||
|
||||
|
||||
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(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()
|
||||
return _decode_gcp_label_value_chunks([value])
|
||||
|
||||
|
||||
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")
|
||||
@ -95,19 +111,32 @@ def _set_litellm_batch_custom_id_labels(labels: Dict[str, str], custom_id: Any)
|
||||
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 (truncated) for
|
||||
``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)
|
||||
labels["litellm_custom_id_raw"] = _encode_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:
|
||||
decoded = _decode_gcp_label_value(str(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)
|
||||
@ -634,6 +663,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
|
||||
)
|
||||
has_success_or_error = (
|
||||
"candidates" in first_line.get("response", {})
|
||||
or "promptFeedback" in first_line.get("response", {})
|
||||
or bool(first_line.get("status"))
|
||||
)
|
||||
|
||||
|
||||
@ -13,6 +13,7 @@ from unittest.mock import MagicMock
|
||||
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
|
||||
@ -449,6 +450,59 @@ class TestVertexBatchOutputTransformation:
|
||||
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 = {
|
||||
@ -491,7 +545,9 @@ class TestVertexBatchOutputTransformation:
|
||||
id(mock_httpx_response),
|
||||
)
|
||||
)
|
||||
return {"custom_id": vertex_output["request"]["labels"]["litellm_custom_id"]}
|
||||
return {
|
||||
"custom_id": vertex_output["request"]["labels"]["litellm_custom_id"]
|
||||
}
|
||||
|
||||
monkeypatch.setattr(
|
||||
config,
|
||||
@ -563,6 +619,42 @@ class TestVertexBatchCustomIdLabels:
|
||||
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()
|
||||
|
||||
Loading…
Reference in New Issue
Block a user