mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
b067e836f8
The cost callback used to look for an existing `<batch id>_batch_cost` row before charging a completed batch, which left a window where concurrent retrieves on any instance all charged the key, and it would honor a row any request had written under that id. The spend update writer now inserts the batch cost row itself with `create_many(skip_duplicates=True)` and only the retrieve whose insert lands charges the key, team, and user. An existing row only takes the charge when it is a successful `aretrieve_batch` row, so a client-chosen `x-litellm-call-id` on another endpoint cannot suppress billing. Batch cost rows no longer get their own immediate flush path `batch_cost_is_final` now treats the proxy's normalized `complete` status like `completed`, which the enterprise batch cost poller relies on when it decides whether a completed batch is safe to retire. Tests build that status with `model_copy` since the OpenAI `Batch` model rejects it The `test-quality-ok` markers sit on the `patch(` lines the gate keys on, and the logging tests no longer wrap the priced retrieve in `contextlib.suppress`
1775 lines
66 KiB
Python
1775 lines
66 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 json
|
|
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, Usage
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# Builders for batch OUTPUT file rows.
|
|
# Shape: {"response": {"status_code": 200, "body": {... "usage": {...}}}}
|
|
# --------------------------------------------------------------------------- #
|
|
|
|
|
|
def _usage(p, c, t=None):
|
|
return {
|
|
"prompt_tokens": p,
|
|
"completion_tokens": c,
|
|
"total_tokens": t if t is not None else p + c,
|
|
}
|
|
|
|
|
|
def _success_row(model="gpt-4o", usage=None, **body_extra):
|
|
body = {"model": model, **body_extra}
|
|
if usage is not None:
|
|
body["usage"] = usage
|
|
return {"response": {"status_code": 200, "body": body}}
|
|
|
|
|
|
def _failed_row(status_code=500, model="gpt-4o"):
|
|
return {"response": {"status_code": status_code, "body": {"model": model}}}
|
|
|
|
|
|
# =========================================================================== #
|
|
# _batch_response_was_successful
|
|
# =========================================================================== #
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"row,expected",
|
|
[
|
|
({"response": {"status_code": 200}}, True),
|
|
({"response": {"status_code": 500}}, False),
|
|
({"response": {"status_code": 429}}, False),
|
|
({"response": {}}, False), # no status_code
|
|
({}, False), # no response
|
|
({"response": None}, False), # null response
|
|
],
|
|
)
|
|
def test_batch_response_was_successful(row, expected):
|
|
assert bu._batch_response_was_successful(row) is expected
|
|
|
|
|
|
# =========================================================================== #
|
|
# _get_response_from_batch_job_output_file
|
|
# =========================================================================== #
|
|
|
|
|
|
def test_get_response_body_present():
|
|
row = {"response": {"body": {"model": "gpt-4o", "usage": {"x": 1}}}}
|
|
assert bu._get_response_from_batch_job_output_file(row) == {
|
|
"model": "gpt-4o",
|
|
"usage": {"x": 1},
|
|
}
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"row",
|
|
[
|
|
{}, # no response
|
|
{"response": {}}, # no body
|
|
{"response": None}, # null response
|
|
{"response": {"body": None}}, # null body
|
|
],
|
|
)
|
|
def test_get_response_body_missing_returns_empty(row):
|
|
assert bu._get_response_from_batch_job_output_file(row) == {}
|
|
|
|
|
|
# =========================================================================== #
|
|
# _get_batch_job_usage_from_response_body
|
|
# =========================================================================== #
|
|
|
|
|
|
def test_get_usage_from_response_body():
|
|
usage = bu._get_batch_job_usage_from_response_body(
|
|
{"usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}}
|
|
)
|
|
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (
|
|
10,
|
|
5,
|
|
15,
|
|
)
|
|
|
|
|
|
def test_get_usage_from_response_body_missing_is_zero():
|
|
usage = bu._get_batch_job_usage_from_response_body({})
|
|
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (
|
|
0,
|
|
0,
|
|
0,
|
|
)
|
|
|
|
|
|
# =========================================================================== #
|
|
# _get_file_content_as_dictionary (JSONL parsing)
|
|
# =========================================================================== #
|
|
|
|
|
|
def test_parse_jsonl_multiple_lines():
|
|
content = b'{"a": 1}\n{"b": 2}\n{"c": 3}'
|
|
assert bu._get_file_content_as_dictionary(content) == [
|
|
{"a": 1},
|
|
{"b": 2},
|
|
{"c": 3},
|
|
]
|
|
|
|
|
|
def test_parse_jsonl_trailing_newline_skipped():
|
|
# outer content is stripped; the trailing-newline empty line is dropped.
|
|
content = b'{"a": 1}\n{"b": 2}\n'
|
|
assert bu._get_file_content_as_dictionary(content) == [{"a": 1}, {"b": 2}]
|
|
|
|
|
|
def test_parse_jsonl_empty_content_is_empty_list():
|
|
assert bu._get_file_content_as_dictionary(b"") == []
|
|
|
|
|
|
def test_parse_jsonl_malformed_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",
|
|
"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_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_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):
|
|
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 1.0)
|
|
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,
|
|
)
|
|
|
|
|
|
# =========================================================================== #
|
|
# _aggregate_batch_cost_usage_models: cost (cost maps mocked)
|
|
# =========================================================================== #
|
|
|
|
|
|
def test_cost_from_content_completion_cost_path(monkeypatch):
|
|
# model_info is None -> litellm.completion_cost per successful row.
|
|
calls = []
|
|
|
|
def _completion_cost(**kw):
|
|
calls.append(kw)
|
|
return 0.5
|
|
|
|
monkeypatch.setattr(litellm, "completion_cost", _completion_cost)
|
|
rows = [
|
|
_success_row(usage=_usage(10, 5)),
|
|
_failed_row(), # excluded -> not costed
|
|
_success_row(usage=_usage(20, 10)),
|
|
]
|
|
|
|
result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
|
|
|
|
assert result.cost == 1.0 # 2 successful * 0.5
|
|
assert len(calls) == 2 # failed row not costed
|
|
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."""
|
|
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.5)
|
|
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: 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: 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_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):
|
|
rows = [_success_row(model="gpt-4o", usage=_usage(10, 5))]
|
|
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 2.5)
|
|
|
|
result = await bu.calculate_batch_cost_and_usage(file_content_dictionary=rows, custom_llm_provider="openai")
|
|
|
|
assert result.cost == 2.5
|
|
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)
|
|
|
|
monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch)
|
|
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 3.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):
|
|
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 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)
|
|
# 3e-06 / 1.5e-05 on-demand, halved for batch.
|
|
assert result.cost == pytest.approx(1800 * 3e-06 / 2 + 1000 * 1.5e-05 / 2)
|
|
|
|
# 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_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
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# 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
|