ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: thread deployment model_info through batch cost calculation

Browse Source 
batch_cost_calculator only checked the global cost map, ignoring
deployment-level custom pricing (input_cost_per_token_batches etc.).
Add optional model_info param through the batch cost chain and pass
it from CheckBatchCost.
This commit is contained in:
Ephrim Stanley 2026-02-15 14:53:30 -05:00
parent a5626768a3
commit 7d794b567c
4 changed files with 180 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py
View File

@ -142,11 +142,15 @@ class CheckBatchCost:
custom_llm_provider=custom_llm_provider,
)
# Pass deployment model_info so custom batch pricing
# (input_cost_per_token_batches etc.) is used for cost calc
deployment_model_info = deployment_info.model_info.model_dump() if deployment_info.model_info else {}
batch_cost, batch_usage, batch_models = (
await calculate_batch_cost_and_usage(
file_content_dictionary=file_content_as_dict,
custom_llm_provider=llm_provider, # type: ignore
model_name=model_name,
model_info=deployment_model_info,
)
)
logging_obj = LiteLLMLogging(

36
litellm/batches/batch_utils.py
View File

@ -16,14 +16,22 @@ async def calculate_batch_cost_and_usage(
file_content_dictionary: List[dict],
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"],
model_name: Optional[str] = None,
model_info: Optional[dict] = None,
) -> Tuple[float, Usage, List[str]]:
"""
Calculate the cost and usage of a batch
Calculate the cost and usage of a batch.
Args:
model_info: Optional deployment-level model info with custom batch
pricing. Threaded through to batch_cost_calculator so that
deployment-specific pricing (e.g. input_cost_per_token_batches)
is used instead of the global cost map.
"""
batch_cost = _batch_cost_calculator(
custom_llm_provider=custom_llm_provider,
file_content_dictionary=file_content_dictionary,
model_name=model_name,
model_info=model_info,
)
batch_usage = _get_batch_job_total_usage_from_file_content(
file_content_dictionary=file_content_dictionary,
@ -94,6 +102,7 @@ def _batch_cost_calculator(
file_content_dictionary: List[dict],
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai",
model_name: Optional[str] = None,
model_info: Optional[dict] = None,
) -> float:
"""
Calculate the cost of a batch based on the output file id
@ -108,6 +117,7 @@ def _batch_cost_calculator(
total_cost = _get_batch_job_cost_from_file_content(
file_content_dictionary=file_content_dictionary,
custom_llm_provider=custom_llm_provider,
model_info=model_info,
)
verbose_logger.debug("total_cost=%s", total_cost)
return total_cost
@ -290,10 +300,13 @@ def _get_file_content_as_dictionary(file_content: bytes) -> List[dict]:
def _get_batch_job_cost_from_file_content(
file_content_dictionary: List[dict],
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai",
model_info: Optional[dict] = None,
) -> float:
"""
Get the cost of a batch job from the file content
"""
from litellm.cost_calculator import batch_cost_calculator
try:
total_cost: float = 0.0
# parse the file content as json
@ -303,11 +316,22 @@ def _get_batch_job_cost_from_file_content(
for _item in file_content_dictionary:
if _batch_response_was_successful(_item):
_response_body = _get_response_from_batch_job_output_file(_item)
total_cost += litellm.completion_cost(
completion_response=_response_body,
custom_llm_provider=custom_llm_provider,
call_type=CallTypes.aretrieve_batch.value,
)
if model_info is not None:
usage = _get_batch_job_usage_from_response_body(_response_body)
model = _response_body.get("model", "")
prompt_cost, completion_cost = batch_cost_calculator(
usage=usage,
model=model,
custom_llm_provider=custom_llm_provider,
model_info=model_info,
)
total_cost += prompt_cost + completion_cost
else:
total_cost += litellm.completion_cost(
completion_response=_response_body,
custom_llm_provider=custom_llm_provider,
call_type=CallTypes.aretrieve_batch.value,
)
verbose_logger.debug("total_cost=%s", total_cost)
return total_cost
except Exception as e:

22
litellm/cost_calculator.py
View File

@ -1892,9 +1892,16 @@ def batch_cost_calculator(
usage: Usage,
model: str,
custom_llm_provider: Optional[str] = None,
model_info: Optional[dict] = None,
) -> Tuple[float, float]:
"""
Calculate the cost of a batch job
Calculate the cost of a batch job.
Args:
model_info: Optional deployment-level model info containing custom
batch pricing (e.g. input_cost_per_token_batches). When provided,
skips the global litellm.get_model_info() lookup so that
deployment-specific pricing is used.
"""
_, custom_llm_provider, _, _ = litellm.get_llm_provider(
@ -1907,12 +1914,13 @@ def batch_cost_calculator(
custom_llm_provider,
)
try:
model_info: Optional[ModelInfo] = litellm.get_model_info(
model=model, custom_llm_provider=custom_llm_provider
)
except Exception:
model_info = None
if model_info is None:
try:
model_info = litellm.get_model_info(
model=model, custom_llm_provider=custom_llm_provider
)
except Exception:
model_info = None
if not model_info:
return 0.0, 0.0

131
tests/batches_tests/test_batch_custom_pricing.py Normal file
View File

@ -0,0 +1,131 @@
"""
Test that batch cost calculation uses custom deployment-level pricing
when model_info is provided.
Reproduces the bug where `input_cost_per_token_batches` /
`output_cost_per_token_batches` set on a proxy deployment's model_info
are ignored by the batch cost pipeline because they are never threaded
through to `batch_cost_calculator`.
"""
import pytest
from litellm.batches.batch_utils import (
_batch_cost_calculator,
_get_batch_job_cost_from_file_content,
calculate_batch_cost_and_usage,
)
from litellm.cost_calculator import batch_cost_calculator
from litellm.types.utils import Usage
# --- helpers ---
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,
}
# --- tests ---
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}"
)
def test_get_batch_job_cost_from_file_content_uses_custom_model_info():
"""_get_batch_job_cost_from_file_content should thread model_info to completion_cost."""
file_content = [_make_batch_output_line(prompt_tokens=10, completion_tokens=5)]
cost = _get_batch_job_cost_from_file_content(
file_content_dictionary=file_content,
custom_llm_provider="openai",
model_info=CUSTOM_MODEL_INFO,
)
expected = (10 * 0.00125) + (5 * 0.005)
assert cost == pytest.approx(expected), (
f"Expected total cost {expected}, got {cost}"
)
def test_batch_cost_calculator_func_uses_custom_model_info():
"""_batch_cost_calculator should thread model_info."""
file_content = [_make_batch_output_line(prompt_tokens=10, completion_tokens=5)]
cost = _batch_cost_calculator(
file_content_dictionary=file_content,
custom_llm_provider="openai",
model_info=CUSTOM_MODEL_INFO,
)
expected = (10 * 0.00125) + (5 * 0.005)
assert cost == pytest.approx(expected), (
f"Expected total cost {expected}, got {cost}"
)
@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)]
batch_cost, batch_usage, batch_models = 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 batch_cost == pytest.approx(expected), (
f"Expected total cost {expected}, got {batch_cost}"
)
assert batch_usage.prompt_tokens == 10
assert batch_usage.completion_tokens == 5