From 5f3fcb96042f3929568f72f313b0d2d7b18ddb1c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 31 Oct 2026 01:33:04 -0700 Subject: [PATCH] fix(vertex_ai): keep the readable custom_id label's value when a row has no raw label --- litellm/llms/vertex_ai/files/transformation.py | 4 ++-- .../vertex_ai/files/test_vertex_ai_files_transformation.py | 3 +++ 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 3c4f428d65e..f0546e18569 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -279,7 +279,7 @@ def _get_litellm_batch_custom_id_from_labels(labels: Mapping[str, object] | None return decoded return raw readable: Final = labels.get("litellm_custom_id") - return readable if isinstance(readable, str) else "unknown" + return "unknown" if readable is None else str(readable) def _is_vertex_embeddings_batch_output_row(vertex_output_row: Mapping[str, object]) -> bool: @@ -1660,7 +1660,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): Vertex AI batch output format (predictions.jsonl): { - "request": {"contents": [...], "labels": {"litellm_custom_id": "request-1", "litellm_custom_id_raw": "..."}}, + "request": {"contents": [...], "labels": {"litellm_custom_id": "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" diff --git a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py index 36a97005aef..9074e3c8ded 100644 --- a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py +++ b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -1081,6 +1081,9 @@ class TestVertexBatchCustomIdLabels: assert _get_litellm_batch_custom_id_from_labels(labels) == "set-by-hand" + def test_readable_label_vertex_typed_as_a_number_still_names_a_row_written_without_raw_labels(self): + assert _get_litellm_batch_custom_id_from_labels({"litellm_custom_id": 7}) == "7" + def test_custom_ids_survive_the_nulls_vertex_adds_for_label_keys_other_rows_set(self): long_custom_id: Final = "request-with-a-long-custom-id-that-needs-a-second-label-0001" vertex_rows: Final = _wrap_entries(