From b0e3eb818f082e17e992b38987f3c01079fedc62 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 17 Jun 2026 11:14:43 +0530 Subject: [PATCH] Add test for litellm/tests/test_litellm/batches/test_batch_utils.py --- .../test_litellm/batches/test_batch_utils.py | 743 ++++++++++++++++++ 1 file changed, 743 insertions(+) diff --git a/tests/test_litellm/batches/test_batch_utils.py b/tests/test_litellm/batches/test_batch_utils.py index e69de29bb2d..16013e4e7a3 100644 --- a/tests/test_litellm/batches/test_batch_utils.py +++ b/tests/test_litellm/batches/test_batch_utils.py @@ -0,0 +1,743 @@ +""" +Unit tests for litellm/batches/batch_utils.py + +batch_utils.py is the batch cost/usage/parsing layer: it turns a batch output +JSONL into spend (cost), token usage, and the list of models seen, and counts +tokens in batch *input* files for rate limiting. A silent bug here mis-bills +real money or lets callers slip past TPM limits, so these tests assert exact +numeric results rather than "ran without error". + +Pure functions (parsing, token math, credential extraction, success checks) run +for real with exact-value assertions. The few true external seams - the cost +maps (litellm.completion_cost, batch_cost_calculator), the tokenizer +(token_counter), and remote file fetch (afile_content) - are mocked with +deterministic stand-ins so the arithmetic under test is the only variable. +""" + +import os +import sys + +import pytest + +sys.path.insert(0, os.path.abspath("../../../..")) + +import litellm +import litellm.batches.batch_utils as bu +from litellm.types.utils import Usage + + +# --------------------------------------------------------------------------- # +# Builders for batch OUTPUT file rows. +# Shape: {"response": {"status_code": 200, "body": {... "usage": {...}}}} +# --------------------------------------------------------------------------- # + + +def _usage(p, c, t=None): + return { + "prompt_tokens": p, + "completion_tokens": c, + "total_tokens": t if t is not None else p + c, + } + + +def _success_row(model="gpt-4o", usage=None, **body_extra): + body = {"model": model, **body_extra} + if usage is not None: + body["usage"] = usage + return {"response": {"status_code": 200, "body": body}} + + +def _failed_row(status_code=500, model="gpt-4o"): + return {"response": {"status_code": status_code, "body": {"model": model}}} + + +# =========================================================================== # +# _batch_response_was_successful +# =========================================================================== # + + +@pytest.mark.parametrize( + "row,expected", + [ + ({"response": {"status_code": 200}}, True), + ({"response": {"status_code": 500}}, False), + ({"response": {"status_code": 429}}, False), + ({"response": {}}, False), # no status_code + ({}, False), # no response + ({"response": None}, False), # null response + ], +) +def test_batch_response_was_successful(row, expected): + assert bu._batch_response_was_successful(row) is expected + + +# =========================================================================== # +# _get_response_from_batch_job_output_file +# =========================================================================== # + + +def test_get_response_body_present(): + row = {"response": {"body": {"model": "gpt-4o", "usage": {"x": 1}}}} + assert bu._get_response_from_batch_job_output_file(row) == { + "model": "gpt-4o", + "usage": {"x": 1}, + } + + +@pytest.mark.parametrize( + "row", + [ + {}, # no response + {"response": {}}, # no body + {"response": None}, # null response + {"response": {"body": None}}, # null body + ], +) +def test_get_response_body_missing_returns_empty(row): + assert bu._get_response_from_batch_job_output_file(row) == {} + + +# =========================================================================== # +# _get_batch_job_usage_from_response_body +# =========================================================================== # + + +def test_get_usage_from_response_body(): + usage = bu._get_batch_job_usage_from_response_body( + {"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}} + ) + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == ( + 10, + 5, + 15, + ) + + +def test_get_usage_from_response_body_missing_is_zero(): + usage = bu._get_batch_job_usage_from_response_body({}) + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == ( + 0, + 0, + 0, + ) + + +# =========================================================================== # +# _get_file_content_as_dictionary (JSONL parsing) +# =========================================================================== # + + +def test_parse_jsonl_multiple_lines(): + content = b'{"a": 1}\n{"b": 2}\n{"c": 3}' + assert bu._get_file_content_as_dictionary(content) == [ + {"a": 1}, + {"b": 2}, + {"c": 3}, + ] + + +def test_parse_jsonl_trailing_newline_skipped(): + # outer content is stripped; the trailing-newline empty line is dropped. + content = b'{"a": 1}\n{"b": 2}\n' + assert bu._get_file_content_as_dictionary(content) == [{"a": 1}, {"b": 2}] + + +def test_parse_jsonl_empty_content_is_empty_list(): + assert bu._get_file_content_as_dictionary(b"") == [] + + +def test_parse_jsonl_malformed_raises(): + with pytest.raises(Exception): + bu._get_file_content_as_dictionary(b"not valid json") + + +# =========================================================================== # +# _get_models_from_batch_input_file_content (distinct body.model, ordered) +# =========================================================================== # + + +def test_input_models_distinct_ordered(): + rows = [ + {"body": {"model": "gpt-4o"}}, + {"body": {"model": "claude-3"}}, + {"body": {"model": "gpt-4o"}}, # dup + {"body": {"model": "o1"}}, + ] + assert bu._get_models_from_batch_input_file_content(rows) == [ + "gpt-4o", + "claude-3", + "o1", + ] + + +def test_input_models_skips_missing_body_and_model(): + rows = [ + {"body": {"model": "gpt-4o"}}, + {}, # no body + {"body": {}}, # no model + {"body": {"model": None}}, # null model + ] + assert bu._get_models_from_batch_input_file_content(rows) == ["gpt-4o"] + + +def test_input_models_empty(): + assert bu._get_models_from_batch_input_file_content([]) == [] + + +# =========================================================================== # +# _get_batch_models_from_file_content (output file) +# =========================================================================== # + + +def test_output_models_uses_model_name_override(): + # model_name short-circuits: content is ignored entirely. + assert bu._get_batch_models_from_file_content( + [_success_row(model="ignored")], model_name="forced-model" + ) == ["forced-model"] + + +def test_output_models_collects_from_successful_only(): + rows = [ + _success_row(model="gpt-4o"), + _failed_row(model="should-be-skipped"), + _success_row(model="claude-3"), + ] + assert bu._get_batch_models_from_file_content(rows) == ["gpt-4o", "claude-3"] + + +def test_output_models_skips_successful_without_model(): + rows = [{"response": {"status_code": 200, "body": {}}}] + assert bu._get_batch_models_from_file_content(rows) == [] + + +# =========================================================================== # +# _extract_file_access_credentials +# =========================================================================== # + + +def test_extract_credentials_only_known_keys(): + params = { + "api_key": "sk-1", + "api_base": "https://b", + "vertex_project": "proj", + "model": "gpt-4o", # not a credential key + "unrelated": "x", + } + assert bu._extract_file_access_credentials(params) == { + "api_key": "sk-1", + "api_base": "https://b", + "vertex_project": "proj", + } + + +@pytest.mark.parametrize("params", [None, {}]) +def test_extract_credentials_empty(params): + assert bu._extract_file_access_credentials(params) == {} + + +def test_extract_credentials_all_supported_keys(): + keys = { + "api_key", + "api_base", + "api_version", + "organization", + "azure_ad_token", + "azure_ad_token_provider", + "vertex_project", + "vertex_location", + "vertex_credentials", + "timeout", + "max_retries", + } + params = {k: f"val-{k}" for k in keys} + assert bu._extract_file_access_credentials(params) == params + + +# =========================================================================== # +# _count_prompt_or_input_tokens (regression-critical: list[list[int]] used to +# count as zero and let callers slip past TPM limits). token_counter stubbed to +# len(text) so every shape has an exact expected value. +# =========================================================================== # + + +@pytest.fixture +def fake_token_counter(monkeypatch): + def _tc(model=None, text=None, messages=None, **kw): + if messages is not None: + return len(messages) + if text is not None: + return len(text) + return 0 + + monkeypatch.setattr(bu, "token_counter", _tc) + return _tc + + +def test_count_tokens_str(fake_token_counter): + assert bu._count_prompt_or_input_tokens("m", "hello") == 5 # len("hello") + + +def test_count_tokens_list_of_str(fake_token_counter): + assert bu._count_prompt_or_input_tokens("m", ["ab", "cde"]) == 5 # 2 + 3 + + +def test_count_tokens_list_of_int(fake_token_counter): + # pre-tokenized prompt: each int counts as one token. + assert bu._count_prompt_or_input_tokens("m", [1, 2, 3, 4]) == 4 + + +def test_count_tokens_list_of_list_of_int(fake_token_counter): + # the bug-fix shape: nested pre-tokenized prompts, each int = 1 token. + assert bu._count_prompt_or_input_tokens("m", [[1, 2, 3], [4, 5]]) == 5 + + +def test_count_tokens_mixed_nested(fake_token_counter): + # nested list with ints + a string: 2 ints (=2) + len("xyz")=3 -> 5 + assert bu._count_prompt_or_input_tokens("m", [[1, 2, "xyz"]]) == 5 + + +def test_count_tokens_unsupported_shape_is_zero(fake_token_counter): + assert bu._count_prompt_or_input_tokens("m", 12345) == 0 + assert bu._count_prompt_or_input_tokens("m", {"a": 1}) == 0 + + +# =========================================================================== # +# _get_batch_job_input_file_usage (rate-limit token counting) +# =========================================================================== # + + +def test_input_usage_messages_path(fake_token_counter): + rows = [ + {"body": {"model": "gpt-4o", "messages": [{"role": "user"}, {"role": "x"}]}} + ] + usage = bu._get_batch_job_input_file_usage(rows) + assert usage.prompt_tokens == 2 # len(messages) + assert usage.completion_tokens == 0 + assert usage.total_tokens == 2 + + +def test_input_usage_prompt_path(fake_token_counter): + rows = [{"body": {"model": "gpt-4o", "prompt": "abcd"}}] + assert bu._get_batch_job_input_file_usage(rows).prompt_tokens == 4 + + +def test_input_usage_input_path(fake_token_counter): + rows = [{"body": {"model": "gpt-4o", "input": "ab"}}] + assert bu._get_batch_job_input_file_usage(rows).prompt_tokens == 2 + + +def test_input_usage_messages_beats_prompt(fake_token_counter): + # messages present -> the prompt field is ignored (continue after messages). + rows = [ + { + "body": { + "model": "gpt-4o", + "messages": [{"role": "user"}], + "prompt": "this-should-be-ignored", + } + } + ] + assert bu._get_batch_job_input_file_usage(rows).prompt_tokens == 1 + + +def test_input_usage_sums_across_rows(fake_token_counter): + rows = [ + {"body": {"model": "m", "prompt": "abc"}}, # 3 + {"body": {"model": "m", "input": "de"}}, # 2 + ] + assert bu._get_batch_job_input_file_usage(rows).prompt_tokens == 5 + + +# =========================================================================== # +# _get_batch_job_total_usage_from_file_content (output usage aggregation) +# =========================================================================== # + + +def test_total_usage_sums_successful_only(): + rows = [ + _success_row(usage=_usage(10, 5)), # 15 + _failed_row(), # excluded + _success_row(usage=_usage(20, 10)), # 30 + ] + usage = bu._get_batch_job_total_usage_from_file_content(rows) + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == ( + 30, + 15, + 45, + ) + + +def test_total_usage_empty_is_zero(): + usage = bu._get_batch_job_total_usage_from_file_content([]) + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == ( + 0, + 0, + 0, + ) + + +# =========================================================================== # +# _get_batch_job_cost_from_file_content (cost maps mocked) +# =========================================================================== # + + +def test_cost_from_content_completion_cost_path(monkeypatch): + # model_info is None -> litellm.completion_cost per successful row. + calls = [] + + def _completion_cost(**kw): + calls.append(kw) + return 0.5 + + monkeypatch.setattr(litellm, "completion_cost", _completion_cost) + rows = [ + _success_row(usage=_usage(10, 5)), + _failed_row(), # excluded -> not costed + _success_row(usage=_usage(20, 10)), + ] + + total = bu._get_batch_job_cost_from_file_content( + rows, custom_llm_provider="openai" + ) + + assert total == 1.0 # 2 successful * 0.5 + assert len(calls) == 2 # failed row not costed + + +def test_cost_from_content_model_info_path(monkeypatch): + # model_info set -> batch_cost_calculator(prompt_cost, completion_cost). + import litellm.cost_calculator as cc + + monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (0.1, 0.2)) + rows = [ + _success_row(usage=_usage(10, 5)), + _success_row(usage=_usage(20, 10)), + ] + + total = bu._get_batch_job_cost_from_file_content( + rows, + custom_llm_provider="openai", + model_info={"input_cost_per_token": 0.0}, # type: ignore[arg-type] # truthy -> model_info path + ) + + assert total == pytest.approx(0.6) # 2 * (0.1 + 0.2) + + +# =========================================================================== # +# _batch_cost_calculator (dispatch: vertex-disable-transform vs generic) +# =========================================================================== # + + +def test_batch_cost_calculator_generic_path(monkeypatch): + monkeypatch.setattr(bu, "_get_batch_job_cost_from_file_content", lambda **kw: 4.2) + assert ( + bu._batch_cost_calculator( + [], custom_llm_provider="openai", model_name="gpt-4o" + ) + == 4.2 + ) + + +def test_batch_cost_calculator_vertex_disable_transform_path(monkeypatch): + monkeypatch.setattr( + litellm, "disable_vertex_batch_output_transformation", True, raising=False + ) + monkeypatch.setattr( + bu, + "calculate_vertex_ai_batch_cost_and_usage", + lambda content, model: (9.9, Usage()), + ) + # generic path must NOT be taken + monkeypatch.setattr( + bu, + "_get_batch_job_cost_from_file_content", + lambda **kw: pytest.fail("generic path should not run"), + ) + + cost = bu._batch_cost_calculator( + [], custom_llm_provider="vertex_ai", model_name="gemini-2.0-flash-001" + ) + assert cost == 9.9 + + +# =========================================================================== # +# calculate_vertex_ai_batch_cost_and_usage (usageMetadata aggregation) +# =========================================================================== # + + +def test_vertex_cost_and_usage_aggregation(monkeypatch): + import litellm.cost_calculator as cc + + monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (0.1, 0.2)) + responses = [ + { + "response": { + "usageMetadata": { + "promptTokenCount": 10, + "candidatesTokenCount": 5, + "totalTokenCount": 15, + } + } + }, + { + "response": { + "usageMetadata": { + "promptTokenCount": 20, + "candidatesTokenCount": 10, + "totalTokenCount": 30, + } + } + }, + ] + + cost, usage = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x") + + assert cost == pytest.approx(0.6) # 2 * (0.1 + 0.2) + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == ( + 30, + 15, + 45, + ) + + +def test_vertex_cost_skips_none_response_body(monkeypatch): + import litellm.cost_calculator as cc + + monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (1.0, 0.0)) + responses = [ + {"response": None}, # skipped + { + "response": { + "usageMetadata": { + "promptTokenCount": 7, + "candidatesTokenCount": 3, + "totalTokenCount": 10, + } + } + }, + ] + + cost, usage = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x") + + assert cost == pytest.approx(1.0) # only one line costed + assert usage.total_tokens == 10 + + +def test_vertex_usage_total_token_fallback(monkeypatch): + # no totalTokenCount -> falls back to prompt + completion. + import litellm.cost_calculator as cc + + monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (0.0, 0.0)) + responses = [ + { + "response": { + "usageMetadata": {"promptTokenCount": 8, "candidatesTokenCount": 4} + } + } + ] + + _, usage = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x") + assert usage.total_tokens == 12 + + +def test_vertex_cost_error_in_line_is_swallowed(monkeypatch): + # a cost error on one line must not abort aggregation; usage still tallies. + import litellm.cost_calculator as cc + + def _boom(**kw): + raise RuntimeError("price map miss") + + monkeypatch.setattr(cc, "batch_cost_calculator", _boom) + responses = [ + { + "response": { + "usageMetadata": { + "promptTokenCount": 5, + "candidatesTokenCount": 5, + "totalTokenCount": 10, + } + } + } + ] + + cost, usage = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x") + assert cost == 0.0 + assert usage.total_tokens == 10 + + +# =========================================================================== # +# calculate_batch_cost_and_usage (async orchestrator) +# =========================================================================== # + + +@pytest.mark.asyncio +async def test_calculate_batch_cost_and_usage_orchestration(monkeypatch): + rows = [_success_row(model="gpt-4o", usage=_usage(10, 5))] + monkeypatch.setattr(bu, "_batch_cost_calculator", lambda **kw: 2.5) + monkeypatch.setattr( + bu, + "_get_batch_job_total_usage_from_file_content", + lambda **kw: Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + ) + + cost, usage, models = await bu.calculate_batch_cost_and_usage( + file_content_dictionary=rows, custom_llm_provider="openai" + ) + + assert cost == 2.5 + assert usage.total_tokens == 15 + assert models == ["gpt-4o"] # real _get_batch_models_from_file_content + + +# =========================================================================== # +# _get_batch_output_file_content_as_dictionary (file fetch + credential merge) +# =========================================================================== # + + +def _batch(output_file_id): + from litellm.types.llms.openai import Batch + + return Batch( + id="b", + completion_window="24h", + created_at=1, + endpoint="/v1/chat/completions", + input_file_id="f", + object="batch", + status="completed", + output_file_id=output_file_id, + ) + + +@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" + ) + + +@pytest.mark.asyncio +async def test_output_file_content_no_output_file_id_raises(): + with pytest.raises(ValueError, match="Output file id is None"): + await bu._get_batch_output_file_content_as_dictionary( + _batch(None), custom_llm_provider="openai" + ) + + +@pytest.mark.asyncio +async def test_output_file_content_fetches_and_parses(monkeypatch): + import litellm.files.main as files_main + import litellm.proxy.openai_files_endpoints.common_utils as cu + + captured: dict = {} + + async def fake_afile_content(**kw): + captured.update(kw) + return type("R", (), {"content": b'{"a": 1}\n{"b": 2}'})() + + monkeypatch.setattr(files_main, "afile_content", fake_afile_content) + monkeypatch.setattr(cu, "_is_base64_encoded_unified_file_id", lambda fid: False) + + result = await bu._get_batch_output_file_content_as_dictionary( + _batch("file-out"), + custom_llm_provider="azure", + litellm_params={"api_key": "sk-az", "api_base": "https://az", "model": "x"}, + ) + + assert result == [{"a": 1}, {"b": 2}] + # afile_content received the file id + extracted credentials (not "model"). + assert captured["file_id"] == "file-out" + assert captured["custom_llm_provider"] == "azure" + assert captured["api_key"] == "sk-az" + assert captured["api_base"] == "https://az" + assert "model" not in captured + + +@pytest.mark.asyncio +async def test_output_file_content_unified_file_id_extraction(monkeypatch): + # a base64 unified id carries the real provider file id inside + # "llm_output_file_id,;" - it must be unwrapped before the fetch. + import litellm.files.main as files_main + import litellm.proxy.openai_files_endpoints.common_utils as cu + + 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) + monkeypatch.setattr( + cu, + "_is_base64_encoded_unified_file_id", + lambda fid: "litellm_proxy;llm_output_file_id,real-file-99;rest", + ) + + await bu._get_batch_output_file_content_as_dictionary( + _batch("encoded-blob"), custom_llm_provider="openai" + ) + + assert captured["file_id"] == "real-file-99" + + +# =========================================================================== # +# _handle_completed_batch (async orchestrator: fetch -> cost/usage/models) +# =========================================================================== # + + +@pytest.mark.asyncio +async def test_handle_completed_batch_orchestration(monkeypatch): + rows = [_success_row(model="gpt-4o", usage=_usage(10, 5))] + + async def fake_get_content(batch, custom_llm_provider, litellm_params=None): + return rows + + monkeypatch.setattr( + bu, "_get_batch_output_file_content_as_dictionary", fake_get_content + ) + monkeypatch.setattr(bu, "_batch_cost_calculator", lambda **kw: 3.3) + monkeypatch.setattr( + bu, + "_get_batch_job_total_usage_from_file_content", + lambda **kw: Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15), + ) + + cost, usage, models = await bu._handle_completed_batch( + _batch("of"), custom_llm_provider="openai" + ) + + assert cost == 3.3 + assert usage.total_tokens == 15 + assert models == ["gpt-4o"] + + +# =========================================================================== # +# Remaining branch: vertex usage disable-transform path. +# +# NOTE: the error path of _get_batch_job_cost_from_file_content (its `raise e`) +# is intentionally NOT tested: the preceding line logs via +# `verbose_logger.error("...", e)`, which passes the exception as a logging +# format-arg with no placeholder and itself raises TypeError under +# logging.raiseExceptions, masking the original error. Asserting that masked +# behavior would lock a source bug; left uncovered on purpose. +# =========================================================================== # + + +def test_total_usage_vertex_disable_transform_path(monkeypatch): + monkeypatch.setattr( + litellm, "disable_vertex_batch_output_transformation", True, raising=False + ) + monkeypatch.setattr( + bu, + "calculate_vertex_ai_batch_cost_and_usage", + lambda content, model: ( + 0.0, + Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3), + ), + ) + + usage = bu._get_batch_job_total_usage_from_file_content( + [], custom_llm_provider="vertex_ai", model_name="gemini-x" + ) + assert usage.total_tokens == 3