add key mapping for vertex batch calls. add tests for openai and vertex batch calls

litellm/batches/batch_utils.py

@@ -207,7 +207,10 @@ async def _get_batch_output_file_content_as_dictionary(
     )
 
     if custom_llm_provider == "vertex_ai":
-        raise ValueError("Vertex AI does not support file content retrieval")
+        return await _get_vertex_ai_batch_output_with_custom_id(
+            batch=batch,
+            litellm_params=litellm_params,
+        )
 
     if batch.output_file_id is None:
         raise ValueError("Output file id is None cannot retrieve file content")
@@ -235,6 +238,60 @@ async def _get_batch_output_file_content_as_dictionary(
     return _get_file_content_as_dictionary(_file_content.content)
 
 
+async def _get_vertex_ai_batch_output_with_custom_id(
+    batch: Batch,
+    litellm_params: Optional[dict] = None,
+) -> List[dict]:
+    """
+    Vertex AI batch outputs do not reliably include custom_id.
+    Map output lines back to custom_id using keyField (preferred).
+    """
+    from litellm.files.main import afile_content
+
+    if batch.output_file_id is None:
+        raise ValueError("Output file id is None cannot retrieve file content")
+
+    if batch.input_file_id is None:
+        raise ValueError("Input file id is None cannot map vertex batch outputs")
+
+    credentials = _extract_file_access_credentials(litellm_params)
+
+    output_content = await afile_content(
+        file_id=batch.output_file_id,
+        custom_llm_provider="vertex_ai",
+        **credentials,
+    )
+    output_lines = _get_file_content_as_dictionary(output_content.content)
+
+    input_content = await afile_content(
+        file_id=batch.input_file_id,
+        custom_llm_provider="vertex_ai",
+        **credentials,
+    )
+    input_lines = _get_file_content_as_dictionary(input_content.content)
+
+    if len(output_lines) != len(input_lines):
+        raise ValueError(
+            "Vertex AI batch output line count does not match input line count"
+        )
+
+    output_has_keys = all(
+        (line.get("custom_id") or line.get("key")) for line in output_lines
+    )
+    if not output_has_keys:
+        raise ValueError(
+            "Vertex AI batch output is missing custom_id/key; cannot safely map results"
+        )
+
+    mapped_lines = []
+    for line in output_lines:
+        key = line.get("custom_id") or line.get("key")
+        if key is not None:
+            line["custom_id"] = key
+        mapped_lines.append(line)
+    return mapped_lines
+
+
 def _extract_file_access_credentials(litellm_params: Optional[dict]) -> dict:
     """
     Extract credentials from litellm_params for file access operations.
@@ -405,4 +462,4 @@ def _batch_response_was_successful(batch_job_output_file: dict) -> bool:
     Check if the batch job response status == 200
     """
     _response: dict = batch_job_output_file.get("response", None) or {}
-    return _response.get("status_code", None) == 200
+    return _response.get("status_code", None) == 200

litellm/llms/vertex_ai/batches/transformation.py

@@ -41,6 +41,7 @@ class VertexAIBatchTransformation:
         return VertexAIBatchPredictionJob(
             inputConfig=input_config,
             outputConfig=output_config,
+            instanceConfig={"keyField": "custom_id"},
             model=model,
             displayName=request_display_name,
         )

litellm/llms/vertex_ai/files/transformation.py

@@ -243,7 +243,11 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig):
                 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 transform_create_file_request(

litellm/types/llms/vertex_ai.py

@@ -2,6 +2,7 @@ from enum import Enum
 from typing import Any, Dict, List, Literal, Optional, Union
 
 from typing_extensions import (
+    NotRequired,
     Required,
     TypedDict,
 )
@@ -568,6 +569,13 @@ class InputConfig(TypedDict):
     gcsSource: GcsSource
 
 
+class InstanceConfig(TypedDict, total=False):
+    instanceType: str
+    keyField: str
+    includedFields: List[str]
+    excludedFields: List[str]
+
+
 class GcsDestination(TypedDict):
     outputUriPrefix: str
 
@@ -630,6 +638,7 @@ class VertexAIBatchPredictionJob(TypedDict):
     displayName: str
     model: str
     inputConfig: InputConfig
+    instanceConfig: NotRequired[InstanceConfig]
     outputConfig: OutputConfig
 
 

tests/batches_tests/test_batch_custom_id_mapping_live.py

@@ -0,0 +1,235 @@
+import asyncio
+import json
+import os
+import tempfile
+from typing import List, Optional
+
+import pytest
+
+import litellm
+
+
+def _write_batch_jsonl(model: str) -> str:
+    records = [
+        {
+            "custom_id": "request-1",
+            "method": "POST",
+            "url": "/v1/chat/completions",
+            "body": {
+                "model": model,
+                "messages": [
+                    {"role": "system", "content": "You are a helpful assistant."},
+                    {"role": "user", "content": "Hello world!"},
+                ],
+                "max_tokens": 10,
+            },
+        },
+        {
+            "custom_id": "request-2",
+            "method": "POST",
+            "url": "/v1/chat/completions",
+            "body": {
+                "model": model,
+                "messages": [
+                    {"role": "system", "content": "You are an unhelpful assistant."},
+                    {"role": "user", "content": "Hello world!"},
+                ],
+                "max_tokens": 10,
+            },
+        },
+        {
+            "custom_id": "request-3",
+            "method": "POST",
+            "url": "/v1/chat/completions",
+            "body": {
+                "model": model,
+                "messages": [
+                    {"role": "system", "content": "Answer with a single word."},
+                    {"role": "user", "content": "Hi"},
+                ],
+                "max_tokens": 5,
+            },
+        },
+    ]
+
+    fd, path = tempfile.mkstemp(suffix=".jsonl")
+    os.close(fd)
+    with open(path, "w", encoding="utf-8") as f:
+        for record in records:
+            f.write(json.dumps(record) + "\n")
+    return path
+
+
+def _load_vertex_ai_credentials_from_env() -> None:
+    os.environ["GCS_FLUSH_INTERVAL"] = "1"
+
+    private_key_id = os.environ.get("GCS_PRIVATE_KEY_ID", "")
+    private_key = os.environ.get("GCS_PRIVATE_KEY", "").replace("\\n", "\n")
+
+    if not private_key_id or not private_key:
+        return
+
+    service_account_key_data = {
+        "private_key_id": private_key_id,
+        "private_key": private_key,
+    }
+
+    with tempfile.NamedTemporaryFile(mode="w+", delete=False) as temp_file:
+        json.dump(service_account_key_data, temp_file, indent=2)
+
+    abs_path = os.path.abspath(temp_file.name)
+    os.environ["GCS_PATH_SERVICE_ACCOUNT"] = abs_path
+    os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = abs_path
+
+
+def _normalize_credential_path(env_key: str) -> Optional[str]:
+    path = os.getenv(env_key)
+    if not path:
+        return None
+    abs_path = os.path.abspath(path)
+    if not os.path.exists(abs_path):
+        return None
+    os.environ[env_key] = abs_path
+    return abs_path
+
+
+def _extract_custom_ids_from_output(lines: List[dict]) -> List[str]:
+    custom_ids: List[str] = []
+    for line in lines:
+        if "custom_id" in line:
+            custom_ids.append(line["custom_id"])
+        elif "key" in line:
+            custom_ids.append(line["key"])
+    return custom_ids
+
+
+async def _wait_for_batch_completion(
+    batch_id: str,
+    provider: str,
+    timeout_seconds: Optional[int] = None,
+    poll_seconds: Optional[int] = None,
+    require_output_file_id: bool = True,
+):
+    timeout_seconds = timeout_seconds or int(
+        os.getenv("BATCH_WAIT_TIMEOUT_SECONDS", "1800")
+    )
+    poll_seconds = poll_seconds or int(os.getenv("BATCH_WAIT_POLL_SECONDS", "10"))
+
+    start = asyncio.get_event_loop().time()
+    while True:
+        batch = await litellm.aretrieve_batch(
+            batch_id=batch_id,
+            custom_llm_provider=provider,
+            litellm_params={
+                "litellm_metadata": {"batch_ignore_default_logging": True},
+            },
+        )
+        if batch.status == "completed":
+            if not require_output_file_id or getattr(batch, "output_file_id", None):
+                return batch
+        if batch.status in ["failed", "cancelled", "expired"]:
+            raise AssertionError(f"Batch ended with status={batch.status}")
+        if asyncio.get_event_loop().time() - start > timeout_seconds:
+            raise AssertionError("Timed out waiting for batch completion")
+        await asyncio.sleep(poll_seconds)
+
+
+@pytest.mark.asyncio
+async def test_openai_batch_custom_id_mapping_live():
+    api_key = os.getenv("OPENAI_API_KEY")
+    if not api_key:
+        pytest.skip("OPENAI_API_KEY not set")
+
+    file_path = _write_batch_jsonl(model="gpt-4o-mini")
+
+    with open(file_path, "rb") as file_handle:
+        file_obj = await litellm.acreate_file(
+            file=file_handle,
+            purpose="batch",
+            custom_llm_provider="openai",
+        )
+
+    create_batch_response = await litellm.acreate_batch(
+        completion_window="24h",
+        endpoint="/v1/chat/completions",
+        input_file_id=file_obj.id,
+        custom_llm_provider="openai",
+    )
+
+    completed_batch = await _wait_for_batch_completion(
+        batch_id=create_batch_response.id,
+        provider="openai",
+    )
+
+    assert completed_batch.output_file_id is not None
+
+    output_content = await litellm.afile_content(
+        file_id=completed_batch.output_file_id,
+        custom_llm_provider="openai",
+    )
+
+    output_lines = [
+        json.loads(line)
+        for line in output_content.content.decode("utf-8").strip().split("\n")
+        if line.strip()
+    ]
+
+    custom_ids = _extract_custom_ids_from_output(output_lines)
+    assert sorted(custom_ids) == ["request-1", "request-2", "request-3"]
+
+
+@pytest.mark.asyncio
+async def test_vertex_batch_custom_id_mapping_live():
+    if not os.getenv("VERTEXAI_PROJECT") or not os.getenv("VERTEXAI_LOCATION"):
+        pytest.skip("VERTEXAI_PROJECT or VERTEXAI_LOCATION not set")
+    if not os.getenv("GCS_BUCKET_NAME"):
+        pytest.skip("GCS_BUCKET_NAME not set")
+
+    _load_vertex_ai_credentials_from_env()
+
+    normalized_gcs_path = _normalize_credential_path("GCS_PATH_SERVICE_ACCOUNT")
+    normalized_google_path = _normalize_credential_path(
+        "GOOGLE_APPLICATION_CREDENTIALS"
+    )
+    if not normalized_gcs_path and not normalized_google_path:
+        pytest.skip(
+            "Vertex credentials not set or file not found "
+            "(GCS_PATH_SERVICE_ACCOUNT/GOOGLE_APPLICATION_CREDENTIALS)"
+        )
+
+    file_path = _write_batch_jsonl(model="gemini-2.5-flash-lite")
+
+    with open(file_path, "rb") as file_handle:
+        file_obj = await litellm.acreate_file(
+            file=file_handle,
+            purpose="batch",
+            custom_llm_provider="vertex_ai",
+        )
+
+    create_batch_response = await litellm.acreate_batch(
+        completion_window="24h",
+        endpoint="/v1/chat/completions",
+        input_file_id=file_obj.id,
+        custom_llm_provider="vertex_ai",
+    )
+
+    completed_batch = await _wait_for_batch_completion(
+        batch_id=create_batch_response.id,
+        provider="vertex_ai",
+    )
+
+    assert completed_batch.output_file_id is not None
+
+    output_content = await litellm.afile_content(
+        file_id=completed_batch.output_file_id,
+        custom_llm_provider="vertex_ai",
+    )
+
+    output_lines = [
+        json.loads(line)
+        for line in output_content.content.decode("utf-8").strip().split("\n")
+        if line.strip()
+    ]
+
+    custom_ids = _extract_custom_ids_from_output(output_lines)
+    assert sorted(custom_ids) == ["request-1", "request-2", "request-3"]

tests/test_litellm/batches/test_vertex_batch_mapping.py

@@ -0,0 +1,127 @@
+import json
+from types import SimpleNamespace
+
+import pytest
+
+from litellm.batches.batch_utils import _get_batch_output_file_content_as_dictionary
+
+
+class DummyBatch(SimpleNamespace):
+    pass
+
+
+@pytest.mark.asyncio
+async def test_should_map_vertex_output_by_key_field(monkeypatch):
+    """Ensure Vertex output lines with key/custom_id are mapped directly."""
+    input_lines = [
+        {"custom_id": "id-1", "body": {"model": "m", "messages": []}},
+        {"custom_id": "id-2", "body": {"model": "m", "messages": []}},
+    ]
+    output_lines = [
+        {"key": "id-2", "response": {"status_code": 200}},
+        {"key": "id-1", "response": {"status_code": 200}},
+    ]
+
+    content_map = {
+        "gs://bucket/input.jsonl": "\n".join(
+            json.dumps(line) for line in input_lines
+        ).encode("utf-8"),
+        "gs://bucket/output.jsonl": "\n".join(
+            json.dumps(line) for line in output_lines
+        ).encode("utf-8"),
+    }
+
+    async def fake_afile_content(*, file_id, **_kwargs):
+        return SimpleNamespace(content=content_map[file_id])
+
+    monkeypatch.setattr("litellm.files.main.afile_content", fake_afile_content)
+
+    batch = DummyBatch(
+        input_file_id="gs://bucket/input.jsonl",
+        output_file_id="gs://bucket/output.jsonl",
+    )
+
+    result = await _get_batch_output_file_content_as_dictionary(
+        batch=batch,
+        custom_llm_provider="vertex_ai",
+    )
+
+    assert [line.get("custom_id") for line in result] == ["id-2", "id-1"]
+
+
+@pytest.mark.asyncio
+async def test_should_map_vertex_output_by_input_order(monkeypatch):
+    """Should error when output lacks key/custom_id fields."""
+    input_lines = [
+        {"custom_id": "id-1", "body": {"model": "m", "messages": []}},
+        {"custom_id": "id-2", "body": {"model": "m", "messages": []}},
+    ]
+    output_lines = [
+        {"response": {"status_code": 200}},
+        {"response": {"status_code": 200}},
+    ]
+
+    content_map = {
+        "gs://bucket/input.jsonl": "\n".join(
+            json.dumps(line) for line in input_lines
+        ).encode("utf-8"),
+        "gs://bucket/output.jsonl": "\n".join(
+            json.dumps(line) for line in output_lines
+        ).encode("utf-8"),
+    }
+
+    async def fake_afile_content(*, file_id, **_kwargs):
+        return SimpleNamespace(content=content_map[file_id])
+
+    monkeypatch.setattr("litellm.files.main.afile_content", fake_afile_content)
+
+    batch = DummyBatch(
+        input_file_id="gs://bucket/input.jsonl",
+        output_file_id="gs://bucket/output.jsonl",
+    )
+
+    with pytest.raises(
+        ValueError,
+        match="Vertex AI batch output is missing custom_id/key",
+    ):
+        await _get_batch_output_file_content_as_dictionary(
+            batch=batch,
+            custom_llm_provider="vertex_ai",
+        )
+
+
+@pytest.mark.asyncio
+async def test_should_raise_on_vertex_output_count_mismatch(monkeypatch):
+    """Guard against mismatched input/output line counts for Vertex."""
+    input_lines = [
+        {"custom_id": "id-1", "body": {"model": "m", "messages": []}},
+        {"custom_id": "id-2", "body": {"model": "m", "messages": []}},
+    ]
+    output_lines = [
+        {"response": {"status_code": 200}},
+    ]
+
+    content_map = {
+        "gs://bucket/input.jsonl": "\n".join(
+            json.dumps(line) for line in input_lines
+        ).encode("utf-8"),
+        "gs://bucket/output.jsonl": "\n".join(
+            json.dumps(line) for line in output_lines
+        ).encode("utf-8"),
+    }
+
+    async def fake_afile_content(*, file_id, **_kwargs):
+        return SimpleNamespace(content=content_map[file_id])
+
+    monkeypatch.setattr("litellm.files.main.afile_content", fake_afile_content)
+
+    batch = DummyBatch(
+        input_file_id="gs://bucket/input.jsonl",
+        output_file_id="gs://bucket/output.jsonl",
+    )
+
+    with pytest.raises(ValueError, match="output line count does not match input"):
+        await _get_batch_output_file_content_as_dictionary(
+            batch=batch,
+            custom_llm_provider="vertex_ai",
+        )