From 692a812f13183b92f9d175abe8dcd5fad33f33ab Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 29 Jul 2026 20:17:39 -0700 Subject: [PATCH 1/3] fix(batches): calculate cost and usage for completed Vertex AI batches --- litellm/batches/batch_utils.py | 3 - .../test_litellm/batches/test_batch_utils.py | 114 +++++++++++++++++- 2 files changed, 111 insertions(+), 6 deletions(-) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 2fcb8455e90..36076d59a35 100644 --- a/litellm/batches/batch_utils.py +++ b/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") diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index c7aecac477e..fea897e9372 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/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 From 27ccc444715ff9c8e98ddc2d102ff92c4e703962 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 29 Jul 2026 20:44:31 -0700 Subject: [PATCH 2/3] fix(batches): forward gcs_bucket_name so vertex batch cost logging can read the output file --- litellm/batches/batch_utils.py | 2 ++ tests/test_litellm/batches/test_batch_utils.py | 8 ++++++++ 2 files changed, 10 insertions(+) diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 36076d59a35..b62331b904b 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -267,6 +267,8 @@ def _extract_file_access_credentials(litellm_params: Optional[dict]) -> dict: "vertex_project", "vertex_location", "vertex_credentials", + "gcs_bucket_name", + "bucket_name", "timeout", "max_retries", ] diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index fea897e9372..70c89f70ceb 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -236,6 +236,8 @@ def test_extract_credentials_only_known_keys(): "api_key": "sk-1", "api_base": "https://b", "vertex_project": "proj", + "gcs_bucket_name": "my-bucket", + "bucket_name": "my-alias-bucket", "model": "gpt-4o", # not a credential key "unrelated": "x", } @@ -243,6 +245,8 @@ def test_extract_credentials_only_known_keys(): "api_key": "sk-1", "api_base": "https://b", "vertex_project": "proj", + "gcs_bucket_name": "my-bucket", + "bucket_name": "my-alias-bucket", } @@ -262,6 +266,8 @@ def test_extract_credentials_all_supported_keys(): "vertex_project", "vertex_location", "vertex_credentials", + "gcs_bucket_name", + "bucket_name", "timeout", "max_retries", } @@ -665,6 +671,7 @@ async def test_output_file_content_vertex_fetches_via_afile_content(monkeypatch) "vertex_project": "proj-1", "vertex_location": "us-central1", "vertex_credentials": "/path/to/creds.json", + "gcs_bucket_name": "litellm-bucket", "model": "vertex_ai/gemini-3.6-flash", }, ) @@ -675,6 +682,7 @@ async def test_output_file_content_vertex_fetches_via_afile_content(monkeypatch) assert captured["vertex_project"] == "proj-1" assert captured["vertex_location"] == "us-central1" assert captured["vertex_credentials"] == "/path/to/creds.json" + assert captured["gcs_bucket_name"] == "litellm-bucket" assert "model" not in captured From 8a300929614cce9b31f6caaa748049590b63ffff Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 29 Jul 2026 23:53:15 -0700 Subject: [PATCH 3/3] test(batches): exercise real GCS validation for vertex batch output reads --- .../test_litellm/batches/test_batch_utils.py | 93 +++++++++++++++++++ 1 file changed, 93 insertions(+) diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index 70c89f70ceb..bf871623a94 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -18,7 +18,9 @@ import json import os import sys +import httpx import pytest +import respx sys.path.insert(0, os.path.abspath("../../../..")) @@ -711,6 +713,97 @@ async def test_output_file_content_vertex_unified_file_id_extracts_gcs_uri(monke assert captured["custom_llm_provider"] == "vertex_ai" +def _vertex_predictions_row(custom_id, prompt_tokens, completion_tokens): + return { + "request": { + "contents": [{"role": "user", "parts": [{"text": "hi"}]}], + "labels": {"litellm_custom_id": custom_id}, + }, + "status": "", + "response": { + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": "ok"}]}, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": prompt_tokens, + "candidatesTokenCount": completion_tokens, + "totalTokenCount": prompt_tokens + completion_tokens, + }, + "modelVersion": "gemini-3.6-flash", + }, + "processed_time": "2026-07-30T00:00:00.000000+00:00", + } + + +@pytest.fixture +def respx_interceptable_httpx_client(monkeypatch): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + yield + litellm.in_memory_llm_clients_cache.flush_cache() + + +@pytest.mark.asyncio +@respx.mock +async def test_output_file_content_vertex_managed_uri_accepted_by_real_validation(respx_interceptable_httpx_client): + managed_output_uri = ( + "gs://litellm-bucket/litellm-vertex-files/publishers/google/models/" + "gemini-3.6-flash/abc-123/prediction-model/predictions.jsonl" + ) + rows = [ + _vertex_predictions_row("request-1", 10, 5), + _vertex_predictions_row("request-2", 20, 10), + ] + route = respx.get(url__regex=r"https://storage\.googleapis\.com/storage/v1/b/litellm-bucket/o/.*").mock( + return_value=httpx.Response(200, content=_vertex_jsonl(rows)) + ) + + result = await bu._get_batch_output_file_content_as_dictionary( + _batch(managed_output_uri), + custom_llm_provider="vertex_ai", + litellm_params={ + "api_key": "test-token", + "vertex_project": "proj-1", + "vertex_location": "us-central1", + "gcs_bucket_name": "litellm-bucket", + }, + ) + + assert route.call_count == 1 + request = route.calls.last.request + assert request.url.raw_path == ( + b"/storage/v1/b/litellm-bucket/o/" + b"litellm-vertex-files%2Fpublishers%2Fgoogle%2Fmodels%2Fgemini-3.6-flash" + b"%2Fabc-123%2Fprediction-model%2Fpredictions.jsonl?alt=media" + ) + assert [row["custom_id"] for row in result] == ["request-1", "request-2"] + assert all(row["response"]["status_code"] == 200 for row in result) + assert all(row["response"]["body"]["model"] == "gemini-3.6-flash" for row in result) + assert [row["response"]["body"]["usage"]["prompt_tokens"] for row in result] == [10, 20] + assert [row["response"]["body"]["usage"]["completion_tokens"] for row in result] == [5, 10] + + +@pytest.mark.asyncio +@respx.mock +async def test_output_file_content_vertex_foreign_bucket_rejected_by_real_validation(): + with pytest.raises(Exception, match="does not match the configured storage bucket"): + await bu._get_batch_output_file_content_as_dictionary( + _batch("gs://attacker-bucket/litellm-vertex-files/x/predictions.jsonl"), + custom_llm_provider="vertex_ai", + litellm_params={ + "api_key": "test-token", + "vertex_project": "proj-1", + "vertex_location": "us-central1", + "gcs_bucket_name": "litellm-bucket", + }, + ) + + assert respx.mock.calls.call_count == 0 + + @pytest.mark.asyncio async def test_handle_completed_vertex_batch_computes_cost_usage_and_models(monkeypatch): import litellm.files.main as files_main