scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-07 02:59:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

test(cost): type the batch_cost_calculator model_info literals instead of suppressing

Browse source
This commit is contained in:
mateo-berri 2026-08-17 15:40:29 -07:00
parent 5a11fe141e
commit e736b59802
1 changed files with 31 additions and 17 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

48
tests/test_litellm/test_cost_calculator.py
View file

@ -20,7 +20,7 @@ from litellm.cost_calculator import (
response_cost_calculator,
)
from litellm.types.llms.openai import OpenAIRealtimeStreamList
from litellm.types.utils import ModelResponse, PromptTokensDetailsWrapper, Usage
from litellm.types.utils import ModelInfo, ModelResponse, PromptTokensDetailsWrapper, Usage
from litellm.utils import TranscriptionResponse
@ -3562,16 +3562,18 @@ def test_batch_cost_calculator_prices_cache_creation_tokens_at_cache_write_rate(
"""
from litellm.cost_calculator import batch_cost_calculator
model_info: ModelInfo = {
"supported_openai_params": [],
"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,
}
prompt_cost, completion_cost_value = batch_cost_calculator(
usage=_batch_cache_usage(),
model="claude-sonnet-4-5-20250929",
custom_llm_provider="anthropic",
model_info={ # type: ignore[arg-type]
"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,
},
model_info=model_info,
)
assert prompt_cost == pytest.approx((1000 * 3e-6 + 8000 * 3e-7 + 2000 * 3.75e-6) / 2)
@ -3581,15 +3583,17 @@ def test_batch_cost_calculator_prices_cache_creation_tokens_at_cache_write_rate(
def test_batch_cost_calculator_cache_creation_falls_back_to_input_rate():
from litellm.cost_calculator import batch_cost_calculator
model_info: ModelInfo = {
"supported_openai_params": [],
"input_cost_per_token": 3e-6,
"output_cost_per_token": 15e-6,
"cache_read_input_token_cost": 3e-7,
}
prompt_cost, _ = batch_cost_calculator(
usage=_batch_cache_usage(),
model="claude-sonnet-4-5-20250929",
custom_llm_provider="anthropic",
model_info={ # type: ignore[arg-type]
"input_cost_per_token": 3e-6,
"output_cost_per_token": 15e-6,
"cache_read_input_token_cost": 3e-7,
},
model_info=model_info,
)
assert prompt_cost == pytest.approx((1000 * 3e-6 + 8000 * 3e-7 + 2000 * 3e-6) / 2)
@ -3616,16 +3620,26 @@ def test_batch_cost_calculator_honors_an_explicitly_zero_batch_rate(
"""
from litellm.cost_calculator import batch_cost_calculator
model_info: dict[str, float] = {"input_cost_per_token": 3e-6, "output_cost_per_token": 15e-6}
if batch_rate is not None:
model_info["input_cost_per_token_batches"] = batch_rate
model_info["output_cost_per_token_batches"] = batch_rate
base_model_info: ModelInfo = {
"supported_openai_params": [],
"input_cost_per_token": 3e-6,
"output_cost_per_token": 15e-6,
}
model_info: ModelInfo = (
base_model_info
if batch_rate is None
else {
**base_model_info,
"input_cost_per_token_batches": batch_rate,
"output_cost_per_token_batches": batch_rate,
}
)
prompt_cost, completion_cost_value = batch_cost_calculator(
usage=Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500),
model="claude-sonnet-4-5-20250929",
custom_llm_provider="anthropic",
model_info=model_info, # type: ignore[arg-type]
model_info=model_info,
)
assert prompt_cost == pytest.approx(expected_prompt)