diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 40e5702b061..a6c48b33d69 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -511,7 +511,11 @@ class VertexAIJsonlFilesTransformation(VertexGeminiConfig): litellm_params={}, cached_content=None, ) - vertex_jsonl_content.append({"request": vertex_request_body}) + vertex_line: Dict[str, Any] = {"request": vertex_request_body} + custom_id = _openai_jsonl_content.get("custom_id") + if custom_id is not None: + vertex_line["custom_id"] = custom_id + vertex_jsonl_content.append(vertex_line) return vertex_jsonl_content def _get_gcs_object_name( diff --git a/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_jsonl_custom_id_preserved.py b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_jsonl_custom_id_preserved.py new file mode 100644 index 00000000000..1f1540dfe72 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/files/test_vertex_ai_jsonl_custom_id_preserved.py @@ -0,0 +1,28 @@ +import json + +from litellm.llms.vertex_ai.files.transformation import VertexAIJsonlFilesTransformation + + +def test_should_preserve_custom_id_in_vertex_jsonl_transformation(): + """Ensure custom_id is preserved in Vertex JSONL batch uploads.""" + transformer = VertexAIJsonlFilesTransformation() + openai_jsonl_content = [ + { + "custom_id": "req-1", + "body": {"model": "gemini-2.5-flash-lite", "messages": []}, + }, + { + "custom_id": "req-2", + "body": {"model": "gemini-2.5-flash-lite", "messages": []}, + }, + ] + + vertex_jsonl_content = transformer._transform_openai_jsonl_content_to_vertex_ai_jsonl_content( + openai_jsonl_content + ) + + assert [item.get("custom_id") for item in vertex_jsonl_content] == [ + "req-1", + "req-2", + ] + assert all("request" in item for item in vertex_jsonl_content)