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