mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-12 03:48:56 +00:00
58258409c9
* refactor(proxy): expose public names for private proxy helpers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): keep original class names behind public aliases Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): preserve internal callback filtering Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): keep _PROXY_ class names for managed files hooks Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): keep old private names in package exports Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): preserve recursive auth helper name Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): align MCP limiter tests with server enforcement Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): match main's MCP limiter tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): keep old private names bound in importing modules Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): preserve compatibility imports through strict lint Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): use exact pyright suppression in password helper test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(proxy): add reasons to compatibility import noqa comments Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(proxy): update IN-list baseline for renamed helpers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
3335 lines
126 KiB
Python
3335 lines
126 KiB
Python
"""
|
|
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 asyncio, json, time
|
|
import logging
|
|
from types import MappingProxyType
|
|
|
|
import httpx
|
|
import pytest
|
|
import respx
|
|
from openai.types.batch import BatchRequestCounts
|
|
|
|
|
|
import litellm
|
|
import litellm.batches.batch_utils as bu
|
|
from litellm.types.utils import LiteLLMBatch, ModelInfo, Usage
|
|
from litellm.batches.batch_utils import(
|
|
_aggregate_batch_cost_usage_models,
|
|
get_file_content_as_dictionary,
|
|
_get_response_from_batch_job_output_file,
|
|
calculate_batch_cost_and_usage,
|
|
)
|
|
from litellm.cost_calculator import batch_cost_calculator
|
|
from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# 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},
|
|
}
|
|
|
|
|
|
def test_get_response_body_is_returned_without_validation():
|
|
response_body = ["provider-specific response"]
|
|
row = {"response": {"body": response_body}}
|
|
|
|
assert bu._get_response_from_batch_job_output_file(row) is response_body
|
|
|
|
|
|
@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_lines_skipped():
|
|
content = b'{"a": 1}\nnot valid json\n{"b": 2}\n'
|
|
assert bu.get_file_content_as_dictionary(content) == [{"a": 1}, {"b": 2}]
|
|
|
|
|
|
# =========================================================================== #
|
|
# _iter_batch_input_lines / _iter_batch_output_entries (JSONL parsing)
|
|
# =========================================================================== #
|
|
|
|
|
|
def test_iter_input_lines_skips_blank_and_strips():
|
|
content = b'{"a":1}\n\n \n{"b":2}\n'
|
|
assert list(bu._iter_batch_input_lines(content)) == [b'{"a":1}', b'{"b":2}']
|
|
|
|
|
|
def test_iter_input_lines_handles_missing_trailing_newline():
|
|
assert list(bu._iter_batch_input_lines(b'{"a":1}')) == [b'{"a":1}']
|
|
|
|
|
|
def test_iter_input_lines_empty():
|
|
assert list(bu._iter_batch_input_lines(b"")) == []
|
|
|
|
|
|
def test_iter_output_entries_parses_each_row():
|
|
content = b'{"body": {"model": "gpt-4o"}}\n{"body": {"model": "claude-3"}}\n'
|
|
assert list(bu._iter_batch_output_entries(content)) == [
|
|
{"body": {"model": "gpt-4o"}},
|
|
{"body": {"model": "claude-3"}},
|
|
]
|
|
|
|
|
|
def test_iter_output_entries_skips_malformed_and_non_object_lines():
|
|
content = b'{"ok": 1}\nnot-json\n[1, 2]\n{"ok": 2}\n'
|
|
assert list(bu._iter_batch_output_entries(content)) == [{"ok": 1}, {"ok": 2}]
|
|
|
|
|
|
def test_iter_output_entries_skips_undecodable_line():
|
|
content = b'{"ok": 1}\n{"note": "\xff-bad"}\n{"ok": 2}\n'
|
|
assert list(bu._iter_batch_output_entries(content)) == [{"ok": 1}, {"ok": 2}]
|
|
|
|
|
|
# =========================================================================== #
|
|
# _estimate_batch_entry_tokens (regression: an uncountable/malformed row must
|
|
# never contribute zero tokens, or a crafted batch could evade the TPM limit)
|
|
# =========================================================================== #
|
|
|
|
|
|
def test_estimate_tokens_scales_with_size():
|
|
# 4 bytes per token, floored, with a minimum of 1.
|
|
assert bu._estimate_batch_entry_tokens(b"a" * 40) == 10
|
|
|
|
|
|
def test_estimate_tokens_never_zero_for_short_rows():
|
|
assert bu._estimate_batch_entry_tokens(b"") == 1
|
|
assert bu._estimate_batch_entry_tokens(b"abc") == 1
|
|
|
|
|
|
# =========================================================================== #
|
|
# _aggregate_batch_cost_usage_models: models (output file)
|
|
# =========================================================================== #
|
|
|
|
|
|
def test_output_models_uses_model_name_override(monkeypatch):
|
|
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0)
|
|
result = bu._aggregate_batch_cost_usage_models(
|
|
entries=[_success_row(model="ignored")], custom_llm_provider="openai", model_name="forced-model"
|
|
)
|
|
assert result.models == ["forced-model"]
|
|
|
|
|
|
def test_output_models_collects_from_successful_only(monkeypatch):
|
|
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0)
|
|
rows = [
|
|
_success_row(model="gpt-4o"),
|
|
_failed_row(model="should-be-skipped"),
|
|
_success_row(model="claude-3"),
|
|
]
|
|
result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
|
|
assert result.models == ["gpt-4o", "claude-3"]
|
|
|
|
|
|
def test_output_models_skips_successful_without_model(monkeypatch):
|
|
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0)
|
|
rows = [{"response": {"status_code": 200, "body": {}}}]
|
|
result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
|
|
assert result.models == []
|
|
|
|
|
|
# =========================================================================== #
|
|
# _extract_file_access_credentials
|
|
# =========================================================================== #
|
|
|
|
|
|
def test_extract_credentials_only_known_keys():
|
|
params = {
|
|
"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",
|
|
}
|
|
assert bu._extract_file_access_credentials(params) == {
|
|
"api_key": "sk-1",
|
|
"api_base": "https://b",
|
|
"vertex_project": "proj",
|
|
"gcs_bucket_name": "my-bucket",
|
|
"bucket_name": "my-alias-bucket",
|
|
}
|
|
|
|
|
|
@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",
|
|
"gcs_bucket_name",
|
|
"bucket_name",
|
|
"s3_endpoint_url",
|
|
"s3_region_name",
|
|
"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
|
|
|
|
|
|
# =========================================================================== #
|
|
# _count_entry_tokens (per-entry rate-limit token counting). The individual
|
|
# prompt/input/embedding shapes are covered in test_batch_file_validation.py;
|
|
# here we pin the body-field precedence and the empty/fallback behavior.
|
|
# =========================================================================== #
|
|
|
|
|
|
def test_count_entry_messages_path(fake_token_counter):
|
|
entry = {"body": {"model": "gpt-4o", "messages": [{"role": "user"}, {"role": "x"}]}}
|
|
assert bu._count_entry_tokens(entry) == 2 # len(messages)
|
|
|
|
|
|
def test_count_entry_uses_dynamic_length_for_messages(fake_token_counter):
|
|
assert bu._count_entry_tokens({"body": {"messages": "abc"}}) == 3
|
|
|
|
|
|
def test_count_entry_prompt_path(fake_token_counter):
|
|
assert bu._count_entry_tokens({"body": {"model": "gpt-4o", "prompt": "abcd"}}) == 4
|
|
|
|
|
|
def test_count_entry_input_path(fake_token_counter):
|
|
assert bu._count_entry_tokens({"body": {"model": "gpt-4o", "input": "ab"}}) == 2
|
|
|
|
|
|
def test_count_entry_messages_beats_prompt(fake_token_counter):
|
|
# messages present -> prompt/input are ignored (messages is checked first).
|
|
entry = {
|
|
"body": {
|
|
"model": "gpt-4o",
|
|
"messages": [{"role": "user"}],
|
|
"prompt": "this-should-be-ignored",
|
|
}
|
|
}
|
|
assert bu._count_entry_tokens(entry) == 1
|
|
|
|
|
|
def test_count_entry_prompt_beats_input(fake_token_counter):
|
|
entry = {"body": {"model": "gpt-4o", "prompt": "abc", "input": "this-is-longer"}}
|
|
assert bu._count_entry_tokens(entry) == 3
|
|
|
|
|
|
def test_count_entry_empty_body_is_zero(fake_token_counter):
|
|
assert bu._count_entry_tokens({"body": {}}) == 0
|
|
assert bu._count_entry_tokens({}) == 0
|
|
|
|
|
|
def test_count_entry_uses_model_name_fallback(monkeypatch):
|
|
# No body.model -> the model_name argument is forwarded to the token counter.
|
|
captured = {}
|
|
|
|
def _tc(model=None, text=None, messages=None, **kw):
|
|
captured["model"] = model
|
|
return len(text or "")
|
|
|
|
monkeypatch.setattr(bu, "token_counter", _tc)
|
|
bu._count_entry_tokens({"body": {"prompt": "ab"}}, model_name="fallback-model")
|
|
assert captured["model"] == "fallback-model"
|
|
|
|
|
|
# =========================================================================== #
|
|
# _aggregate_batch_cost_usage_models: usage (output usage aggregation)
|
|
# =========================================================================== #
|
|
|
|
|
|
def test_total_usage_sums_successful_only(monkeypatch):
|
|
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0)
|
|
rows = [
|
|
_success_row(usage=_usage(10, 5)), # 15
|
|
_failed_row(), # excluded
|
|
_success_row(usage=_usage(20, 10)), # 30
|
|
]
|
|
result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
|
|
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (
|
|
30,
|
|
15,
|
|
45,
|
|
)
|
|
|
|
|
|
def test_total_usage_and_cost_normalize_mixed_responses_and_chat():
|
|
responses_row = _success_row(
|
|
usage={
|
|
"input_tokens": 20,
|
|
"output_tokens": 7,
|
|
"total_tokens": 27,
|
|
"input_tokens_details": {"cached_tokens": 3},
|
|
}
|
|
)
|
|
chat_row = _success_row(usage=_usage(10, 5))
|
|
|
|
result = bu._aggregate_batch_cost_usage_models(
|
|
entries=[responses_row, chat_row],
|
|
custom_llm_provider="openai",
|
|
model_info={
|
|
"input_cost_per_token_batches": 0.00125,
|
|
"output_cost_per_token_batches": 0.005,
|
|
},
|
|
)
|
|
|
|
assert result.usage.prompt_tokens == 30
|
|
assert result.usage.completion_tokens == 12
|
|
assert result.usage.total_tokens == 42
|
|
assert result.usage.cache_read_input_tokens == 3
|
|
assert result.cost == pytest.approx((30 * 0.00125) + (12 * 0.005))
|
|
|
|
|
|
def test_total_cost_applies_the_long_context_batch_tier_per_line():
|
|
long_row = _success_row(usage=_usage(300_000, 10))
|
|
short_row = _success_row(usage=_usage(100, 10))
|
|
|
|
result = bu._aggregate_batch_cost_usage_models(
|
|
entries=[long_row, short_row],
|
|
custom_llm_provider="openai",
|
|
model_info=ModelInfo(
|
|
key="lit-batch-tier",
|
|
max_tokens=None,
|
|
max_input_tokens=None,
|
|
max_output_tokens=None,
|
|
input_cost_per_token=2e-6,
|
|
output_cost_per_token=8e-6,
|
|
litellm_provider="openai",
|
|
mode="chat",
|
|
supported_openai_params=None,
|
|
input_cost_per_token_batches=1e-6,
|
|
output_cost_per_token_batches=4e-6,
|
|
input_cost_per_token_above_272k_tokens_batches=2e-6,
|
|
output_cost_per_token_above_272k_tokens_batches=6e-6,
|
|
),
|
|
)
|
|
|
|
assert result.cost == pytest.approx((300_000 * 2e-6) + (10 * 6e-6) + (100 * 1e-6) + (10 * 4e-6))
|
|
|
|
|
|
def test_xai_output_lines_bill_reasoning_tokens_as_completion_tokens():
|
|
row = _success_row(
|
|
model="grok-4.3",
|
|
usage={
|
|
"prompt_tokens": 615,
|
|
"completion_tokens": 3,
|
|
"total_tokens": 993,
|
|
"completion_tokens_details": {"reasoning_tokens": 375},
|
|
},
|
|
)
|
|
|
|
result = bu._aggregate_batch_cost_usage_models(
|
|
entries=[row],
|
|
custom_llm_provider="xai",
|
|
model_info=ModelInfo(
|
|
key="xai/grok-4.3",
|
|
max_tokens=None,
|
|
max_input_tokens=None,
|
|
max_output_tokens=None,
|
|
input_cost_per_token=1.25e-6,
|
|
output_cost_per_token=2.5e-6,
|
|
litellm_provider="xai",
|
|
mode="chat",
|
|
supported_openai_params=None,
|
|
input_cost_per_token_batches=1e-6,
|
|
output_cost_per_token_batches=2e-6,
|
|
),
|
|
)
|
|
|
|
assert result.usage.completion_tokens == 378
|
|
assert result.usage.total_tokens == 993
|
|
assert result.cost == pytest.approx((615 * 1e-6) + (378 * 2e-6))
|
|
|
|
|
|
def test_total_usage_empty_is_zero():
|
|
result = bu._aggregate_batch_cost_usage_models(entries=[], custom_llm_provider="openai")
|
|
assert result.cost == 0.0
|
|
assert result.models == []
|
|
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (
|
|
0,
|
|
0,
|
|
0,
|
|
)
|
|
assert result.successful_requests == 0
|
|
assert result.failed_requests == 0
|
|
|
|
|
|
def test_total_usage_includes_reasoning_tokens(monkeypatch):
|
|
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0)
|
|
rows = [
|
|
_success_row(
|
|
usage={
|
|
"prompt_tokens": 10,
|
|
"completion_tokens": 50,
|
|
"total_tokens": 60,
|
|
"completion_tokens_details": {"reasoning_tokens": 30},
|
|
}
|
|
),
|
|
_success_row(
|
|
usage={
|
|
"prompt_tokens": 5,
|
|
"completion_tokens": 20,
|
|
"total_tokens": 25,
|
|
"completion_tokens_details": {"reasoning_tokens": 8},
|
|
}
|
|
),
|
|
_failed_row(), # excluded, must not contribute reasoning tokens either
|
|
]
|
|
result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
|
|
assert result.usage.completion_tokens_details is not None
|
|
assert result.usage.completion_tokens_details.reasoning_tokens == 38
|
|
|
|
|
|
def test_aggregate_counts_successful_and_failed_requests(monkeypatch):
|
|
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0)
|
|
rows = [
|
|
_success_row(usage=_usage(10, 5)),
|
|
_failed_row(),
|
|
_success_row(usage=_usage(20, 10)),
|
|
_failed_row(),
|
|
_failed_row(),
|
|
]
|
|
result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
|
|
assert result.successful_requests == 2
|
|
assert result.failed_requests == 3
|
|
assert result.successful_requests + result.failed_requests == len(rows)
|
|
|
|
|
|
def test_aggregate_returns_batch_cost_usage_result_dataclass(monkeypatch):
|
|
import litellm.cost_calculator as cc
|
|
|
|
monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (0.4, 0.6))
|
|
result = bu._aggregate_batch_cost_usage_models(
|
|
entries=[_success_row(usage=_usage(10, 5))], custom_llm_provider="openai"
|
|
)
|
|
assert isinstance(result, bu.BatchCostUsageResult)
|
|
assert (result.cost, result.models, result.successful_requests, result.failed_requests) == (
|
|
1.0,
|
|
["gpt-4o"],
|
|
1,
|
|
0,
|
|
)
|
|
assert (result.prompt_cost, result.completion_cost) == (0.4, 0.6)
|
|
|
|
|
|
# =========================================================================== #
|
|
# _aggregate_batch_cost_usage_models: cost (cost maps mocked)
|
|
# =========================================================================== #
|
|
|
|
|
|
def test_cost_without_model_info_prices_each_row_by_its_response_model(monkeypatch):
|
|
# model_info is None -> batch_cost_calculator per successful row, model from the response body.
|
|
import litellm.cost_calculator as cc
|
|
|
|
calls = []
|
|
|
|
def _batch_cost(**kw):
|
|
calls.append(kw)
|
|
return (0.3, 0.2)
|
|
|
|
monkeypatch.setattr(cc, "batch_cost_calculator", _batch_cost)
|
|
rows = [
|
|
_success_row(usage=_usage(10, 5)),
|
|
_failed_row(), # excluded -> not costed
|
|
_success_row(usage=_usage(20, 10)),
|
|
]
|
|
|
|
result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
|
|
|
|
assert result.cost == pytest.approx(1.0) # 2 successful * (0.3 + 0.2)
|
|
assert (result.prompt_cost, result.completion_cost) == (pytest.approx(0.6), pytest.approx(0.4))
|
|
assert len(calls) == 2 # failed row not costed
|
|
assert all(call["model"] == "gpt-4o" and call["model_info"] is None for call in calls)
|
|
assert result.successful_requests == 2
|
|
assert result.failed_requests == 1
|
|
|
|
|
|
def test_empty_body_line_does_not_zero_whole_batch():
|
|
"""A status-200 row with an empty body makes litellm.completion_cost raise;
|
|
that line must be skipped from pricing instead of zeroing the whole batch.
|
|
|
|
The provider still reported it as a success, so it stays in
|
|
successful_requests and out of failed_requests - otherwise the counts stop
|
|
reconciling with the provider's own request_counts over a litellm-side
|
|
pricing gap the customer never caused."""
|
|
rows = [
|
|
_success_row(usage=_usage(10, 5)),
|
|
{
|
|
"custom_id": "request-poison-empty",
|
|
"response": {"status_code": 200, "request_id": "inject-empty-body", "body": {}},
|
|
},
|
|
_success_row(usage=_usage(20, 10)),
|
|
]
|
|
|
|
result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
|
|
|
|
assert result.cost > 0.0
|
|
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (30, 15, 45)
|
|
assert result.models == ["gpt-4o", "gpt-4o"]
|
|
assert (result.successful_requests, result.failed_requests) == (3, 0)
|
|
|
|
|
|
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)),
|
|
]
|
|
|
|
result = bu._aggregate_batch_cost_usage_models(
|
|
entries=rows,
|
|
custom_llm_provider="openai",
|
|
model_info={"input_cost_per_token": 0.0}, # type: ignore[arg-type] # truthy -> model_info path
|
|
)
|
|
|
|
assert result.cost == pytest.approx(0.6) # 2 * (0.1 + 0.2)
|
|
|
|
|
|
def test_aggregate_consumes_entries_in_a_single_pass(monkeypatch):
|
|
"""A one-shot generator: any implementation that iterates the entries twice
|
|
(e.g. separate cost and usage passes) sees nothing on the second pass and
|
|
returns wrong totals for at least one of cost/usage/models."""
|
|
import litellm.cost_calculator as cc
|
|
|
|
monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (0.25, 0.25))
|
|
one_shot = (row for row in [_success_row(usage=_usage(10, 5)), _failed_row(), _success_row(usage=_usage(20, 10))])
|
|
|
|
result = bu._aggregate_batch_cost_usage_models(entries=one_shot, custom_llm_provider="openai")
|
|
|
|
assert result.cost == 1.0
|
|
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (30, 15, 45)
|
|
assert result.models == ["gpt-4o", "gpt-4o"]
|
|
assert result.successful_requests == 2
|
|
assert result.failed_requests == 1
|
|
|
|
|
|
# =========================================================================== #
|
|
# calculate_batch_cost_and_usage (dispatch: vertex-disable-transform vs generic)
|
|
# =========================================================================== #
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_calculate_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, model_info=None: bu.BatchCostUsageResult(
|
|
cost=9.9,
|
|
usage=Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3),
|
|
models=["gemini-2.0-flash-001"],
|
|
successful_requests=1,
|
|
failed_requests=0,
|
|
),
|
|
)
|
|
# generic path must NOT be taken
|
|
monkeypatch.setattr(
|
|
bu,
|
|
"_aggregate_batch_cost_usage_models",
|
|
lambda **kw: pytest.fail("generic path should not run"),
|
|
)
|
|
|
|
result = await bu.calculate_batch_cost_and_usage(
|
|
file_content_dictionary=[], custom_llm_provider="vertex_ai", model_name="gemini-2.0-flash-001"
|
|
)
|
|
assert result.cost == 9.9
|
|
assert result.usage.total_tokens == 3
|
|
assert result.models == ["gemini-2.0-flash-001"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_calculate_vertex_disable_transform_needs_model_name(monkeypatch):
|
|
"""Without a model_name the raw-vertex path cannot price lines; the generic
|
|
aggregation path must run even with the disable flag set."""
|
|
monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", True, raising=False)
|
|
monkeypatch.setattr(
|
|
bu,
|
|
"calculate_vertex_ai_batch_cost_and_usage",
|
|
lambda content, model, model_info=None: pytest.fail("raw vertex path should not run"),
|
|
)
|
|
|
|
result = await bu.calculate_batch_cost_and_usage(file_content_dictionary=[], custom_llm_provider="vertex_ai")
|
|
assert result.cost == 0.0
|
|
assert result.usage.total_tokens == 0
|
|
assert result.models == []
|
|
|
|
|
|
# =========================================================================== #
|
|
# 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,
|
|
}
|
|
}
|
|
},
|
|
]
|
|
|
|
result = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x")
|
|
|
|
assert result.cost == pytest.approx(0.6) # 2 * (0.1 + 0.2)
|
|
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (
|
|
30,
|
|
15,
|
|
45,
|
|
)
|
|
assert result.successful_requests == 2
|
|
assert result.failed_requests == 0
|
|
|
|
|
|
def test_vertex_batch_usage_preserves_modality_token_details(monkeypatch):
|
|
monkeypatch.setitem(
|
|
litellm.model_cost,
|
|
"vertex_ai/gemini-embedding-2",
|
|
{
|
|
"input_cost_per_token_batches": 1e-7,
|
|
"input_cost_per_audio_token_batches": 3.25e-6,
|
|
"input_cost_per_image_token_batches": 2.25e-7,
|
|
"input_cost_per_video_token_batches": 6e-6,
|
|
},
|
|
)
|
|
responses = [
|
|
{
|
|
"key": "id_1",
|
|
"status": "",
|
|
"request": {"content": {"parts": [{"text": "hello"}, {"fileData": {"mimeType": "audio/wav"}}]}},
|
|
"response": {
|
|
"embedding": {"values": [0.1, 0.2]},
|
|
"usageMetadata": {
|
|
"promptTokenCount": 84,
|
|
"candidatesTokenCount": 0,
|
|
"totalTokenCount": 84,
|
|
"promptTokensDetails": [
|
|
{"modality": "AUDIO", "tokenCount": 64},
|
|
{"modality": "TEXT", "tokenCount": 20},
|
|
],
|
|
},
|
|
},
|
|
}
|
|
]
|
|
|
|
result = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-embedding-2")
|
|
|
|
assert (result.successful_requests, result.usage.prompt_tokens) == (1, 84)
|
|
assert result.prompt_cost == pytest.approx(64 * 3.25e-6 + 20 * 1e-7)
|
|
|
|
|
|
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,
|
|
}
|
|
}
|
|
},
|
|
]
|
|
|
|
result = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x")
|
|
|
|
assert result.cost == pytest.approx(1.0) # only one line costed
|
|
assert result.usage.total_tokens == 10
|
|
assert result.successful_requests == 1
|
|
assert result.failed_requests == 1
|
|
|
|
|
|
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}}}]
|
|
|
|
result = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x")
|
|
assert result.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,
|
|
}
|
|
}
|
|
}
|
|
]
|
|
|
|
result = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x")
|
|
assert result.cost == 0.0
|
|
assert result.usage.total_tokens == 10
|
|
|
|
|
|
# =========================================================================== #
|
|
# calculate_batch_cost_and_usage (async orchestrator)
|
|
# =========================================================================== #
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_calculate_batch_cost_and_usage_orchestration(monkeypatch):
|
|
import litellm.cost_calculator as cc
|
|
|
|
rows = [_success_row(model="gpt-4o", usage=_usage(10, 5))]
|
|
monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (1.5, 1.0))
|
|
|
|
result = await bu.calculate_batch_cost_and_usage(file_content_dictionary=rows, custom_llm_provider="openai")
|
|
|
|
assert result.cost == 2.5
|
|
assert (result.prompt_cost, result.completion_cost) == (1.5, 1.0)
|
|
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (10, 5, 15)
|
|
assert result.models == ["gpt-4o"]
|
|
|
|
|
|
# =========================================================================== #
|
|
# _fetch_batch_output_file_content (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,
|
|
)
|
|
|
|
|
|
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_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._fetch_batch_output_file_content(
|
|
_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",
|
|
"gcs_bucket_name": "litellm-bucket",
|
|
"model": "vertex_ai/gemini-3.6-flash",
|
|
},
|
|
)
|
|
|
|
assert bu.get_file_content_as_dictionary(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 captured["gcs_bucket_name"] == "litellm-bucket"
|
|
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._fetch_batch_output_file_content(_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_output_file_content_model_encoded_file_id_decoded_to_provider_id(monkeypatch):
|
|
import litellm.files.main as files_main
|
|
from litellm.proxy.openai_files_endpoints.common_utils import encode_file_id_with_model
|
|
|
|
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)
|
|
encoded_id = encode_file_id_with_model("file-Y3FHrMpi7uCkDpY6fgWGeR", "my-batch-model")
|
|
|
|
await bu._fetch_batch_output_file_content(_batch(encoded_id), custom_llm_provider="openai")
|
|
|
|
assert captured["file_id"] == "file-Y3FHrMpi7uCkDpY6fgWGeR"
|
|
assert captured["custom_llm_provider"] == "openai"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_output_file_content_raw_openai_file_id_passes_through(monkeypatch):
|
|
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)
|
|
|
|
await bu._fetch_batch_output_file_content(_batch("file-abc123"), custom_llm_provider="openai")
|
|
|
|
assert captured["file_id"] == "file-abc123"
|
|
|
|
|
|
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))
|
|
)
|
|
|
|
file_content = await bu._fetch_batch_output_file_content(
|
|
_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",
|
|
},
|
|
)
|
|
result = bu.get_file_content_as_dictionary(file_content)
|
|
|
|
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._fetch_batch_output_file_content(
|
|
_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
|
|
|
|
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)
|
|
|
|
result = 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"},
|
|
)
|
|
|
|
pricing = litellm.model_cost["vertex_ai/gemini-3.6-flash"]
|
|
batch_input = pricing["input_cost_per_token_batches"]
|
|
batch_output = pricing["output_cost_per_token_batches"]
|
|
|
|
assert batch_input < pricing["input_cost_per_token"]
|
|
assert batch_output < pricing["output_cost_per_token"]
|
|
assert result.cost > 0
|
|
assert result.cost == pytest.approx(30 * batch_input + 15 * batch_output)
|
|
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (30, 15, 45)
|
|
assert result.models == ["gemini-3.6-flash", "gemini-3.6-flash"]
|
|
assert result.successful_requests == 2
|
|
assert result.failed_requests == 0
|
|
|
|
|
|
@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._fetch_batch_output_file_content(_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._fetch_batch_output_file_content(
|
|
_batch("file-out"),
|
|
custom_llm_provider="azure",
|
|
litellm_params={"api_key": "sk-az", "api_base": "https://az", "model": "x"},
|
|
)
|
|
|
|
assert result == b'{"a": 1}\n{"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,<FID>;" - 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._fetch_batch_output_file_content(_batch("encoded-blob"), custom_llm_provider="openai")
|
|
|
|
assert captured["file_id"] == "real-file-99"
|
|
|
|
|
|
# =========================================================================== #
|
|
# handle_completed_batch (async orchestrator: fetch -> single-pass aggregate)
|
|
# =========================================================================== #
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_completed_batch_orchestration(monkeypatch):
|
|
rows = [_success_row(model="gpt-4o", usage=_usage(10, 5))]
|
|
|
|
async def fake_fetch(batch, custom_llm_provider, litellm_params=None):
|
|
return _vertex_jsonl(rows)
|
|
|
|
import litellm.cost_calculator as cc
|
|
|
|
monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch)
|
|
monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (2.0, 1.3))
|
|
|
|
result = await bu.handle_completed_batch(_batch("of"), custom_llm_provider="openai")
|
|
|
|
assert result.cost == 3.3
|
|
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (10, 5, 15)
|
|
assert result.models == ["gpt-4o"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_completed_batch_counts_error_file_failures(monkeypatch):
|
|
"""Regression test: OpenAI writes per-request failures (e.g. a rejected param)
|
|
to a separate error_file_id, never into the output file - so failed_requests
|
|
must include them or it silently undercounts real batch failures."""
|
|
from litellm.types.llms.openai import Batch
|
|
|
|
rows = [_success_row(model="gpt-5-mini", usage=_usage(24, 107))]
|
|
error_rows = [
|
|
{
|
|
"id": "batch_req_err1",
|
|
"custom_id": "req-2-bad",
|
|
"response": {"status_code": 400, "body": {"error": {"message": "Invalid 'temperature'"}}},
|
|
"error": None,
|
|
}
|
|
]
|
|
|
|
async def fake_fetch(batch, custom_llm_provider, litellm_params=None):
|
|
return _vertex_jsonl(rows)
|
|
|
|
async def fake_afile_content(**kw):
|
|
return type("R", (), {"content": _vertex_jsonl(error_rows)})()
|
|
|
|
import litellm.files.main as files_main
|
|
|
|
monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch)
|
|
monkeypatch.setattr(files_main, "afile_content", fake_afile_content)
|
|
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0)
|
|
|
|
batch = Batch(
|
|
id="b",
|
|
completion_window="24h",
|
|
created_at=1,
|
|
endpoint="/v1/chat/completions",
|
|
input_file_id="f",
|
|
object="batch",
|
|
status="completed",
|
|
output_file_id="of",
|
|
error_file_id="ef",
|
|
)
|
|
|
|
result = await bu.handle_completed_batch(batch, custom_llm_provider="openai")
|
|
|
|
assert result.successful_requests == 1
|
|
assert result.failed_requests == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_completed_batch_decodes_model_encoded_error_file_id(monkeypatch):
|
|
"""A model-encoded error file id must be decoded to the raw provider id before
|
|
the fetch, exactly like the output file id. Sending the encoded id straight to
|
|
the provider 404s, and the swallowed fetch failure silently reports 0 failures."""
|
|
import base64
|
|
|
|
from litellm.types.llms.openai import Batch
|
|
|
|
provider_error_file_id = "file-real-error-id"
|
|
encoded_error_file_id = "file-" + base64.urlsafe_b64encode(
|
|
f"litellm:{provider_error_file_id};model,model-abc".encode()
|
|
).decode().rstrip("=")
|
|
|
|
requested_file_ids = []
|
|
|
|
async def fake_fetch(batch, custom_llm_provider, litellm_params=None):
|
|
return _vertex_jsonl([_success_row(model="gpt-4o", usage=_usage(10, 5))])
|
|
|
|
async def fake_afile_content(**kw):
|
|
requested_file_ids.append(kw["file_id"])
|
|
return type("R", (), {"content": _vertex_jsonl([{"custom_id": "bad-1"}])})()
|
|
|
|
import litellm.files.main as files_main
|
|
|
|
monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch)
|
|
monkeypatch.setattr(files_main, "afile_content", fake_afile_content)
|
|
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0)
|
|
|
|
batch = Batch(
|
|
id="b",
|
|
completion_window="24h",
|
|
created_at=1,
|
|
endpoint="/v1/chat/completions",
|
|
input_file_id="f",
|
|
object="batch",
|
|
status="completed",
|
|
output_file_id="of",
|
|
error_file_id=encoded_error_file_id,
|
|
)
|
|
|
|
result = await bu.handle_completed_batch(batch, custom_llm_provider="openai")
|
|
|
|
assert requested_file_ids == [provider_error_file_id]
|
|
assert result.failed_requests == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_completed_batch_no_error_file_id_reports_zero_error_failures(monkeypatch):
|
|
rows = [_success_row(model="gpt-4o", usage=_usage(10, 5))]
|
|
|
|
async def fake_fetch(batch, custom_llm_provider, litellm_params=None):
|
|
return _vertex_jsonl(rows)
|
|
|
|
monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch)
|
|
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0)
|
|
|
|
result = await bu.handle_completed_batch(_batch("of"), custom_llm_provider="openai")
|
|
|
|
assert result.successful_requests == 1
|
|
assert result.failed_requests == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_completed_batch_no_output_file_is_zero(monkeypatch):
|
|
"""
|
|
Regression: an all-error batch completes with output_file_id=None (results go
|
|
to a separate error_file_id). handle_completed_batch must report an empty
|
|
result set - zero cost, zero usage, no models - instead of letting the file
|
|
fetch raise "Output file id is None" on every aretrieve_batch logging poll.
|
|
"""
|
|
|
|
# The output-file fetch must not even be attempted when there is no output file.
|
|
async def _must_not_fetch(*args, **kwargs):
|
|
pytest.fail("_fetch_batch_output_file_content should not be called")
|
|
|
|
monkeypatch.setattr(bu, "_fetch_batch_output_file_content", _must_not_fetch)
|
|
|
|
result = await bu.handle_completed_batch(_batch(None), custom_llm_provider="openai")
|
|
|
|
assert result.cost == 0.0
|
|
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (0, 0, 0)
|
|
assert result.models == []
|
|
assert result.successful_requests == 0
|
|
assert result.failed_requests == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_completed_batch_vertex_disable_transform_path(monkeypatch):
|
|
raw_rows = [{"response": {"usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 2}}}]
|
|
|
|
async def fake_fetch(batch, custom_llm_provider, litellm_params=None):
|
|
return _vertex_jsonl(raw_rows)
|
|
|
|
monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch)
|
|
monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", True, raising=False)
|
|
seen: dict = {}
|
|
|
|
def fake_vertex_calc(content, model, model_info=None):
|
|
seen["content"] = content
|
|
seen["model"] = model
|
|
return bu.BatchCostUsageResult(
|
|
cost=7.7,
|
|
usage=Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3),
|
|
models=["gemini-x"],
|
|
successful_requests=1,
|
|
failed_requests=0,
|
|
)
|
|
|
|
monkeypatch.setattr(bu, "calculate_vertex_ai_batch_cost_and_usage", fake_vertex_calc)
|
|
|
|
result = await bu.handle_completed_batch(
|
|
_batch("gs://litellm-bucket/output/predictions.jsonl"),
|
|
custom_llm_provider="vertex_ai",
|
|
model_name="gemini-x",
|
|
)
|
|
|
|
assert result.cost == 7.7
|
|
assert result.usage.total_tokens == 3
|
|
assert result.models == ["gemini-x"]
|
|
assert list(seen["content"]) == raw_rows
|
|
assert seen["model"] == "gemini-x"
|
|
|
|
|
|
def _anthropic_usage(input_tokens, output_tokens, cache_creation=0, cache_read=0):
|
|
return {
|
|
"input_tokens": input_tokens,
|
|
"output_tokens": output_tokens,
|
|
"cache_creation_input_tokens": cache_creation,
|
|
"cache_read_input_tokens": cache_read,
|
|
}
|
|
|
|
|
|
def _anthropic_succeeded_row(model="claude-sonnet-4-5-20250929", usage=None):
|
|
return {
|
|
"custom_id": "req-1",
|
|
"result": {
|
|
"type": "succeeded",
|
|
"message": {
|
|
"id": "msg_1",
|
|
"type": "message",
|
|
"role": "assistant",
|
|
"model": model,
|
|
"content": [{"type": "text", "text": "ok"}],
|
|
"stop_reason": "end_turn",
|
|
"usage": usage or _anthropic_usage(10, 5),
|
|
},
|
|
},
|
|
}
|
|
|
|
|
|
def _anthropic_errored_row():
|
|
return {
|
|
"custom_id": "req-2",
|
|
"result": {
|
|
"type": "errored",
|
|
"error": {"type": "invalid_request_error", "message": "bad request"},
|
|
},
|
|
}
|
|
|
|
|
|
_ANTHROPIC_MODEL_INFO = {
|
|
"input_cost_per_token": 3e-6,
|
|
"output_cost_per_token": 15e-6,
|
|
"cache_read_input_token_cost": 3e-7,
|
|
"cache_creation_input_token_cost": 3.75e-6,
|
|
}
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"row,expected",
|
|
[
|
|
(_anthropic_succeeded_row(), True),
|
|
(_anthropic_errored_row(), False),
|
|
({"custom_id": "x", "result": {"type": "canceled"}}, False),
|
|
({"custom_id": "x", "result": {"type": "expired"}}, False),
|
|
({"custom_id": "x"}, False),
|
|
({"custom_id": "x", "result": None}, False),
|
|
],
|
|
)
|
|
def test_anthropic_result_line_success_check(row, expected):
|
|
"""
|
|
LIT-4008 regression: anthropic batch results JSONL lines are not
|
|
OpenAI-shaped; success is result.type == "succeeded", not
|
|
response.status_code == 200. Pre-fix every anthropic line parsed as
|
|
unsuccessful, so completed batches were billed $0 forever.
|
|
"""
|
|
assert bu._batch_response_was_successful(row, custom_llm_provider="anthropic") is expected
|
|
|
|
|
|
def test_anthropic_response_body_is_result_message():
|
|
row = _anthropic_succeeded_row(model="claude-sonnet-4-5-20250929")
|
|
body = bu._get_response_from_batch_job_output_file(row, custom_llm_provider="anthropic")
|
|
assert body["model"] == "claude-sonnet-4-5-20250929"
|
|
assert body["usage"] == _anthropic_usage(10, 5)
|
|
|
|
|
|
def test_anthropic_usage_conversion_includes_cache_tokens():
|
|
body = {
|
|
"model": "claude-sonnet-4-5-20250929",
|
|
"usage": _anthropic_usage(1000, 200, cache_creation=2000, cache_read=8000),
|
|
}
|
|
usage = bu._get_batch_job_usage_from_response_body(body, custom_llm_provider="anthropic")
|
|
assert usage.prompt_tokens == 11000
|
|
assert usage.completion_tokens == 200
|
|
assert usage.total_tokens == 11200
|
|
assert usage.prompt_tokens_details.cached_tokens == 8000
|
|
assert usage.prompt_tokens_details.cache_creation_tokens == 2000
|
|
|
|
|
|
def test_bedrock_model_output_line_success_check():
|
|
row = {
|
|
"recordId": "1",
|
|
"modelOutput": {"model": "claude-sonnet-4-6", "usage": {"input_tokens": 13, "output_tokens": 5}},
|
|
}
|
|
assert bu._batch_response_was_successful(row, custom_llm_provider="bedrock") is True
|
|
assert (
|
|
bu._get_response_from_batch_job_output_file(row, custom_llm_provider="bedrock")["model"] == "claude-sonnet-4-6"
|
|
)
|
|
|
|
|
|
def test_bedrock_cost_uses_deployment_model_name():
|
|
row = {
|
|
"recordId": "1",
|
|
"modelOutput": {"model": "claude-sonnet-4-6", "usage": {"input_tokens": 13, "output_tokens": 5}},
|
|
}
|
|
result = bu._aggregate_batch_cost_usage_models(
|
|
entries=[row],
|
|
custom_llm_provider="bedrock",
|
|
model_name="us.anthropic.claude-sonnet-4-6",
|
|
model_info={},
|
|
)
|
|
assert result.cost > 0
|
|
assert result.models == ["us.anthropic.claude-sonnet-4-6"]
|
|
|
|
|
|
def test_anthropic_total_usage_sums_succeeded_only(monkeypatch):
|
|
import litellm.cost_calculator as cc
|
|
|
|
monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (0.0, 0.0))
|
|
rows = [
|
|
_anthropic_succeeded_row(usage=_anthropic_usage(10, 5)),
|
|
_anthropic_errored_row(),
|
|
_anthropic_succeeded_row(usage=_anthropic_usage(20, 10, cache_read=100)),
|
|
]
|
|
result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="anthropic")
|
|
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (130, 15, 145)
|
|
assert result.successful_requests == 2
|
|
assert result.failed_requests == 1
|
|
|
|
|
|
def test_anthropic_total_usage_aggregates_cache_token_details(monkeypatch):
|
|
import litellm.cost_calculator as cc
|
|
|
|
monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (0.0, 0.0))
|
|
rows = [
|
|
_anthropic_succeeded_row(usage=_anthropic_usage(1000, 200, cache_creation=2000, cache_read=8000)),
|
|
_anthropic_errored_row(),
|
|
_anthropic_succeeded_row(usage=_anthropic_usage(50, 20, cache_creation=300, cache_read=700)),
|
|
]
|
|
result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="anthropic")
|
|
assert result.usage.prompt_tokens_details.cached_tokens == 8700
|
|
assert result.usage.prompt_tokens_details.cache_creation_tokens == 2300
|
|
assert result.usage.cache_read_input_tokens == 8700
|
|
assert result.usage.cache_creation_input_tokens == 2300
|
|
|
|
|
|
def test_total_usage_without_cache_tokens_has_no_prompt_details(monkeypatch):
|
|
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0)
|
|
rows = [
|
|
{
|
|
"custom_id": "req-1",
|
|
"response": {
|
|
"status_code": 200,
|
|
"body": {
|
|
"model": "gpt-5.2",
|
|
"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
|
|
},
|
|
},
|
|
}
|
|
]
|
|
result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
|
|
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (10, 5, 15)
|
|
assert result.usage.prompt_tokens_details is None
|
|
|
|
|
|
def test_anthropic_cost_applies_batch_discount_and_cache_pricing():
|
|
"""Anthropic batches bill at 50% of the regular rate for base input,
|
|
cache reads, cache writes, and output tokens alike."""
|
|
rows = [
|
|
_anthropic_succeeded_row(usage=_anthropic_usage(1000, 200, cache_creation=2000, cache_read=8000)),
|
|
_anthropic_errored_row(),
|
|
]
|
|
|
|
result = bu._aggregate_batch_cost_usage_models(
|
|
entries=rows,
|
|
custom_llm_provider="anthropic",
|
|
model_info=_ANTHROPIC_MODEL_INFO, # type: ignore[arg-type]
|
|
)
|
|
|
|
expected_half_price = (1000 * 3e-6 + 8000 * 3e-7 + 2000 * 3.75e-6 + 200 * 15e-6) / 2
|
|
assert result.cost == pytest.approx(expected_half_price)
|
|
|
|
|
|
def test_anthropic_cost_without_model_info_uses_batch_cost_calculator(monkeypatch):
|
|
import litellm.cost_calculator as cc
|
|
|
|
seen = []
|
|
|
|
def _fake_batch_cost_calculator(**kw):
|
|
seen.append(kw)
|
|
return (0.1, 0.2)
|
|
|
|
monkeypatch.setattr(cc, "batch_cost_calculator", _fake_batch_cost_calculator)
|
|
monkeypatch.setattr(
|
|
litellm,
|
|
"completion_cost",
|
|
lambda **kw: pytest.fail("anthropic rows must not go through completion_cost"),
|
|
)
|
|
|
|
result = bu._aggregate_batch_cost_usage_models(
|
|
entries=[_anthropic_succeeded_row()], custom_llm_provider="anthropic"
|
|
)
|
|
|
|
assert result.cost == pytest.approx(0.3)
|
|
assert seen[0]["model"] == "claude-sonnet-4-5-20250929"
|
|
assert seen[0]["custom_llm_provider"] == "anthropic"
|
|
assert seen[0]["usage"].prompt_tokens == 10
|
|
|
|
|
|
def test_anthropic_batch_models_collected_from_succeeded_rows(monkeypatch):
|
|
import litellm.cost_calculator as cc
|
|
|
|
monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (0.0, 0.0))
|
|
rows = [
|
|
_anthropic_succeeded_row(model="claude-sonnet-4-5-20250929"),
|
|
_anthropic_errored_row(),
|
|
]
|
|
result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="anthropic")
|
|
assert result.models == ["claude-sonnet-4-5-20250929"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_calculate_batch_cost_and_usage_anthropic_end_to_end():
|
|
rows = [
|
|
_anthropic_succeeded_row(usage=_anthropic_usage(1000, 200, cache_creation=2000, cache_read=8000)),
|
|
_anthropic_errored_row(),
|
|
]
|
|
|
|
result = await bu.calculate_batch_cost_and_usage(
|
|
file_content_dictionary=rows,
|
|
custom_llm_provider="anthropic",
|
|
model_name="claude-sonnet-4-5",
|
|
model_info=_ANTHROPIC_MODEL_INFO, # type: ignore[arg-type]
|
|
)
|
|
|
|
assert result.cost == pytest.approx(1000 * 3e-6 / 2 + 8000 * 3e-7 / 2 + 2000 * 3.75e-6 / 2 + 200 * 15e-6 / 2)
|
|
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (
|
|
11000,
|
|
200,
|
|
11200,
|
|
)
|
|
assert result.models == ["claude-sonnet-4-5"]
|
|
|
|
|
|
def test_extract_credentials_forwards_the_trusted_model_credential_snapshot():
|
|
"""Bedrock resolves a batch's output bucket only from the immutable server-side
|
|
snapshot, never from a request param, so cost accounting on the retrieve path cannot
|
|
read the output file unless this key is forwarded. Without it the accounting raises
|
|
"S3 bucket_name is required" for a bucket the deployment has configured, and the
|
|
batch's cost is never recorded."""
|
|
snapshot = MappingProxyType({"s3_bucket_name": "configured-bucket", "aws_region_name": "us-east-1"})
|
|
|
|
credentials = bu._extract_file_access_credentials({"_litellm_internal_model_credentials": snapshot})
|
|
|
|
assert credentials["_litellm_internal_model_credentials"] is snapshot
|
|
|
|
|
|
def test_extract_credentials_forwards_the_deployment_aws_credentials():
|
|
"""The retrieve path's logging object carries the deployment's AWS keys in its
|
|
litellm_params, and the S3 read of the output file signs with whatever afile_content
|
|
receives. Dropping them here sent the read to the ambient credential chain, so a
|
|
deployment whose only AWS credentials live in its litellm_params never recorded
|
|
batch cost on retrieve even once the bucket resolved."""
|
|
params = {
|
|
"aws_access_key_id": "AKIA-deployment",
|
|
"aws_secret_access_key": "secret-deployment",
|
|
"aws_session_token": "token-deployment",
|
|
"aws_region_name": "us-west-2",
|
|
"aws_role_name": "arn:aws:iam::123456789012:role/batch-reader",
|
|
"model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
|
|
}
|
|
|
|
credentials = bu._extract_file_access_credentials(params)
|
|
|
|
assert credentials == {key: value for key, value in params.items() if key != "model"}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_output_file_content_bedrock_reads_with_deployment_aws_credentials(monkeypatch):
|
|
import litellm.files.main as files_main
|
|
|
|
captured: dict = {}
|
|
|
|
async def fake_afile_content(**kw):
|
|
captured.update(kw)
|
|
return type("R", (), {"content": b""})()
|
|
|
|
monkeypatch.setattr(files_main, "afile_content", fake_afile_content)
|
|
snapshot = MappingProxyType({"s3_bucket_name": "configured-bucket", "aws_region_name": "us-west-2"})
|
|
|
|
await bu._fetch_batch_output_file_content(
|
|
_batch("s3://configured-bucket/litellm-batch-outputs/job-1/out.jsonl.out"),
|
|
custom_llm_provider="bedrock",
|
|
litellm_params={
|
|
"aws_access_key_id": "AKIA-deployment",
|
|
"aws_secret_access_key": "secret-deployment",
|
|
"aws_session_token": "token-deployment",
|
|
"aws_region_name": "us-west-2",
|
|
"_litellm_internal_model_credentials": snapshot,
|
|
"model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0",
|
|
},
|
|
)
|
|
|
|
assert captured["file_id"] == "s3://configured-bucket/litellm-batch-outputs/job-1/out.jsonl.out"
|
|
assert captured["custom_llm_provider"] == "bedrock"
|
|
assert captured["aws_access_key_id"] == "AKIA-deployment"
|
|
assert captured["aws_secret_access_key"] == "secret-deployment"
|
|
assert captured["aws_session_token"] == "token-deployment"
|
|
assert captured["aws_region_name"] == "us-west-2"
|
|
assert captured["_litellm_internal_model_credentials"] is snapshot
|
|
assert "model" not in captured
|
|
|
|
|
|
# =========================================================================== #
|
|
# handle_completed_batch threads the deployment's model identity + pricing
|
|
# =========================================================================== #
|
|
|
|
|
|
def _bedrock_row(model: str, input_tokens: int, output_tokens: int) -> dict[str, object]:
|
|
return {
|
|
"modelInput": {"messages": [{"role": "user", "content": [{"type": "text", "text": "hi"}]}]},
|
|
"modelOutput": {
|
|
"model": model,
|
|
"id": "msg_1",
|
|
"type": "message",
|
|
"role": "assistant",
|
|
"content": [{"type": "text", "text": "ok"}],
|
|
"stop_reason": "end_turn",
|
|
"usage": {
|
|
"input_tokens": input_tokens,
|
|
"output_tokens": output_tokens,
|
|
"cache_creation_input_tokens": 0,
|
|
"cache_read_input_tokens": 0,
|
|
},
|
|
},
|
|
"recordId": "r",
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_completed_bedrock_batch_prices_from_deployment_model(monkeypatch) -> None:
|
|
"""A bedrock batch must price from the deployment model, not the response model."""
|
|
rows = [_bedrock_row("claude-sonnet-4-6", 18, 10)] * 100
|
|
|
|
async def fake_fetch(batch: object, custom_llm_provider: str, litellm_params: dict | None = None) -> bytes:
|
|
return _vertex_jsonl(rows)
|
|
|
|
monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch)
|
|
|
|
result = await bu.handle_completed_batch(
|
|
_batch("of"),
|
|
custom_llm_provider="bedrock",
|
|
model_name="bedrock/global.anthropic.claude-sonnet-4-6",
|
|
)
|
|
|
|
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (1800, 1000, 2800)
|
|
|
|
# The response model alone cannot price a bedrock batch: this is the $0 bug.
|
|
zero_result = await bu.handle_completed_batch(
|
|
_batch("of"),
|
|
custom_llm_provider="bedrock",
|
|
model_name=None,
|
|
)
|
|
assert zero_result.cost == 0.0
|
|
assert zero_result.usage.total_tokens == 2800
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_completed_batch_honors_deployment_pricing(monkeypatch) -> None:
|
|
"""A deployment's configured rates must win over the global cost map."""
|
|
rows = [_success_row(model="gemini-2.5-flash", usage=_usage(60, 75))]
|
|
|
|
async def fake_fetch(batch: object, custom_llm_provider: str, litellm_params: dict | None = None) -> bytes:
|
|
return _vertex_jsonl(rows)
|
|
|
|
monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch)
|
|
|
|
free_result = await bu.handle_completed_batch(
|
|
_batch("of"),
|
|
custom_llm_provider="vertex_ai",
|
|
model_name="vertex_ai/gemini-2.5-flash",
|
|
model_info={
|
|
"input_cost_per_token": 0.0,
|
|
"output_cost_per_token": 0.0,
|
|
"input_cost_per_token_batches": 0.0,
|
|
"output_cost_per_token_batches": 0.0,
|
|
},
|
|
)
|
|
assert free_result.cost == 0.0
|
|
|
|
billed_result = await bu.handle_completed_batch(
|
|
_batch("of"),
|
|
custom_llm_provider="vertex_ai",
|
|
model_name="vertex_ai/gemini-2.5-flash",
|
|
model_info=None,
|
|
)
|
|
assert billed_result.cost > 0.0
|
|
|
|
|
|
# =========================================================================== #
|
|
# _get_batch_job_usage_from_response_body: bedrock usage shapes
|
|
# =========================================================================== #
|
|
|
|
|
|
def test_bedrock_converse_shaped_batch_usage_is_parsed():
|
|
body = {
|
|
"model": "us.amazon.nova-lite-v1:0",
|
|
"usage": {"inputTokens": 2202, "outputTokens": 540, "totalTokens": 2742},
|
|
}
|
|
usage = bu._get_batch_job_usage_from_response_body(body, custom_llm_provider="bedrock")
|
|
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (2202, 540, 2742)
|
|
|
|
|
|
def test_bedrock_converse_batch_usage_totals_default_when_absent():
|
|
body = {"model": "us.amazon.nova-lite-v1:0", "usage": {"inputTokens": 10, "outputTokens": 4}}
|
|
usage = bu._get_batch_job_usage_from_response_body(body, custom_llm_provider="bedrock")
|
|
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (10, 4, 14)
|
|
|
|
|
|
def test_bedrock_converse_batch_usage_includes_cache_tokens():
|
|
body = {
|
|
"model": "us.amazon.nova-lite-v1:0",
|
|
"usage": {
|
|
"inputTokens": 100,
|
|
"outputTokens": 20,
|
|
"totalTokens": 120,
|
|
"cacheReadInputTokens": 800,
|
|
"cacheWriteInputTokens": 200,
|
|
},
|
|
}
|
|
usage = bu._get_batch_job_usage_from_response_body(body, custom_llm_provider="bedrock")
|
|
assert usage.prompt_tokens == 1100
|
|
assert usage.completion_tokens == 20
|
|
assert usage.prompt_tokens_details.cached_tokens == 800
|
|
assert usage.prompt_tokens_details.cache_creation_tokens == 200
|
|
|
|
|
|
def test_bedrock_anthropic_shaped_batch_usage_still_parsed():
|
|
"""Anthropic-shaped bedrock output (what an Anthropic model's batch emits) must not regress."""
|
|
body = {"model": "claude-sonnet-4-6", "usage": {"input_tokens": 18, "output_tokens": 10}}
|
|
usage = bu._get_batch_job_usage_from_response_body(body, custom_llm_provider="bedrock")
|
|
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (18, 10, 28)
|
|
|
|
|
|
def test_bedrock_titan_embedding_batch_usage_is_parsed():
|
|
"""Titan embedding batch lines carry a top-level inputTextTokenCount and no usage block."""
|
|
body = {"embedding": [0.1, 0.2], "embeddingsByType": {"float": [0.1, 0.2]}, "inputTextTokenCount": 17}
|
|
usage = bu._get_batch_job_usage_from_response_body(body, custom_llm_provider="bedrock")
|
|
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (17, 0, 17)
|
|
|
|
|
|
def test_bedrock_titan_embedding_batch_is_billed():
|
|
"""Binary embedding rows carry only embeddingsByType and must bill like float rows."""
|
|
rows = [
|
|
{"recordId": "0", "modelOutput": {"embedding": [0.1], "inputTextTokenCount": 10}},
|
|
{"recordId": "1", "modelOutput": {"embeddingsByType": {"binary": [1, 0]}, "inputTextTokenCount": 7}},
|
|
]
|
|
result = bu._aggregate_batch_cost_usage_models(
|
|
entries=rows,
|
|
custom_llm_provider="bedrock",
|
|
model_name="amazon.titan-embed-text-v2:0",
|
|
model_info={"input_cost_per_token_batches": 1e-6, "output_cost_per_token_batches": 0.0},
|
|
)
|
|
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (17, 0, 17)
|
|
assert result.cost == pytest.approx(17 * 1e-6)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"body",
|
|
[
|
|
{"embedding": [0.1], "inputTextTokenCount": "17"},
|
|
{"embedding": [0.1], "inputTextTokenCount": True},
|
|
{"embedding": [0.1], "inputTextTokenCount": None},
|
|
{"results": [{"outputText": "hi", "tokenCount": 2}], "inputTextTokenCount": 17},
|
|
],
|
|
)
|
|
def test_bedrock_input_text_token_count_outside_embedding_lines_is_not_billed(body):
|
|
"""Only embedding lines are parsed here; Titan text generation lines are left as they were."""
|
|
usage = bu._get_batch_job_usage_from_response_body(body, custom_llm_provider="bedrock")
|
|
assert usage.total_tokens == 0
|
|
|
|
|
|
def test_unparsable_bedrock_batch_usage_warns(caplog):
|
|
"""An unrecognized usage shape must be visible, not a silent $0."""
|
|
body = {"model": "amazon.titan-text-lite-v1", "usage": {"inputTextTokenCount": 42}}
|
|
with caplog.at_level(logging.WARNING):
|
|
usage = bu._get_batch_job_usage_from_response_body(body, custom_llm_provider="bedrock")
|
|
assert usage.total_tokens == 0
|
|
assert "does not understand" in caplog.text
|
|
assert "inputTextTokenCount" in caplog.text
|
|
|
|
|
|
class TestFileAccessCredentialsCarryFederation:
|
|
"""A federated deployment holds no api_key, so the fetch that reads a finished batch's output
|
|
has to inherit the federation fields or it cannot authenticate and the batch is never billed."""
|
|
|
|
def test_federation_fields_survive_extraction(self):
|
|
from litellm.batches.batch_utils import extract_file_access_credentials
|
|
|
|
credentials = extract_file_access_credentials(
|
|
{
|
|
"model": "anthropic/claude-sonnet-4-5",
|
|
"anthropic_federation_rule_id": "fdrl_x",
|
|
"anthropic_organization_id": "org-x",
|
|
"anthropic_identity_token_file": "/var/run/secrets/anthropic.com/token",
|
|
"something_unrelated": "dropped",
|
|
}
|
|
)
|
|
|
|
assert credentials["anthropic_federation_rule_id"] == "fdrl_x"
|
|
assert credentials["anthropic_organization_id"] == "org-x"
|
|
assert credentials["anthropic_identity_token_file"] == "/var/run/secrets/anthropic.com/token"
|
|
assert "something_unrelated" not in credentials
|
|
|
|
def test_every_federation_field_is_carried(self):
|
|
"""Derived from the kwargs set, so a new federation field is carried without an edit here."""
|
|
from litellm.batches.batch_utils import extract_file_access_credentials
|
|
from litellm.litellm_core_utils.get_litellm_params import ANTHROPIC_WIF_KWARGS_KEYS
|
|
|
|
params = {name: f"value-{name}" for name in ANTHROPIC_WIF_KWARGS_KEYS}
|
|
|
|
credentials = extract_file_access_credentials(params)
|
|
|
|
assert set(credentials) == set(ANTHROPIC_WIF_KWARGS_KEYS)
|
|
|
|
|
|
def test_total_cost_bills_cached_tokens_per_line_at_the_batch_cached_rate():
|
|
responses_row = _success_row(
|
|
usage={
|
|
"input_tokens": 300_000,
|
|
"output_tokens": 10,
|
|
"total_tokens": 300_010,
|
|
"input_tokens_details": {"cached_tokens": 299_000},
|
|
}
|
|
)
|
|
chat_row = _success_row(usage={**_usage(100, 10), "prompt_tokens_details": {"cached_tokens": 60}})
|
|
|
|
result = bu._aggregate_batch_cost_usage_models(
|
|
entries=[responses_row, chat_row],
|
|
custom_llm_provider="openai",
|
|
model_info=ModelInfo(
|
|
key="lit-batch-cached-tier",
|
|
max_tokens=None,
|
|
max_input_tokens=None,
|
|
max_output_tokens=None,
|
|
input_cost_per_token=2e-6,
|
|
output_cost_per_token=8e-6,
|
|
cache_read_input_token_cost=1e-6,
|
|
litellm_provider="openai",
|
|
mode="chat",
|
|
supported_openai_params=None,
|
|
input_cost_per_token_batches=1e-6,
|
|
output_cost_per_token_batches=4e-6,
|
|
cache_read_input_token_cost_batches=5e-7,
|
|
input_cost_per_token_above_272k_tokens_batches=2e-6,
|
|
output_cost_per_token_above_272k_tokens_batches=6e-6,
|
|
cache_read_input_token_cost_above_272k_tokens_batches=1e-6,
|
|
),
|
|
)
|
|
|
|
long_line = 1_000 * 2e-6 + 299_000 * 1e-6 + 10 * 6e-6
|
|
short_line = 40 * 1e-6 + 60 * 5e-7 + 10 * 4e-6
|
|
assert result.cost == pytest.approx(long_line + short_line)
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# batch_cost_is_final
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
def _retrieved_batch(
|
|
status: str, output_file_id: str | None = None, counts: BatchRequestCounts | None = None
|
|
) -> LiteLLMBatch:
|
|
return LiteLLMBatch(
|
|
id="batch_abc",
|
|
completion_window="24h",
|
|
created_at=1,
|
|
endpoint="/v1/chat/completions",
|
|
input_file_id="file-in",
|
|
object="batch",
|
|
status="validating",
|
|
output_file_id=output_file_id,
|
|
request_counts=counts,
|
|
).model_copy(update={"status": status})
|
|
|
|
|
|
class TestBatchCostIsFinal:
|
|
"""Every retrieve of one batch writes the same spend row, so the first retrieve
|
|
that prices it decides the row for good. A poll before the output exists must
|
|
therefore not count as final: pricing it recorded $0 and pinned it (LIT-7048)."""
|
|
|
|
@pytest.mark.parametrize("status", ["validating", "in_progress", "finalizing", "cancelling"])
|
|
def test_in_flight_batch_is_not_final(self, status):
|
|
assert bu.batch_cost_is_final(_retrieved_batch(status)) is False
|
|
|
|
@pytest.mark.parametrize("status", ["completed", "complete"])
|
|
def test_completed_with_output_is_final(self, status):
|
|
assert bu.batch_cost_is_final(_retrieved_batch(status, output_file_id="file-out")) is True
|
|
|
|
def test_completed_without_output_and_unknown_counts_is_not_final(self):
|
|
assert bu.batch_cost_is_final(_retrieved_batch("completed")) is False
|
|
|
|
def test_completed_without_output_and_zero_counts_is_not_final(self):
|
|
counts = BatchRequestCounts(total=0, completed=0, failed=0)
|
|
assert bu.batch_cost_is_final(_retrieved_batch("completed", counts=counts)) is False
|
|
|
|
def test_completed_without_output_but_successful_lines_is_not_final(self):
|
|
counts = BatchRequestCounts(total=2, completed=2, failed=0)
|
|
assert bu.batch_cost_is_final(_retrieved_batch("completed", counts=counts)) is False
|
|
|
|
@pytest.mark.parametrize("status", ["completed", "complete"])
|
|
def test_completed_without_output_and_every_line_failed_is_final(self, status):
|
|
counts = BatchRequestCounts(total=2, completed=0, failed=2)
|
|
assert bu.batch_cost_is_final(_retrieved_batch(status, counts=counts)) is True
|
|
|
|
@pytest.mark.parametrize("status", ["failed", "expired", "cancelled"])
|
|
def test_other_terminal_statuses_are_final(self, status):
|
|
assert bu.batch_cost_is_final(_retrieved_batch(status)) is True
|
|
|
|
|
|
def _ocr_row(pages_processed, annotation_pages=None, model="mistral-ocr-latest"):
|
|
usage_info = {"pages_processed": pages_processed, "doc_size_bytes": 4096}
|
|
if annotation_pages is not None:
|
|
usage_info["pages_processed_annotation"] = annotation_pages
|
|
return _success_row(
|
|
model=model, pages=[{"index": i, "markdown": "x"} for i in range(pages_processed)], usage_info=usage_info
|
|
)
|
|
|
|
|
|
def test_ocr_rows_are_priced_per_page_at_batch_rate(monkeypatch):
|
|
monkeypatch.setattr(
|
|
litellm,
|
|
"get_model_info",
|
|
lambda model, custom_llm_provider=None: {"ocr_cost_per_page": 0.004, "ocr_cost_per_page_batches": 0.002},
|
|
)
|
|
result = bu._aggregate_batch_cost_usage_models(
|
|
entries=[_ocr_row(3), _ocr_row(5), _failed_row(model="mistral-ocr-latest")],
|
|
custom_llm_provider="mistral",
|
|
model_name="mistral/mistral-ocr-latest",
|
|
)
|
|
assert result.cost == pytest.approx(8 * 0.002)
|
|
assert result.prompt_cost == pytest.approx(8 * 0.002)
|
|
assert result.completion_cost == 0.0
|
|
assert (result.successful_requests, result.failed_requests) == (2, 1)
|
|
assert result.usage.total_tokens == 0
|
|
assert result.models == ["mistral/mistral-ocr-latest"]
|
|
|
|
|
|
def test_ocr_rows_fall_back_to_sync_page_rate_without_batch_price(monkeypatch):
|
|
monkeypatch.setattr(litellm, "get_model_info", lambda model, custom_llm_provider=None: {"ocr_cost_per_page": 0.004})
|
|
result = bu._aggregate_batch_cost_usage_models(entries=[_ocr_row(2)], custom_llm_provider="mistral")
|
|
assert result.cost == pytest.approx(2 * 0.004)
|
|
|
|
|
|
def test_ocr_rows_bill_annotation_pages_separately(monkeypatch):
|
|
monkeypatch.setattr(
|
|
litellm,
|
|
"get_model_info",
|
|
lambda model, custom_llm_provider=None: {
|
|
"ocr_cost_per_page_batches": 0.002,
|
|
"annotation_cost_per_page_batches": 0.0025,
|
|
},
|
|
)
|
|
result = bu._aggregate_batch_cost_usage_models(
|
|
entries=[_ocr_row(4, annotation_pages=4)], custom_llm_provider="mistral"
|
|
)
|
|
assert result.cost == pytest.approx(4 * 0.002 + 4 * 0.0025)
|
|
|
|
|
|
def test_ocr_rows_use_deployment_model_info_pricing_over_cost_map(monkeypatch):
|
|
monkeypatch.setattr(
|
|
litellm, "get_model_info", lambda model, custom_llm_provider=None: pytest.fail("cost map must not be consulted")
|
|
)
|
|
result = bu._aggregate_batch_cost_usage_models(
|
|
entries=[_ocr_row(10)],
|
|
custom_llm_provider="mistral",
|
|
model_info={"ocr_cost_per_page_batches": 0.001},
|
|
)
|
|
assert result.cost == pytest.approx(0.01)
|
|
|
|
|
|
def test_ocr_rows_keep_the_published_page_rate_when_the_deployment_prices_only_annotations(monkeypatch):
|
|
monkeypatch.setattr(
|
|
litellm,
|
|
"get_model_info",
|
|
lambda model, custom_llm_provider=None: {
|
|
"ocr_cost_per_page_batches": 0.002,
|
|
"annotation_cost_per_page_batches": 0.0025,
|
|
},
|
|
)
|
|
result = bu._aggregate_batch_cost_usage_models(
|
|
entries=[_ocr_row(4, annotation_pages=4)],
|
|
custom_llm_provider="mistral",
|
|
model_info={"annotation_cost_per_page_batches": 0.01},
|
|
)
|
|
assert result.cost == pytest.approx(4 * 0.002 + 4 * 0.01)
|
|
|
|
|
|
def test_ocr_rows_keep_the_deployment_page_rate_when_the_unmapped_model_has_no_annotation_price(monkeypatch):
|
|
def _unmapped(model, custom_llm_provider=None):
|
|
raise Exception(f"This model isn't mapped yet: {model}")
|
|
|
|
monkeypatch.setattr(litellm, "get_model_info", _unmapped)
|
|
result = bu._aggregate_batch_cost_usage_models(
|
|
entries=[_ocr_row(4, annotation_pages=4, model="my-private-ocr-model")],
|
|
custom_llm_provider="mistral",
|
|
model_info={"ocr_cost_per_page_batches": 0.001},
|
|
)
|
|
assert result.cost == pytest.approx(4 * 0.001 + 4 * 0.001)
|
|
|
|
|
|
def test_ocr_rows_bill_the_deployment_sync_page_rate_over_the_published_batch_rate(monkeypatch):
|
|
monkeypatch.setattr(
|
|
litellm, "get_model_info", lambda model, custom_llm_provider=None: pytest.fail("cost map must not be consulted")
|
|
)
|
|
result = bu._aggregate_batch_cost_usage_models(
|
|
entries=[_ocr_row(3)],
|
|
custom_llm_provider="mistral",
|
|
model_info={"ocr_cost_per_page": 0.0912},
|
|
)
|
|
assert result.cost == pytest.approx(3 * 0.0912)
|
|
|
|
|
|
def test_ocr_rows_without_pricing_bill_zero_but_count_as_successful(monkeypatch):
|
|
monkeypatch.setattr(litellm, "get_model_info", lambda model, custom_llm_provider=None: {"mode": "ocr"})
|
|
result = bu._aggregate_batch_cost_usage_models(entries=[_ocr_row(3)], custom_llm_provider="mistral")
|
|
assert result.cost == 0.0
|
|
assert (result.successful_requests, result.failed_requests) == (1, 0)
|
|
|
|
|
|
def test_chat_rows_from_mistral_still_use_token_pricing(monkeypatch):
|
|
monkeypatch.setattr(
|
|
litellm,
|
|
"get_model_info",
|
|
lambda model, custom_llm_provider=None: {"input_cost_per_token": 0.001, "output_cost_per_token": 0.002},
|
|
)
|
|
result = bu._aggregate_batch_cost_usage_models(
|
|
entries=[_success_row(model="mistral-small-latest", usage=_usage(10, 5))],
|
|
custom_llm_provider="mistral",
|
|
)
|
|
assert result.cost == pytest.approx((10 * 0.001 + 5 * 0.002) / 2)
|
|
assert result.usage.total_tokens == 15
|
|
|
|
|
|
GROUNDED_USAGE_METADATA = {
|
|
"promptTokenCount": 19,
|
|
"candidatesTokenCount": 59,
|
|
"thoughtsTokenCount": 406,
|
|
"toolUsePromptTokenCount": 73,
|
|
"totalTokenCount": 557,
|
|
"promptTokensDetails": [{"modality": "TEXT", "tokenCount": 19}],
|
|
"candidatesTokensDetails": [{"modality": "TEXT", "tokenCount": 59}],
|
|
"toolUsePromptTokensDetails": [{"modality": "TEXT", "tokenCount": 73}],
|
|
"trafficType": "ON_DEMAND",
|
|
}
|
|
|
|
|
|
PASSTHROUGH_OUTPUT_URI = (
|
|
"gs://litellm-bucket/litellm-vertex-files/passthrough/publishers/google/models/gemini-2.5-flash/u/"
|
|
"predictions.jsonl"
|
|
)
|
|
|
|
|
|
UNGROUNDED_USAGE_METADATA = {
|
|
"promptTokenCount": 20,
|
|
"candidatesTokenCount": 48,
|
|
"thoughtsTokenCount": 195,
|
|
"toolUsePromptTokenCount": 73,
|
|
"totalTokenCount": 336,
|
|
"promptTokensDetails": [{"modality": "TEXT", "tokenCount": 20}],
|
|
"trafficType": "ON_DEMAND",
|
|
}
|
|
|
|
|
|
def _native_vertex_row(usage_metadata: dict, *, grounded: bool, model_version: str | None = "gemini-2.5-flash"):
|
|
candidate = {"content": {"role": "model", "parts": [{"text": "ok"}]}, "finishReason": "STOP"}
|
|
grounding = {"groundingMetadata": {"webSearchQueries": ["q"]}} if grounded else {}
|
|
response = {"candidates": [{**candidate, **grounding}], "usageMetadata": usage_metadata}
|
|
return {
|
|
"request": {"contents": [{"role": "user", "parts": [{"text": "q"}]}], "tools": [{"googleSearch": {}}]},
|
|
"status": "",
|
|
"response": {**response, **({"modelVersion": model_version} if model_version else {})},
|
|
"processed_time": "2026-09-23T19:02:00.000+00:00",
|
|
}
|
|
|
|
|
|
def _capture_cost_calls(monkeypatch, prompt_cost=0.5, completion_cost=0.25) -> list:
|
|
import litellm.cost_calculator as cc
|
|
|
|
calls: list = []
|
|
|
|
def _calc(**kw):
|
|
calls.append(kw)
|
|
return (prompt_cost, completion_cost)
|
|
|
|
monkeypatch.setattr(cc, "batch_cost_calculator", _calc)
|
|
return calls
|
|
|
|
|
|
def test_vertex_native_cost_bills_embedding_rows(monkeypatch):
|
|
monkeypatch.setitem(litellm.model_cost, "vertex_ai/gemini-embedding-2", {"input_cost_per_token_batches": 1e-7})
|
|
rows = [
|
|
{
|
|
"key": "id_1",
|
|
"status": "",
|
|
"request": {"content": {"parts": [{"text": "hello world"}]}},
|
|
"response": {"embedding": {"values": [0.1, 0.2]}, "usageMetadata": {"promptTokenCount": 2}},
|
|
},
|
|
{
|
|
"key": "id_2",
|
|
"status": "",
|
|
"request": {"content": {"parts": [{"text": "hello"}]}},
|
|
"response": {"embedding": {"values": [0.3]}, "tokenCount": "3"},
|
|
},
|
|
{"key": "id_3", "status": "INVALID_ARGUMENT", "request": {"content": {"parts": [{"text": ""}]}}},
|
|
]
|
|
|
|
result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-embedding-2")
|
|
|
|
assert (result.successful_requests, result.failed_requests) == (2, 1)
|
|
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (5, 0, 5)
|
|
assert result.cost == pytest.approx(5 * 1e-7)
|
|
assert result.models == ["gemini-embedding-2"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_native_vertex_rows_route_to_vertex_cost_path_without_flag(monkeypatch):
|
|
monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", False, raising=False)
|
|
monkeypatch.setattr(
|
|
bu, "_aggregate_batch_cost_usage_models", lambda **kw: pytest.fail("generic path should not run")
|
|
)
|
|
calls = _capture_cost_calls(monkeypatch)
|
|
rows = [
|
|
_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True),
|
|
_native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False),
|
|
]
|
|
|
|
result = await bu.calculate_batch_cost_and_usage(
|
|
file_content_dictionary=rows, custom_llm_provider="vertex_ai", model_name="gemini-2.5-flash"
|
|
)
|
|
|
|
assert result.cost == pytest.approx(1.5)
|
|
assert (result.successful_requests, result.failed_requests) == (2, 0)
|
|
assert result.models == ["gemini-2.5-flash"]
|
|
assert {(call["model"], call["custom_llm_provider"]) for call in calls} == {("gemini-2.5-flash", "vertex_ai")}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_openai_shaped_vertex_rows_keep_the_generic_path_without_flag(monkeypatch):
|
|
monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", False, raising=False)
|
|
monkeypatch.setattr(
|
|
bu, "calculate_vertex_ai_batch_cost_and_usage", lambda *a, **kw: pytest.fail("native path should not run")
|
|
)
|
|
_capture_cost_calls(monkeypatch)
|
|
rows = [_vertex_openai_row("request-1", "gemini-2.5-flash", 10, 5)]
|
|
|
|
result = await bu.calculate_batch_cost_and_usage(
|
|
file_content_dictionary=rows, custom_llm_provider="vertex_ai", model_name="gemini-2.5-flash"
|
|
)
|
|
|
|
assert result.successful_requests == 1
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_native_vertex_rows_on_another_provider_keep_the_generic_path(monkeypatch):
|
|
monkeypatch.setattr(
|
|
bu, "calculate_vertex_ai_batch_cost_and_usage", lambda *a, **kw: pytest.fail("native path should not run")
|
|
)
|
|
_capture_cost_calls(monkeypatch)
|
|
|
|
result = await bu.calculate_batch_cost_and_usage(
|
|
file_content_dictionary=[_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)],
|
|
custom_llm_provider="openai",
|
|
)
|
|
|
|
assert result.successful_requests == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_handle_completed_batch_routes_native_rows_without_flag(monkeypatch):
|
|
monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", False, raising=False)
|
|
raw_rows = [_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)]
|
|
|
|
async def fake_fetch(batch, custom_llm_provider, litellm_params=None):
|
|
return _vertex_jsonl(raw_rows)
|
|
|
|
monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch)
|
|
monkeypatch.setattr(
|
|
bu, "_aggregate_batch_cost_usage_models", lambda **kw: pytest.fail("generic path should not run")
|
|
)
|
|
calls = _capture_cost_calls(monkeypatch, prompt_cost=0.7, completion_cost=0.3)
|
|
deployment_model_info = {"input_cost_per_token_batches": 1e-6, "output_cost_per_token_batches": 2e-6}
|
|
|
|
result = await bu.handle_completed_batch(
|
|
_batch(PASSTHROUGH_OUTPUT_URI),
|
|
custom_llm_provider="vertex_ai",
|
|
model_name="gemini-2.5-flash",
|
|
model_info=deployment_model_info,
|
|
)
|
|
|
|
assert result.cost == pytest.approx(1.0)
|
|
assert result.usage.total_tokens == 557
|
|
assert [call["model_info"] for call in calls] == [deployment_model_info]
|
|
|
|
|
|
def test_native_vertex_usage_is_billed_like_the_online_path(monkeypatch):
|
|
calls = _capture_cost_calls(monkeypatch)
|
|
grounded = _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)
|
|
ungrounded = _native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False)
|
|
|
|
result = bu.calculate_vertex_ai_batch_cost_and_usage([grounded, ungrounded], "gemini-2.5-flash")
|
|
|
|
grounded_usage, ungrounded_usage = (call["usage"] for call in calls)
|
|
assert grounded_usage.prompt_tokens == 19
|
|
assert grounded_usage.completion_tokens == 59 + 406
|
|
assert grounded_usage.completion_tokens_details.reasoning_tokens == 406
|
|
assert ungrounded_usage.prompt_tokens == 20 + 73
|
|
assert ungrounded_usage.completion_tokens == 48 + 195
|
|
assert (result.usage.prompt_tokens, result.usage.completion_tokens, result.usage.total_tokens) == (
|
|
19 + 93,
|
|
465 + 243,
|
|
557 + 336,
|
|
)
|
|
|
|
|
|
def test_native_vertex_rows_are_priced_by_model_version_without_a_model_name(monkeypatch):
|
|
calls = _capture_cost_calls(monkeypatch)
|
|
rows = [
|
|
_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-flash"),
|
|
_native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version="gemini-2.5-pro"),
|
|
_native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version=None),
|
|
]
|
|
|
|
result = bu.calculate_vertex_ai_batch_cost_and_usage(rows)
|
|
|
|
assert [call["model"] for call in calls] == ["gemini-2.5-flash", "gemini-2.5-pro"]
|
|
assert result.models == ["gemini-2.5-flash", "gemini-2.5-pro"]
|
|
assert result.cost == pytest.approx(1.5)
|
|
assert result.successful_requests == 3
|
|
assert result.usage.total_tokens == 557 + 336 + 336
|
|
|
|
|
|
def test_native_vertex_rows_without_usage_metadata_count_as_failed(monkeypatch):
|
|
_capture_cost_calls(monkeypatch)
|
|
rows = [
|
|
{"request": {"contents": []}, "status": "Error: bad request", "processed_time": "t"},
|
|
{"request": {"contents": []}, "response": {"candidates": []}},
|
|
_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True),
|
|
]
|
|
|
|
result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash")
|
|
|
|
assert (result.successful_requests, result.failed_requests) == (1, 2)
|
|
assert result.usage.total_tokens == 557
|
|
|
|
|
|
def test_native_vertex_batch_whose_rows_all_failed_still_names_the_deployment_model(monkeypatch):
|
|
calls = _capture_cost_calls(monkeypatch)
|
|
rows = [{"request": {"contents": []}, "status": "Error: quota exceeded", "processed_time": "t"}] * 2
|
|
|
|
result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash")
|
|
|
|
assert result.models == ["gemini-2.5-flash"]
|
|
assert (result.successful_requests, result.failed_requests, result.cost) == (0, 2, 0.0)
|
|
assert calls == []
|
|
|
|
|
|
def test_native_vertex_rows_are_priced_with_the_deployment_model_info(monkeypatch):
|
|
calls = _capture_cost_calls(monkeypatch)
|
|
deployment_model_info = {"input_cost_per_token_batches": 1e-6, "output_cost_per_token_batches": 2e-6}
|
|
|
|
bu.calculate_vertex_ai_batch_cost_and_usage(
|
|
[_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)],
|
|
"gemini-2.5-flash",
|
|
model_info=deployment_model_info,
|
|
)
|
|
|
|
assert [call["model_info"] for call in calls] == [deployment_model_info]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_native_vertex_rows_keep_the_deployment_model_info_through_the_batch_entrypoint(monkeypatch):
|
|
calls = _capture_cost_calls(monkeypatch)
|
|
deployment_model_info = {"input_cost_per_token_batches": 1e-6}
|
|
|
|
await bu.calculate_batch_cost_and_usage(
|
|
file_content_dictionary=[_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True)],
|
|
custom_llm_provider="vertex_ai",
|
|
model_name="gemini-2.5-flash",
|
|
model_info=deployment_model_info,
|
|
)
|
|
|
|
assert [call["model_info"] for call in calls] == [deployment_model_info]
|
|
|
|
|
|
def test_native_vertex_rows_are_priced_by_the_deployment_model_over_model_version(monkeypatch):
|
|
calls = _capture_cost_calls(monkeypatch)
|
|
rows = [_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-pro")]
|
|
|
|
result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash")
|
|
|
|
assert [call["model"] for call in calls] == ["gemini-2.5-flash"]
|
|
assert result.models == ["gemini-2.5-flash"]
|
|
|
|
|
|
def test_native_vertex_rows_that_fail_response_validation_count_as_failed(monkeypatch):
|
|
calls = _capture_cost_calls(monkeypatch)
|
|
rows = [
|
|
{"request": {"contents": []}, "response": {"candidates": "nope", "usageMetadata": GROUNDED_USAGE_METADATA}},
|
|
_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True),
|
|
]
|
|
|
|
result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, "gemini-2.5-flash")
|
|
|
|
assert (result.successful_requests, result.failed_requests) == (1, 1)
|
|
assert result.usage.total_tokens == 557
|
|
assert len(calls) == 1
|
|
|
|
|
|
@pytest.mark.parametrize("wildcard_model", ["*", "vertex_ai/*"])
|
|
def test_native_vertex_rows_under_a_wildcard_deployment_are_priced_by_model_version(monkeypatch, wildcard_model):
|
|
calls = _capture_cost_calls(monkeypatch)
|
|
rows = [
|
|
_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-flash"),
|
|
_native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version=None),
|
|
]
|
|
|
|
result = bu.calculate_vertex_ai_batch_cost_and_usage(rows, wildcard_model)
|
|
|
|
assert [call["model"] for call in calls] == ["gemini-2.5-flash", wildcard_model]
|
|
assert result.cost == pytest.approx(1.5)
|
|
assert (result.successful_requests, result.failed_requests) == (2, 0)
|
|
assert result.usage.total_tokens == 557 + 336
|
|
|
|
|
|
def test_native_vertex_row_without_model_version_under_a_wildcard_deployment_bills_its_explicit_prices():
|
|
deployment_model_info = {"input_cost_per_token_batches": 1e-6, "output_cost_per_token_batches": 2e-6}
|
|
with_version = _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-2.5-flash")
|
|
without_version = _native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version=None)
|
|
|
|
twin = bu.calculate_vertex_ai_batch_cost_and_usage([with_version], "vertex_ai/*", model_info=deployment_model_info)
|
|
both = bu.calculate_vertex_ai_batch_cost_and_usage(
|
|
[with_version, without_version], "vertex_ai/*", model_info=deployment_model_info
|
|
)
|
|
|
|
assert twin.cost > 0
|
|
assert both.cost == pytest.approx(2 * twin.cost)
|
|
assert (both.successful_requests, both.failed_requests) == (2, 0)
|
|
|
|
|
|
def test_native_vertex_row_the_cost_map_cannot_price_is_billed_at_zero_and_the_rest_still_bills(monkeypatch):
|
|
import litellm.cost_calculator as cc
|
|
|
|
def _calc(**kw):
|
|
if kw["model"] == "gemini-unpriced":
|
|
raise ValueError("no pricing")
|
|
return (0.5, 0.25)
|
|
|
|
monkeypatch.setattr(cc, "batch_cost_calculator", _calc)
|
|
rows = [
|
|
_native_vertex_row(GROUNDED_USAGE_METADATA, grounded=True, model_version="gemini-unpriced"),
|
|
_native_vertex_row(UNGROUNDED_USAGE_METADATA, grounded=False, model_version="gemini-2.5-flash"),
|
|
]
|
|
|
|
result = bu.calculate_vertex_ai_batch_cost_and_usage(rows)
|
|
|
|
assert result.cost == pytest.approx(0.75)
|
|
assert (result.successful_requests, result.failed_requests) == (2, 0)
|
|
assert result.usage.total_tokens == 557 + 336
|
|
assert result.models == ["gemini-unpriced", "gemini-2.5-flash"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_flag_sends_every_vertex_row_down_the_native_path_when_a_model_is_known(monkeypatch):
|
|
monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", True, raising=False)
|
|
monkeypatch.setattr(
|
|
bu, "_aggregate_batch_cost_usage_models", lambda **kw: pytest.fail("generic path should not run")
|
|
)
|
|
calls = _capture_cost_calls(monkeypatch)
|
|
rows = [_vertex_openai_row("request-1", "gemini-2.5-flash", 10, 5)]
|
|
|
|
result = await bu.calculate_batch_cost_and_usage(
|
|
file_content_dictionary=rows, custom_llm_provider="vertex_ai", model_name="gemini-2.5-flash"
|
|
)
|
|
|
|
assert calls == []
|
|
assert (result.successful_requests, result.failed_requests) == (0, 1)
|
|
|
|
|
|
@pytest.fixture()
|
|
def _vcr_outcome_gate(request, vcr):
|
|
install_live_call_probe(request, vcr)
|
|
yield
|
|
record_vcr_outcome(request, vcr)
|
|
|
|
@pytest.fixture(scope="session")
|
|
def event_loop():
|
|
try:
|
|
loop = asyncio.get_running_loop()
|
|
except RuntimeError:
|
|
loop = asyncio.new_event_loop()
|
|
yield loop
|
|
loop.close()
|
|
|
|
@pytest.fixture(scope="function")
|
|
def setup_and_teardown(event_loop):
|
|
original_state = _copy_litellm_state()
|
|
_clear_logging_queue(event_loop)
|
|
_reset_litellm_callbacks()
|
|
asyncio.set_event_loop(event_loop)
|
|
yield
|
|
_clear_logging_queue(event_loop)
|
|
_reset_litellm_callbacks()
|
|
_restore_litellm_state(original_state)
|
|
pending = asyncio.all_tasks(event_loop)
|
|
for task in pending:
|
|
task.cancel()
|
|
if pending:
|
|
event_loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True))
|
|
|
|
def _copy_litellm_state():
|
|
state = {}
|
|
for attr in _CALLBACK_ATTRS:
|
|
if hasattr(litellm, attr):
|
|
value = getattr(litellm, attr)
|
|
state[attr] = value.copy() if isinstance(value, list) else value
|
|
for attr in _SCALAR_ATTRS:
|
|
if hasattr(litellm, attr):
|
|
state[attr] = getattr(litellm, attr)
|
|
return state
|
|
|
|
_CALLBACK_ATTRS = (
|
|
"callbacks",
|
|
"success_callback",
|
|
"failure_callback",
|
|
"_async_success_callback",
|
|
"_async_failure_callback",
|
|
)
|
|
|
|
_SCALAR_ATTRS = (
|
|
"num_retries",
|
|
"set_verbose",
|
|
"cache",
|
|
"allowed_fails",
|
|
"disable_aiohttp_transport",
|
|
"force_ipv4",
|
|
"drop_params",
|
|
"modify_params",
|
|
"api_base",
|
|
"api_key",
|
|
"cohere_key",
|
|
)
|
|
|
|
def _clear_logging_queue(loop=None) -> None:
|
|
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
|
|
|
if loop is not None and (not loop.is_closed()) and (not loop.is_running()):
|
|
loop.run_until_complete(GLOBAL_LOGGING_WORKER.clear_queue())
|
|
return
|
|
asyncio.run(GLOBAL_LOGGING_WORKER.clear_queue())
|
|
|
|
def _reset_litellm_callbacks() -> None:
|
|
for attr in _CALLBACK_ATTRS:
|
|
if hasattr(litellm, attr):
|
|
setattr(litellm, attr, [])
|
|
manager = getattr(litellm, "logging_callback_manager", None)
|
|
reset = getattr(manager, "_reset_all_callbacks", None)
|
|
if callable(reset):
|
|
reset()
|
|
|
|
def _restore_litellm_state(state) -> None:
|
|
for attr, value in state.items():
|
|
if hasattr(litellm, attr):
|
|
setattr(litellm, attr, value)
|
|
|
|
def _make_batch_output_line(prompt_tokens: int = 10, completion_tokens: int = 5):
|
|
"""Return a single successful batch output line (OpenAI JSONL format)."""
|
|
return {
|
|
"id": "batch_req_1",
|
|
"custom_id": "req-1",
|
|
"response": {
|
|
"status_code": 200,
|
|
"body": {
|
|
"id": "chatcmpl-test",
|
|
"object": "chat.completion",
|
|
"model": "fake-batch-model",
|
|
"usage": {
|
|
"prompt_tokens": prompt_tokens,
|
|
"completion_tokens": completion_tokens,
|
|
"total_tokens": prompt_tokens + completion_tokens,
|
|
},
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"message": {"role": "assistant", "content": "Hello"},
|
|
"finish_reason": "stop",
|
|
}
|
|
],
|
|
},
|
|
},
|
|
"error": None,
|
|
}
|
|
|
|
CUSTOM_MODEL_INFO = {
|
|
"input_cost_per_token_batches": 0.00125,
|
|
"output_cost_per_token_batches": 0.005,
|
|
}
|
|
|
|
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
|
|
def test_batch_cost_calculator_explicit_zero_pricing_not_overridden_by_global(
|
|
monkeypatch,
|
|
):
|
|
"""
|
|
Explicit ``0`` / ``0.0`` pricing must count as present so we do not fall back
|
|
to the global pricing table (truthiness would treat zero as missing).
|
|
"""
|
|
usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500)
|
|
|
|
def fake_get_model_info(*args, **kwargs):
|
|
return {
|
|
"input_cost_per_token_batches": 1e-3,
|
|
"output_cost_per_token_batches": 2e-3,
|
|
}
|
|
|
|
monkeypatch.setattr(litellm, "get_model_info", fake_get_model_info)
|
|
|
|
prompt_cost, completion_cost = batch_cost_calculator(
|
|
usage=usage,
|
|
model="any-model",
|
|
custom_llm_provider="openai",
|
|
model_info={
|
|
"input_cost_per_token_batches": 0.0,
|
|
"output_cost_per_token_batches": 0.0,
|
|
},
|
|
)
|
|
|
|
assert prompt_cost == 0.0
|
|
assert completion_cost == 0.0
|
|
|
|
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
|
|
def test_batch_cost_calculator_uses_custom_model_info():
|
|
"""batch_cost_calculator should use model_info override when provided."""
|
|
usage = Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15)
|
|
|
|
prompt_cost, completion_cost = batch_cost_calculator(
|
|
usage=usage,
|
|
model="fake-batch-model",
|
|
custom_llm_provider="openai",
|
|
model_info=CUSTOM_MODEL_INFO,
|
|
)
|
|
|
|
expected_prompt = 10 * 0.00125
|
|
expected_completion = 5 * 0.005
|
|
assert prompt_cost == pytest.approx(expected_prompt), f"Expected prompt cost {expected_prompt}, got {prompt_cost}"
|
|
assert completion_cost == pytest.approx(expected_completion), (
|
|
f"Expected completion cost {expected_completion}, got {completion_cost}"
|
|
)
|
|
|
|
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
|
|
def test_aggregate_batch_cost_uses_custom_model_info():
|
|
"""_aggregate_batch_cost_usage_models should thread model_info to batch_cost_calculator."""
|
|
file_content = [_make_batch_output_line(prompt_tokens=10, completion_tokens=5)]
|
|
|
|
result = _aggregate_batch_cost_usage_models(
|
|
entries=file_content,
|
|
custom_llm_provider="openai",
|
|
model_info=CUSTOM_MODEL_INFO,
|
|
)
|
|
|
|
expected = (10 * 0.00125) + (5 * 0.005)
|
|
assert result.cost == pytest.approx(expected), f"Expected total cost {expected}, got {result.cost}"
|
|
|
|
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
|
|
@pytest.mark.parametrize("data_residency", ["eu", "us"])
|
|
def test_batch_cost_calculator_applies_data_residency_uplift(data_residency, monkeypatch):
|
|
"""batch_cost_calculator should apply the regional uplift multiplier when
|
|
data_residency is set and the model carries a configured multiplier."""
|
|
monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True")
|
|
prev_model_cost = litellm.model_cost
|
|
litellm.model_cost = litellm.get_model_cost_map(url="")
|
|
try:
|
|
usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500)
|
|
|
|
base_prompt, base_completion = batch_cost_calculator(
|
|
usage=usage,
|
|
model="gpt-5.4",
|
|
custom_llm_provider="openai",
|
|
)
|
|
regional_prompt, regional_completion = batch_cost_calculator(
|
|
usage=usage,
|
|
model="gpt-5.4",
|
|
custom_llm_provider="openai",
|
|
data_residency=data_residency,
|
|
)
|
|
|
|
assert base_prompt > 0 and base_completion > 0
|
|
assert regional_prompt == pytest.approx(base_prompt * 1.10, rel=1e-9)
|
|
assert regional_completion == pytest.approx(base_completion * 1.10, rel=1e-9)
|
|
finally:
|
|
litellm.model_cost = prev_model_cost
|
|
|
|
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
|
|
@pytest.mark.asyncio
|
|
async def test_calculate_batch_cost_and_usage_uses_custom_model_info():
|
|
"""calculate_batch_cost_and_usage should thread model_info."""
|
|
file_content = [_make_batch_output_line(prompt_tokens=10, completion_tokens=5)]
|
|
|
|
result = await calculate_batch_cost_and_usage(
|
|
file_content_dictionary=file_content,
|
|
custom_llm_provider="openai",
|
|
model_info=CUSTOM_MODEL_INFO,
|
|
)
|
|
|
|
expected = (10 * 0.00125) + (5 * 0.005)
|
|
assert result.cost == pytest.approx(expected), f"Expected total cost {expected}, got {result.cost}"
|
|
assert result.usage.prompt_tokens == 10
|
|
assert result.usage.completion_tokens == 5
|
|
|
|
@pytest.fixture
|
|
def sample_file_content():
|
|
return b"""
|
|
{"id": "batch_req_6769ca596b38819093d7ae9f522de924", "custom_id": "request-1", "response": {"status_code": 200, "request_id": "07bc45ab4e7e26ac23a0c949973327e7", "body": {"id": "chatcmpl-AhjSMl7oZ79yIPHLRYgmgXSixTJr7", "object": "chat.completion", "created": 1734986202, "model": "gpt-4o-mini-2024-07-18", "choices": [{"index": 0, "message": {"role": "assistant", "content": "Hello! How can I assist you today?", "refusal": null}, "logprobs": null, "finish_reason": "stop"}], "usage": {"prompt_tokens": 20, "completion_tokens": 10, "total_tokens": 30, "prompt_tokens_details": {"cached_tokens": 0, "audio_tokens": 0}, "completion_tokens_details": {"reasoning_tokens": 0, "audio_tokens": 0, "accepted_prediction_tokens": 0, "rejected_prediction_tokens": 0}}, "system_fingerprint": "fp_0aa8d3e20b"}}, "error": null}
|
|
{"id": "batch_req_6769ca597e588190920666612634e2b4", "custom_id": "request-2", "response": {"status_code": 200, "request_id": "82e04f4c001fe2c127cbad199f5fd31b", "body": {"id": "chatcmpl-AhjSNgVB4Oa4Hq0NruTRsBaEbRWUP", "object": "chat.completion", "created": 1734986203, "model": "gpt-4o-mini-2024-07-18", "choices": [{"index": 0, "message": {"role": "assistant", "content": "Hello! What can I do for you today?", "refusal": null}, "logprobs": null, "finish_reason": "length"}], "usage": {"prompt_tokens": 22, "completion_tokens": 10, "total_tokens": 32, "prompt_tokens_details": {"cached_tokens": 0, "audio_tokens": 0}, "completion_tokens_details": {"reasoning_tokens": 0, "audio_tokens": 0, "accepted_prediction_tokens": 0, "rejected_prediction_tokens": 0}}, "system_fingerprint": "fp_0aa8d3e20b"}}, "error": null}
|
|
"""
|
|
|
|
@pytest.fixture
|
|
def sample_file_content_dict():
|
|
return [
|
|
{
|
|
"id": "batch_req_6769ca596b38819093d7ae9f522de924",
|
|
"custom_id": "request-1",
|
|
"response": {
|
|
"status_code": 200,
|
|
"request_id": "07bc45ab4e7e26ac23a0c949973327e7",
|
|
"body": {
|
|
"id": "chatcmpl-AhjSMl7oZ79yIPHLRYgmgXSixTJr7",
|
|
"object": "chat.completion",
|
|
"created": 1734986202,
|
|
"model": "gpt-4o-mini-2024-07-18",
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": "Hello! How can I assist you today?",
|
|
"refusal": None,
|
|
},
|
|
"logprobs": None,
|
|
"finish_reason": "stop",
|
|
}
|
|
],
|
|
"usage": {
|
|
"prompt_tokens": 20,
|
|
"completion_tokens": 10,
|
|
"total_tokens": 30,
|
|
"prompt_tokens_details": {
|
|
"cached_tokens": 0,
|
|
"audio_tokens": 0,
|
|
},
|
|
"completion_tokens_details": {
|
|
"reasoning_tokens": 0,
|
|
"audio_tokens": 0,
|
|
"accepted_prediction_tokens": 0,
|
|
"rejected_prediction_tokens": 0,
|
|
},
|
|
},
|
|
"system_fingerprint": "fp_0aa8d3e20b",
|
|
},
|
|
},
|
|
"error": None,
|
|
},
|
|
{
|
|
"id": "batch_req_6769ca597e588190920666612634e2b4",
|
|
"custom_id": "request-2",
|
|
"response": {
|
|
"status_code": 200,
|
|
"request_id": "82e04f4c001fe2c127cbad199f5fd31b",
|
|
"body": {
|
|
"id": "chatcmpl-AhjSNgVB4Oa4Hq0NruTRsBaEbRWUP",
|
|
"object": "chat.completion",
|
|
"created": 1734986203,
|
|
"model": "gpt-4o-mini-2024-07-18",
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": "Hello! What can I do for you today?",
|
|
"refusal": None,
|
|
},
|
|
"logprobs": None,
|
|
"finish_reason": "length",
|
|
}
|
|
],
|
|
"usage": {
|
|
"prompt_tokens": 22,
|
|
"completion_tokens": 10,
|
|
"total_tokens": 32,
|
|
"prompt_tokens_details": {
|
|
"cached_tokens": 0,
|
|
"audio_tokens": 0,
|
|
},
|
|
"completion_tokens_details": {
|
|
"reasoning_tokens": 0,
|
|
"audio_tokens": 0,
|
|
"accepted_prediction_tokens": 0,
|
|
"rejected_prediction_tokens": 0,
|
|
},
|
|
},
|
|
"system_fingerprint": "fp_0aa8d3e20b",
|
|
},
|
|
},
|
|
"error": None,
|
|
},
|
|
]
|
|
|
|
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
|
|
def test_get_file_content_as_dictionary(sample_file_content):
|
|
result = get_file_content_as_dictionary(sample_file_content)
|
|
assert len(result) == 2
|
|
assert result[0]["id"] == "batch_req_6769ca596b38819093d7ae9f522de924"
|
|
assert result[0]["custom_id"] == "request-1"
|
|
assert result[0]["response"]["status_code"] == 200
|
|
assert result[0]["response"]["body"]["usage"]["total_tokens"] == 30
|
|
|
|
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
|
|
def test_get_batch_job_total_usage_from_file_content(sample_file_content_dict):
|
|
with patch("litellm.completion_cost", return_value=0.0):
|
|
result = _aggregate_batch_cost_usage_models(entries=sample_file_content_dict, custom_llm_provider="openai")
|
|
assert result.usage.total_tokens == 62 # 30 + 32
|
|
assert result.usage.prompt_tokens == 42 # 20 + 22
|
|
assert result.usage.completion_tokens == 20 # 10 + 10
|
|
|
|
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
|
|
@pytest.mark.asyncio
|
|
async def test_batch_cost_calculator(sample_file_content_dict):
|
|
"""
|
|
mock batch_cost_calculator to return (0.3, 0.2) per line
|
|
|
|
we know sample_file_content_dict has 2 successful responses
|
|
|
|
so we expect the cost to be (0.3 + 0.2) * 2 = 1.0, split 0.6 / 0.4
|
|
"""
|
|
with patch("litellm.cost_calculator.batch_cost_calculator", return_value=(0.3, 0.2)):
|
|
result = _aggregate_batch_cost_usage_models(
|
|
entries=sample_file_content_dict,
|
|
custom_llm_provider="openai",
|
|
)
|
|
assert result.cost == pytest.approx(1.0) # (0.3 + 0.2) * 2 successful responses
|
|
assert result.prompt_cost == pytest.approx(0.6)
|
|
assert result.completion_cost == pytest.approx(0.4)
|
|
|
|
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
|
|
def test_get_response_from_batch_job_output_file(sample_file_content_dict):
|
|
result = _get_response_from_batch_job_output_file(sample_file_content_dict[0])
|
|
assert result["id"] == "chatcmpl-AhjSMl7oZ79yIPHLRYgmgXSixTJr7"
|
|
assert result["object"] == "chat.completion"
|
|
assert result["usage"]["total_tokens"] == 30
|
|
|
|
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
|
|
@pytest.mark.asyncio
|
|
async def test_batch_retrieve_cost_tracking_with_completed_batch_no_explicit_cost():
|
|
"""
|
|
Test that cost is calculated for completed batches when no explicit cost data is provided.
|
|
|
|
Regression test for: When batch status is "completed" and explicit batch_cost/batch_usage/batch_models
|
|
are not provided, the system should compute batch data by calling _handle_completed_batch.
|
|
"""
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
from litellm.litellm_core_utils.litellm_logging import Logging
|
|
from litellm.types.utils import CallTypes, LiteLLMBatch
|
|
|
|
# Mock batch result with completed status
|
|
mock_batch = LiteLLMBatch(
|
|
id="batch-test-123",
|
|
object="batch",
|
|
endpoint="/v1/chat/completions",
|
|
errors=None,
|
|
input_file_id="file-input-123",
|
|
completion_window="24h",
|
|
status="completed",
|
|
output_file_id="file-output-123",
|
|
error_file_id=None,
|
|
created_at=1234567890,
|
|
in_progress_at=1234567900,
|
|
expires_at=1234654290,
|
|
finalizing_at=1234568000,
|
|
completed_at=1234568100,
|
|
failed_at=None,
|
|
expired_at=None,
|
|
cancelling_at=None,
|
|
cancelled_at=None,
|
|
request_counts={
|
|
"total": 10,
|
|
"completed": 10,
|
|
"failed": 0,
|
|
},
|
|
metadata=None,
|
|
)
|
|
mock_batch._hidden_params = {}
|
|
|
|
# Create logging object
|
|
logging_obj = Logging(
|
|
model="gpt-5-mini",
|
|
messages=[{"role": "user", "content": "test"}],
|
|
stream=False,
|
|
call_type=CallTypes.aretrieve_batch.value,
|
|
litellm_call_id="test-call-123",
|
|
function_id="test-function",
|
|
start_time=time.time(),
|
|
dynamic_success_callbacks=[],
|
|
)
|
|
logging_obj.custom_llm_provider = "openai"
|
|
|
|
# Mock handle_completed_batch to return cost data
|
|
from litellm.batches.batch_utils import BatchCostUsageResult
|
|
|
|
expected_cost = 0.05
|
|
expected_usage = litellm.Usage(
|
|
prompt_tokens=100,
|
|
completion_tokens=50,
|
|
total_tokens=150,
|
|
)
|
|
expected_models = ["gpt-5-mini"]
|
|
|
|
with patch(
|
|
"litellm.litellm_core_utils.litellm_logging.handle_completed_batch",
|
|
new=AsyncMock(
|
|
return_value=BatchCostUsageResult(
|
|
cost=expected_cost,
|
|
usage=expected_usage,
|
|
models=expected_models,
|
|
successful_requests=10,
|
|
failed_requests=0,
|
|
)
|
|
),
|
|
) as mock_handle_batch:
|
|
# Call async_success_handler
|
|
await logging_obj.async_success_handler(
|
|
result=mock_batch,
|
|
start_time=time.time(),
|
|
end_time=time.time() + 1,
|
|
)
|
|
|
|
# Verify handle_completed_batch was called
|
|
mock_handle_batch.assert_called_once()
|
|
|
|
# Verify cost and usage were set on the batch result
|
|
assert mock_batch._hidden_params["response_cost"] == expected_cost
|
|
assert mock_batch._hidden_params["batch_models"] == expected_models
|
|
assert mock_batch._hidden_params["batch_successful_requests"] == 10
|
|
assert mock_batch._hidden_params["batch_failed_requests"] == 0
|
|
assert mock_batch.usage == expected_usage
|
|
|
|
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
|
|
@pytest.mark.asyncio
|
|
async def test_handle_completed_batch_computes_real_cost_from_output_file(
|
|
sample_file_content_dict,
|
|
):
|
|
"""Integration: a completed batch's cost and usage are computed from its output
|
|
file via the real cost-calc chain (only the file download is stubbed). This is
|
|
the function the retrieve handler invokes on completion; a dropped output line, a
|
|
wrong token sum, or mispriced model fails this test.
|
|
"""
|
|
from litellm.batches.batch_utils import handle_completed_batch
|
|
from litellm.types.utils import LiteLLMBatch
|
|
|
|
batch = LiteLLMBatch(
|
|
id="batch-real-cost-123",
|
|
object="batch",
|
|
endpoint="/v1/chat/completions",
|
|
input_file_id="file-input-123",
|
|
completion_window="24h",
|
|
status="completed",
|
|
output_file_id="file-output-123",
|
|
created_at=1234567890,
|
|
)
|
|
|
|
sample_file_content_bytes = "\n".join(json.dumps(row) for row in sample_file_content_dict).encode()
|
|
with patch(
|
|
"litellm.batches.batch_utils._fetch_batch_output_file_content",
|
|
new=AsyncMock(return_value=sample_file_content_bytes),
|
|
):
|
|
result = await handle_completed_batch(batch=batch, custom_llm_provider="openai")
|
|
|
|
pricing = litellm.model_cost["gpt-4o-mini-2024-07-18"]
|
|
expected_cost = 42 * pricing["input_cost_per_token_batches"] + 20 * pricing["output_cost_per_token_batches"]
|
|
|
|
assert result.cost == pytest.approx(expected_cost)
|
|
assert result.cost > 0
|
|
assert result.cost < 42 * pricing["input_cost_per_token"] + 20 * pricing["output_cost_per_token"]
|
|
assert result.usage.prompt_tokens == 42
|
|
assert result.usage.completion_tokens == 20
|
|
assert result.usage.total_tokens == 62
|
|
assert result.models == ["gpt-4o-mini-2024-07-18", "gpt-4o-mini-2024-07-18"]
|
|
assert result.successful_requests == 2
|
|
assert result.failed_requests == 0
|
|
|
|
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
|
|
@pytest.mark.asyncio
|
|
async def test_batch_retrieve_cost_tracking_with_explicit_cost_data():
|
|
"""
|
|
Test that explicit cost data is used when provided, skipping computation.
|
|
|
|
Regression test for: When batch_cost, batch_usage, and batch_models are explicitly
|
|
provided in kwargs, they should be used directly without calling _handle_completed_batch.
|
|
"""
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
from litellm.litellm_core_utils.litellm_logging import Logging
|
|
from litellm.types.utils import CallTypes, LiteLLMBatch
|
|
|
|
# Mock batch result with completed status
|
|
mock_batch = LiteLLMBatch(
|
|
id="batch-test-456",
|
|
object="batch",
|
|
endpoint="/v1/chat/completions",
|
|
errors=None,
|
|
input_file_id="file-input-456",
|
|
completion_window="24h",
|
|
status="completed",
|
|
output_file_id="file-output-456",
|
|
error_file_id=None,
|
|
created_at=1234567890,
|
|
in_progress_at=1234567900,
|
|
expires_at=1234654290,
|
|
finalizing_at=1234568000,
|
|
completed_at=1234568100,
|
|
failed_at=None,
|
|
expired_at=None,
|
|
cancelling_at=None,
|
|
cancelled_at=None,
|
|
request_counts={
|
|
"total": 5,
|
|
"completed": 5,
|
|
"failed": 0,
|
|
},
|
|
metadata=None,
|
|
)
|
|
mock_batch._hidden_params = {}
|
|
|
|
# Create logging object
|
|
logging_obj = Logging(
|
|
model="gpt-5-mini",
|
|
messages=[{"role": "user", "content": "test"}],
|
|
stream=False,
|
|
call_type=CallTypes.aretrieve_batch.value,
|
|
litellm_call_id="test-call-456",
|
|
function_id="test-function",
|
|
start_time=time.time(),
|
|
dynamic_success_callbacks=[],
|
|
)
|
|
logging_obj.custom_llm_provider = "openai"
|
|
|
|
# Explicit cost data to pass in kwargs
|
|
explicit_cost = 0.10
|
|
explicit_usage = litellm.Usage(
|
|
prompt_tokens=200,
|
|
completion_tokens=100,
|
|
total_tokens=300,
|
|
)
|
|
explicit_models = ["gpt-5-mini", "gpt-5.5"]
|
|
|
|
with patch(
|
|
"litellm.litellm_core_utils.litellm_logging.handle_completed_batch",
|
|
new=AsyncMock(),
|
|
) as mock_handle_batch:
|
|
# Call async_success_handler with explicit cost data
|
|
await logging_obj.async_success_handler(
|
|
result=mock_batch,
|
|
start_time=time.time(),
|
|
end_time=time.time() + 1,
|
|
batch_cost=explicit_cost,
|
|
batch_usage=explicit_usage,
|
|
batch_models=explicit_models,
|
|
)
|
|
|
|
# Verify handle_completed_batch was NOT called (since explicit data provided)
|
|
mock_handle_batch.assert_not_called()
|
|
|
|
# Verify explicit cost data was used
|
|
assert mock_batch._hidden_params["response_cost"] == explicit_cost
|
|
assert mock_batch._hidden_params["batch_models"] == explicit_models
|
|
assert mock_batch.usage == explicit_usage
|
|
|
|
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
|
|
@pytest.mark.asyncio
|
|
async def test_batch_retrieve_explicit_cost_split_sets_cost_breakdown():
|
|
"""The poller passes the batch's prompt/completion cost split so the spend row's
|
|
cost_breakdown carries real input/output costs; without it the UI's Cost Breakdown
|
|
card renders blank for every batch. Regression for the split being dropped."""
|
|
from litellm.litellm_core_utils.litellm_logging import Logging
|
|
from litellm.types.utils import CallTypes, LiteLLMBatch
|
|
|
|
mock_batch = LiteLLMBatch(
|
|
id="batch-breakdown-1",
|
|
object="batch",
|
|
endpoint="/v1/chat/completions",
|
|
errors=None,
|
|
input_file_id="file-input-1",
|
|
completion_window="24h",
|
|
status="completed",
|
|
output_file_id="file-output-1",
|
|
created_at=1234567890,
|
|
)
|
|
mock_batch._hidden_params = {}
|
|
|
|
logging_obj = Logging(
|
|
model="gpt-5-mini",
|
|
messages=[{"role": "user", "content": "test"}],
|
|
stream=False,
|
|
call_type=CallTypes.aretrieve_batch.value,
|
|
litellm_call_id="test-call-breakdown",
|
|
function_id="test-function",
|
|
start_time=time.time(),
|
|
dynamic_success_callbacks=[],
|
|
)
|
|
logging_obj.custom_llm_provider = "openai"
|
|
|
|
await logging_obj.async_success_handler(
|
|
result=mock_batch,
|
|
start_time=time.time(),
|
|
end_time=time.time() + 1,
|
|
batch_cost=0.10,
|
|
batch_usage=litellm.Usage(prompt_tokens=200, completion_tokens=100, total_tokens=300),
|
|
batch_models=["gpt-5-mini"],
|
|
batch_prompt_cost=0.06,
|
|
batch_completion_cost=0.04,
|
|
)
|
|
|
|
assert logging_obj.cost_breakdown is not None
|
|
assert logging_obj.cost_breakdown["input_cost"] == 0.06
|
|
assert logging_obj.cost_breakdown["output_cost"] == 0.04
|
|
assert logging_obj.cost_breakdown["total_cost"] == 0.10
|
|
|
|
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
|
|
@pytest.mark.asyncio
|
|
async def test_batch_retrieve_cost_tracking_with_unified_file_id_incomplete_batch():
|
|
"""
|
|
Test that cost computation is skipped for unified file IDs with non-completed batches.
|
|
|
|
Regression test for: For unified file IDs (base64 encoded), cost should only be computed
|
|
when batch status is "completed" and explicit data is not provided.
|
|
"""
|
|
import base64
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
from litellm.litellm_core_utils.litellm_logging import Logging
|
|
from litellm.types.utils import CallTypes, LiteLLMBatch, SpecialEnums
|
|
|
|
# Create a proper unified file ID by encoding the correct prefix
|
|
unified_id_str = f"{SpecialEnums.LITELM_MANAGED_FILE_ID_PREFIX.value}:test_file_789;unified_id:batch-789"
|
|
encoded_unified_id = base64.urlsafe_b64encode(unified_id_str.encode()).decode().rstrip("=")
|
|
|
|
# Mock batch result with in_progress status and unified file ID
|
|
mock_batch = LiteLLMBatch(
|
|
id=encoded_unified_id, # Properly encoded unified ID
|
|
object="batch",
|
|
endpoint="/v1/chat/completions",
|
|
errors=None,
|
|
input_file_id="file-input-789",
|
|
completion_window="24h",
|
|
status="in_progress", # Not completed
|
|
output_file_id=None,
|
|
error_file_id=None,
|
|
created_at=1234567890,
|
|
in_progress_at=1234567900,
|
|
expires_at=1234654290,
|
|
finalizing_at=None,
|
|
completed_at=None,
|
|
failed_at=None,
|
|
expired_at=None,
|
|
cancelling_at=None,
|
|
cancelled_at=None,
|
|
request_counts={
|
|
"total": 10,
|
|
"completed": 3,
|
|
"failed": 0,
|
|
},
|
|
metadata=None,
|
|
)
|
|
mock_batch._hidden_params = {}
|
|
|
|
# Create logging object
|
|
logging_obj = Logging(
|
|
model="gpt-5-mini",
|
|
messages=[{"role": "user", "content": "test"}],
|
|
stream=False,
|
|
call_type=CallTypes.aretrieve_batch.value,
|
|
litellm_call_id="test-call-789",
|
|
function_id="test-function",
|
|
start_time=time.time(),
|
|
dynamic_success_callbacks=[],
|
|
)
|
|
logging_obj.custom_llm_provider = "openai"
|
|
|
|
with patch(
|
|
"litellm.litellm_core_utils.litellm_logging.handle_completed_batch",
|
|
new=AsyncMock(),
|
|
) as mock_handle_batch:
|
|
# Call async_success_handler with in_progress batch (unified file ID)
|
|
await logging_obj.async_success_handler(
|
|
result=mock_batch,
|
|
start_time=time.time(),
|
|
end_time=time.time() + 1,
|
|
)
|
|
|
|
# Verify handle_completed_batch was NOT called (batch not completed and is unified file ID)
|
|
mock_handle_batch.assert_not_called()
|
|
|
|
# Verify cost data was not set
|
|
assert "response_cost" not in mock_batch._hidden_params
|
|
assert "batch_models" not in mock_batch._hidden_params
|
|
assert not hasattr(mock_batch, "usage") or mock_batch.usage is None
|
|
|
|
@pytest.mark.usefixtures("_vcr_outcome_gate", "setup_and_teardown")
|
|
@pytest.mark.asyncio
|
|
async def test_batch_retrieve_cost_tracking_with_partial_explicit_data():
|
|
"""
|
|
Test that cost is computed when only partial explicit data is provided.
|
|
|
|
Regression test for: If batch_cost, batch_usage, or batch_models is missing
|
|
(not all three provided), and batch is completed, system should compute the data.
|
|
"""
|
|
from unittest.mock import AsyncMock, patch
|
|
|
|
from litellm.litellm_core_utils.litellm_logging import Logging
|
|
from litellm.types.utils import CallTypes, LiteLLMBatch
|
|
|
|
# Mock batch result with completed status
|
|
mock_batch = LiteLLMBatch(
|
|
id="batch-test-partial",
|
|
object="batch",
|
|
endpoint="/v1/chat/completions",
|
|
errors=None,
|
|
input_file_id="file-input-partial",
|
|
completion_window="24h",
|
|
status="completed",
|
|
output_file_id="file-output-partial",
|
|
error_file_id=None,
|
|
created_at=1234567890,
|
|
in_progress_at=1234567900,
|
|
expires_at=1234654290,
|
|
finalizing_at=1234568000,
|
|
completed_at=1234568100,
|
|
failed_at=None,
|
|
expired_at=None,
|
|
cancelling_at=None,
|
|
cancelled_at=None,
|
|
request_counts={
|
|
"total": 8,
|
|
"completed": 8,
|
|
"failed": 0,
|
|
},
|
|
metadata=None,
|
|
)
|
|
mock_batch._hidden_params = {}
|
|
|
|
# Create logging object
|
|
logging_obj = Logging(
|
|
model="gpt-5-mini",
|
|
messages=[{"role": "user", "content": "test"}],
|
|
stream=False,
|
|
call_type=CallTypes.aretrieve_batch.value,
|
|
litellm_call_id="test-call-partial",
|
|
function_id="test-function",
|
|
start_time=time.time(),
|
|
dynamic_success_callbacks=[],
|
|
)
|
|
|
|
logging_obj.custom_llm_provider = "openai"
|
|
|
|
# Only provide batch_cost, missing batch_usage and batch_models
|
|
partial_cost = 0.08
|
|
|
|
expected_cost = 0.06
|
|
expected_usage = litellm.Usage(
|
|
prompt_tokens=150,
|
|
completion_tokens=75,
|
|
total_tokens=225,
|
|
)
|
|
expected_models = ["gpt-5-mini"]
|
|
|
|
from litellm.batches.batch_utils import BatchCostUsageResult
|
|
|
|
with patch(
|
|
"litellm.litellm_core_utils.litellm_logging.handle_completed_batch",
|
|
new=AsyncMock(
|
|
return_value=BatchCostUsageResult(
|
|
cost=expected_cost,
|
|
usage=expected_usage,
|
|
models=expected_models,
|
|
successful_requests=8,
|
|
failed_requests=0,
|
|
)
|
|
),
|
|
) as mock_handle_batch:
|
|
# Call async_success_handler with partial explicit data
|
|
await logging_obj.async_success_handler(
|
|
result=mock_batch,
|
|
start_time=time.time(),
|
|
end_time=time.time() + 1,
|
|
batch_cost=partial_cost, # Only cost provided, not usage or models
|
|
)
|
|
|
|
# Verify handle_completed_batch WAS called (since not all data provided)
|
|
mock_handle_batch.assert_called_once()
|
|
|
|
# Verify computed cost data was used (not partial explicit data)
|
|
assert mock_batch._hidden_params["response_cost"] == expected_cost
|
|
assert mock_batch._hidden_params["batch_models"] == expected_models
|
|
assert mock_batch._hidden_params["batch_successful_requests"] == 8
|
|
assert mock_batch._hidden_params["batch_failed_requests"] == 0
|
|
assert mock_batch.usage == expected_usage
|