fix(batches): calculate cost and usage for completed Vertex AI batches

litellm/batches/batch_utils.py

@@ -212,9 +212,6 @@ async def _get_batch_output_file_content_as_dictionary(
         _is_base64_encoded_unified_file_id,
     )
 
-    if custom_llm_provider == "vertex_ai":
-        raise ValueError("Vertex AI does not support file content retrieval")
-
     if batch.output_file_id is None:
         raise ValueError("Output file id is None cannot retrieve file content")
 

tests/test_litellm/batches/test_batch_utils.py

@@ -14,6 +14,7 @@ maps (litellm.completion_cost, batch_cost_calculator), the tokenizer
 deterministic stand-ins so the arithmetic under test is the only variable.
 """
 
+import json
 import os
 import sys
 
@@ -615,10 +616,117 @@ def _batch(output_file_id):
     )
 
 
+def _vertex_openai_row(custom_id, model, prompt_tokens, completion_tokens):
+    return {
+        "id": f"batch_req_{custom_id}",
+        "custom_id": custom_id,
+        "response": {
+            "status_code": 200,
+            "request_id": custom_id,
+            "body": {
+                "id": f"chatcmpl-{custom_id}",
+                "object": "chat.completion",
+                "model": model,
+                "choices": [
+                    {
+                        "index": 0,
+                        "message": {"role": "assistant", "content": "ok"},
+                        "finish_reason": "stop",
+                    }
+                ],
+                "usage": _usage(prompt_tokens, completion_tokens),
+            },
+        },
+        "error": None,
+    }
+
+
+def _vertex_jsonl(rows):
+    return "\n".join(json.dumps(row) for row in rows).encode()
+
+
 @pytest.mark.asyncio
-async def test_output_file_content_vertex_raises():
-    with pytest.raises(ValueError, match="Vertex AI does not support"):
-        await bu._get_batch_output_file_content_as_dictionary(_batch("of"), custom_llm_provider="vertex_ai")
+async def test_output_file_content_vertex_fetches_via_afile_content(monkeypatch):
+    import litellm.files.main as files_main
+
+    rows = [_vertex_openai_row("request-1", "gemini-3.6-flash", 10, 5)]
+    captured: dict = {}
+
+    async def fake_afile_content(**kw):
+        captured.update(kw)
+        return type("R", (), {"content": _vertex_jsonl(rows)})()
+
+    monkeypatch.setattr(files_main, "afile_content", fake_afile_content)
+
+    result = await bu._get_batch_output_file_content_as_dictionary(
+        _batch("gs://litellm-bucket/output/predictions.jsonl"),
+        custom_llm_provider="vertex_ai",
+        litellm_params={
+            "vertex_project": "proj-1",
+            "vertex_location": "us-central1",
+            "vertex_credentials": "/path/to/creds.json",
+            "model": "vertex_ai/gemini-3.6-flash",
+        },
+    )
+
+    assert result == rows
+    assert captured["file_id"] == "gs://litellm-bucket/output/predictions.jsonl"
+    assert captured["custom_llm_provider"] == "vertex_ai"
+    assert captured["vertex_project"] == "proj-1"
+    assert captured["vertex_location"] == "us-central1"
+    assert captured["vertex_credentials"] == "/path/to/creds.json"
+    assert "model" not in captured
+
+
+@pytest.mark.asyncio
+async def test_output_file_content_vertex_unified_file_id_extracts_gcs_uri(monkeypatch):
+    import base64
+
+    import litellm.files.main as files_main
+
+    captured: dict = {}
+
+    async def fake_afile_content(**kw):
+        captured.update(kw)
+        return type("R", (), {"content": b'{"a": 1}'})()
+
+    monkeypatch.setattr(files_main, "afile_content", fake_afile_content)
+    unified_id = (
+        "litellm_proxy:application/jsonl;unified_id,uuid-1;target_model_names,vertex-model;"
+        "llm_output_file_id,gs://litellm-bucket/output/predictions.jsonl;llm_output_file_model_id,model-1"
+    )
+    encoded_id = base64.urlsafe_b64encode(unified_id.encode()).decode().rstrip("=")
+
+    await bu._get_batch_output_file_content_as_dictionary(_batch(encoded_id), custom_llm_provider="vertex_ai")
+
+    assert captured["file_id"] == "gs://litellm-bucket/output/predictions.jsonl"
+    assert captured["custom_llm_provider"] == "vertex_ai"
+
+
+@pytest.mark.asyncio
+async def test_handle_completed_vertex_batch_computes_cost_usage_and_models(monkeypatch):
+    import litellm.files.main as files_main
+
+    rows = [
+        _vertex_openai_row("request-1", "gemini-3.6-flash", 10, 5),
+        _vertex_openai_row("request-2", "gemini-3.6-flash", 20, 10),
+    ]
+
+    async def fake_afile_content(**kw):
+        return type("R", (), {"content": _vertex_jsonl(rows)})()
+
+    monkeypatch.setattr(files_main, "afile_content", fake_afile_content)
+
+    cost, usage, models = await bu._handle_completed_batch(
+        _batch("gs://litellm-bucket/output/predictions.jsonl"),
+        custom_llm_provider="vertex_ai",
+        litellm_params={"vertex_project": "proj-1", "vertex_location": "us-central1"},
+    )
+
+    assert cost > 0
+    assert cost == pytest.approx(30 * 7.5e-07 + 15 * 3.75e-06)
+    assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (30, 15, 45)
+    assert models == ["gemini-3.6-flash", "gemini-3.6-flash"]
 
 
 @pytest.mark.asyncio