diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 7878c5467d..5eda1c836d 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -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")) ) 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 120d5b48bc..239621f415 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, + _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()