fix(vertex_ai): keep the readable custom_id label's value when a row has no raw label

This commit is contained in:
mateo-berri 2026-10-31 01:33:04 -07:00
parent e9c1018870
commit 5f3fcb9604
2 changed files with 5 additions and 2 deletions

View file

@ -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"

View file

@ -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(