mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
refactor(batches): aggregate batch output cost, usage, and models in a single pass
Completed-batch cost tracking parsed the whole output file into a list of dicts, pretty-printed it into debug strings even with debug logging off, and walked the list three times (cost, usage, models), so a large batch output could pin a worker's memory. The output is now folded line by line into small per-line stats records via _aggregate_batch_cost_usage_models, the eager json.dumps debug calls are gone, and the raw-vertex path computes cost and usage in one call instead of two. _get_batch_output_file_content_as_dictionary becomes _fetch_batch_output_file_content (returns bytes); the superseded three-pass helpers are deleted and their tests migrated
This commit is contained in:
mateo-berri
parent
27ccc44471
commit
0b09588685
9 changed files with 294 additions and 344 deletions
|
|
@ -1,6 +1,6 @@
|
|||
{
|
||||
"reportAny": {
|
||||
"limit": 33216
|
||||
"limit": 33210
|
||||
},
|
||||
"reportArgumentType": {
|
||||
"limit": 2648
|
||||
|
|
@ -57,7 +57,7 @@
|
|||
"limit": 5893
|
||||
},
|
||||
"reportMissingTypeArgument": {
|
||||
"limit": 15886
|
||||
"limit": 15883
|
||||
},
|
||||
"reportMissingTypeStubs": {
|
||||
"limit": 41
|
||||
|
|
@ -105,19 +105,19 @@
|
|||
"limit": 113
|
||||
},
|
||||
"reportUnknownMemberType": {
|
||||
"limit": 40525
|
||||
"limit": 40523
|
||||
},
|
||||
"reportUnknownParameterType": {
|
||||
"limit": 20384
|
||||
"limit": 20381
|
||||
},
|
||||
"reportUnknownVariableType": {
|
||||
"limit": 32099
|
||||
"limit": 32095
|
||||
},
|
||||
"reportUnnecessaryCast": {
|
||||
"limit": 177
|
||||
},
|
||||
"reportUnnecessaryComparison": {
|
||||
"limit": 1023
|
||||
"limit": 1022
|
||||
},
|
||||
"reportUnnecessaryContains": {
|
||||
"limit": 7
|
||||
|
|
@ -135,7 +135,7 @@
|
|||
"limit": 33
|
||||
},
|
||||
"reportUnusedFunction": {
|
||||
"limit": 206
|
||||
"limit": 205
|
||||
},
|
||||
"reportUnusedImport": {
|
||||
"limit": 1005
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import json
|
||||
from typing import Any, Iterator, List, Literal, Optional, Tuple
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Iterable, Iterator, List, Literal, Optional, Tuple
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_logger
|
||||
|
|
@ -24,20 +25,20 @@ async def calculate_batch_cost_and_usage(
|
|||
deployment-specific pricing (e.g. input_cost_per_token_batches)
|
||||
is used instead of the global cost map.
|
||||
"""
|
||||
batch_cost = _batch_cost_calculator(
|
||||
if (
|
||||
custom_llm_provider == "vertex_ai"
|
||||
and model_name
|
||||
and getattr(litellm, "disable_vertex_batch_output_transformation", False)
|
||||
):
|
||||
batch_cost, batch_usage = calculate_vertex_ai_batch_cost_and_usage(file_content_dictionary, model_name)
|
||||
return batch_cost, batch_usage, [model_name]
|
||||
|
||||
return _aggregate_batch_cost_usage_models(
|
||||
entries=file_content_dictionary,
|
||||
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,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
model_name=model_name,
|
||||
)
|
||||
batch_models = _get_batch_models_from_file_content(file_content_dictionary, model_name, custom_llm_provider)
|
||||
|
||||
return batch_cost, batch_usage, batch_models
|
||||
|
||||
|
||||
async def _handle_completed_batch(
|
||||
|
|
@ -46,7 +47,9 @@ async def _handle_completed_batch(
|
|||
model_name: Optional[str] = None,
|
||||
litellm_params: Optional[dict] = None,
|
||||
) -> Tuple[float, Usage, List[str]]:
|
||||
"""Helper function to process a completed batch and handle logging
|
||||
"""Fetch a completed batch's output file and aggregate its cost, usage, and
|
||||
models in a single pass over the JSONL lines, so the parsed file content is
|
||||
never materialized in memory.
|
||||
|
||||
Args:
|
||||
batch: The batch object
|
||||
|
|
@ -54,75 +57,109 @@ async def _handle_completed_batch(
|
|||
model_name: Optional model name
|
||||
litellm_params: Optional litellm parameters containing credentials (api_key, api_base, etc.)
|
||||
"""
|
||||
# Get batch results
|
||||
file_content_dictionary = await _get_batch_output_file_content_as_dictionary(
|
||||
batch, custom_llm_provider, litellm_params=litellm_params
|
||||
)
|
||||
file_content = await _fetch_batch_output_file_content(batch, custom_llm_provider, litellm_params=litellm_params)
|
||||
|
||||
# Calculate costs and usage
|
||||
batch_cost = _batch_cost_calculator(
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
file_content_dictionary=file_content_dictionary,
|
||||
model_name=model_name,
|
||||
)
|
||||
batch_usage = _get_batch_job_total_usage_from_file_content(
|
||||
file_content_dictionary=file_content_dictionary,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
model_name=model_name,
|
||||
)
|
||||
|
||||
batch_models = _get_batch_models_from_file_content(file_content_dictionary, model_name, custom_llm_provider)
|
||||
|
||||
return batch_cost, batch_usage, batch_models
|
||||
|
||||
|
||||
def _get_batch_models_from_file_content(
|
||||
file_content_dictionary: List[dict],
|
||||
model_name: Optional[str] = None,
|
||||
custom_llm_provider: str = "openai",
|
||||
) -> List[str]:
|
||||
"""
|
||||
Get the models from the file content
|
||||
"""
|
||||
if model_name:
|
||||
return [model_name]
|
||||
batch_models = []
|
||||
for _item in file_content_dictionary:
|
||||
if _batch_response_was_successful(_item, custom_llm_provider):
|
||||
_response_body = _get_response_from_batch_job_output_file(_item, custom_llm_provider)
|
||||
_model = _response_body.get("model")
|
||||
if _model:
|
||||
batch_models.append(_model)
|
||||
return batch_models
|
||||
|
||||
|
||||
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[ModelInfo] = None,
|
||||
) -> float:
|
||||
"""
|
||||
Calculate the cost of a batch based on the output file id
|
||||
"""
|
||||
if (
|
||||
custom_llm_provider == "vertex_ai"
|
||||
and model_name
|
||||
and getattr(litellm, "disable_vertex_batch_output_transformation", False)
|
||||
):
|
||||
batch_cost, _ = calculate_vertex_ai_batch_cost_and_usage(file_content_dictionary, model_name)
|
||||
verbose_logger.debug("vertex_ai_total_cost=%s", batch_cost)
|
||||
return batch_cost
|
||||
batch_cost, batch_usage = calculate_vertex_ai_batch_cost_and_usage(
|
||||
_get_file_content_as_dictionary(file_content), model_name
|
||||
)
|
||||
return batch_cost, batch_usage, [model_name]
|
||||
|
||||
# For other providers, use the existing logic
|
||||
total_cost = _get_batch_job_cost_from_file_content(
|
||||
file_content_dictionary=file_content_dictionary,
|
||||
return _aggregate_batch_cost_usage_models(
|
||||
entries=_iter_batch_input_entries(file_content),
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
model_name=model_name,
|
||||
model_info=model_info,
|
||||
)
|
||||
verbose_logger.debug("total_cost=%s", total_cost)
|
||||
return total_cost
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _BatchOutputLineStats:
|
||||
cost: float
|
||||
prompt_tokens: int
|
||||
completion_tokens: int
|
||||
total_tokens: int
|
||||
cache_read_tokens: int
|
||||
cache_creation_tokens: int
|
||||
model: Optional[str]
|
||||
|
||||
|
||||
def _iter_successful_output_line_stats(
|
||||
entries: Iterable[dict],
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
|
||||
model_name: Optional[str],
|
||||
model_info: Optional[ModelInfo],
|
||||
) -> Iterator[_BatchOutputLineStats]:
|
||||
from litellm.cost_calculator import batch_cost_calculator
|
||||
|
||||
for entry in entries:
|
||||
if not _batch_response_was_successful(entry, custom_llm_provider):
|
||||
continue
|
||||
response_body = _get_response_from_batch_job_output_file(entry, custom_llm_provider)
|
||||
usage = _get_batch_job_usage_from_response_body(response_body, custom_llm_provider)
|
||||
prompt_details = _parse_prompt_tokens_details(usage)
|
||||
raw_model = response_body.get("model")
|
||||
response_model = raw_model if isinstance(raw_model, str) and raw_model else None
|
||||
if model_info is not None or custom_llm_provider in ("anthropic", "bedrock"):
|
||||
if custom_llm_provider == "bedrock" and model_name:
|
||||
cost_model = model_name
|
||||
else:
|
||||
cost_model = response_model or model_name or ""
|
||||
prompt_cost, completion_cost = batch_cost_calculator(
|
||||
usage=usage,
|
||||
model=cost_model,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
model_info=model_info,
|
||||
)
|
||||
line_cost = prompt_cost + completion_cost
|
||||
else:
|
||||
line_cost = litellm.completion_cost(
|
||||
completion_response=response_body,
|
||||
custom_llm_provider=custom_llm_provider,
|
||||
call_type=CallTypes.aretrieve_batch.value,
|
||||
)
|
||||
yield _BatchOutputLineStats(
|
||||
cost=line_cost,
|
||||
prompt_tokens=usage.prompt_tokens,
|
||||
completion_tokens=usage.completion_tokens,
|
||||
total_tokens=usage.total_tokens,
|
||||
cache_read_tokens=prompt_details["cache_hit_tokens"],
|
||||
cache_creation_tokens=prompt_details["cache_creation_tokens"],
|
||||
model=response_model,
|
||||
)
|
||||
|
||||
|
||||
def _aggregate_batch_cost_usage_models(
|
||||
entries: Iterable[dict],
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
|
||||
model_name: Optional[str] = None,
|
||||
model_info: Optional[ModelInfo] = None,
|
||||
) -> Tuple[float, Usage, List[str]]:
|
||||
"""Aggregate cost, usage, and models from batch output entries in a single
|
||||
pass, holding one small stats record per line instead of the parsed file."""
|
||||
line_stats = tuple(_iter_successful_output_line_stats(entries, custom_llm_provider, model_name, model_info))
|
||||
|
||||
cache_token_params = {
|
||||
key: tokens
|
||||
for key, tokens in (
|
||||
("cache_read_input_tokens", sum(stats.cache_read_tokens for stats in line_stats)),
|
||||
("cache_creation_input_tokens", sum(stats.cache_creation_tokens for stats in line_stats)),
|
||||
)
|
||||
if tokens > 0
|
||||
}
|
||||
batch_usage = Usage(
|
||||
total_tokens=sum(stats.total_tokens for stats in line_stats),
|
||||
prompt_tokens=sum(stats.prompt_tokens for stats in line_stats),
|
||||
completion_tokens=sum(stats.completion_tokens for stats in line_stats),
|
||||
**cache_token_params,
|
||||
)
|
||||
batch_models = [model_name] if model_name else [stats.model for stats in line_stats if stats.model]
|
||||
total_cost = sum((stats.cost for stats in line_stats), 0.0)
|
||||
verbose_logger.debug("batch output aggregate: cost=%s usage=%s models=%s", total_cost, batch_usage, batch_models)
|
||||
return total_cost, batch_usage, batch_models
|
||||
|
||||
|
||||
def calculate_vertex_ai_batch_cost_and_usage(
|
||||
|
|
@ -193,13 +230,13 @@ def calculate_vertex_ai_batch_cost_and_usage(
|
|||
)
|
||||
|
||||
|
||||
async def _get_batch_output_file_content_as_dictionary(
|
||||
async def _fetch_batch_output_file_content(
|
||||
batch: Batch,
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai",
|
||||
litellm_params: Optional[dict] = None,
|
||||
) -> List[dict]:
|
||||
) -> bytes:
|
||||
"""
|
||||
Get the batch output file content as a list of dictionaries
|
||||
Fetch the batch output file and return its raw JSONL bytes
|
||||
|
||||
Args:
|
||||
batch: The batch object
|
||||
|
|
@ -237,7 +274,7 @@ async def _get_batch_output_file_content_as_dictionary(
|
|||
file_content_kwargs.update(credentials)
|
||||
|
||||
_file_content = await afile_content(**file_content_kwargs) # type: ignore[reportArgumentType]
|
||||
return _get_file_content_as_dictionary(_file_content.content)
|
||||
return _file_content.content
|
||||
|
||||
|
||||
def _extract_file_access_credentials(litellm_params: Optional[dict]) -> dict:
|
||||
|
|
@ -283,17 +320,7 @@ def _get_file_content_as_dictionary(file_content: bytes) -> List[dict]:
|
|||
"""
|
||||
Get the file content as a list of dictionaries from JSON Lines format
|
||||
"""
|
||||
try:
|
||||
_file_content_str = file_content.decode("utf-8")
|
||||
# Split by newlines and parse each line as a separate JSON object
|
||||
json_objects = []
|
||||
for line in _file_content_str.strip().split("\n"):
|
||||
if line: # Skip empty lines
|
||||
json_objects.append(json.loads(line))
|
||||
verbose_logger.debug("json_objects=%s", json.dumps(json_objects, indent=4))
|
||||
return json_objects
|
||||
except Exception as e:
|
||||
raise e
|
||||
return list(_iter_batch_input_entries(file_content))
|
||||
|
||||
|
||||
def _iter_batch_input_lines(file_content: bytes) -> Iterator[bytes]:
|
||||
|
|
@ -360,101 +387,6 @@ def _count_entry_tokens(
|
|||
return 0
|
||||
|
||||
|
||||
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_name: Optional[str] = None,
|
||||
model_info: Optional[ModelInfo] = 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
|
||||
verbose_logger.debug("file_content_dictionary=%s", json.dumps(file_content_dictionary, indent=4))
|
||||
for _item in file_content_dictionary:
|
||||
if _batch_response_was_successful(_item, custom_llm_provider):
|
||||
_response_body = _get_response_from_batch_job_output_file(_item, custom_llm_provider)
|
||||
if model_info is not None or custom_llm_provider in ("anthropic", "bedrock"):
|
||||
usage = _get_batch_job_usage_from_response_body(_response_body, custom_llm_provider)
|
||||
# Bedrock batch output lines report a short internal model id
|
||||
# (e.g. "claude-sonnet-4-6") that is not in the cost map; use the
|
||||
# deployment model name for pricing when available.
|
||||
if custom_llm_provider == "bedrock" and model_name:
|
||||
model = model_name
|
||||
else:
|
||||
model = _response_body.get("model") or model_name or ""
|
||||
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:
|
||||
verbose_logger.error("error in _get_batch_job_cost_from_file_content", e)
|
||||
raise e
|
||||
|
||||
|
||||
def _get_batch_job_total_usage_from_file_content(
|
||||
file_content_dictionary: List[dict],
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai",
|
||||
model_name: Optional[str] = None,
|
||||
) -> Usage:
|
||||
"""
|
||||
Get the tokens of a batch job from the file content
|
||||
"""
|
||||
if (
|
||||
custom_llm_provider == "vertex_ai"
|
||||
and model_name
|
||||
and getattr(litellm, "disable_vertex_batch_output_transformation", False)
|
||||
):
|
||||
_, batch_usage = calculate_vertex_ai_batch_cost_and_usage(file_content_dictionary, model_name)
|
||||
return batch_usage
|
||||
|
||||
# For other providers, use the existing logic
|
||||
total_tokens: int = 0
|
||||
prompt_tokens: int = 0
|
||||
completion_tokens: int = 0
|
||||
cache_read_tokens: int = 0
|
||||
cache_creation_tokens: int = 0
|
||||
for _item in file_content_dictionary:
|
||||
if _batch_response_was_successful(_item, custom_llm_provider):
|
||||
_response_body = _get_response_from_batch_job_output_file(_item, custom_llm_provider)
|
||||
usage: Usage = _get_batch_job_usage_from_response_body(_response_body, custom_llm_provider)
|
||||
total_tokens += usage.total_tokens
|
||||
prompt_tokens += usage.prompt_tokens
|
||||
completion_tokens += usage.completion_tokens
|
||||
prompt_details = _parse_prompt_tokens_details(usage)
|
||||
cache_read_tokens += prompt_details["cache_hit_tokens"]
|
||||
cache_creation_tokens += prompt_details["cache_creation_tokens"]
|
||||
cache_token_params = {
|
||||
key: tokens
|
||||
for key, tokens in (
|
||||
("cache_read_input_tokens", cache_read_tokens),
|
||||
("cache_creation_input_tokens", cache_creation_tokens),
|
||||
)
|
||||
if tokens > 0
|
||||
}
|
||||
return Usage(
|
||||
total_tokens=total_tokens,
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=completion_tokens,
|
||||
**cache_token_params,
|
||||
)
|
||||
|
||||
|
||||
def _count_prompt_or_input_tokens(model: str, value: Any) -> int:
|
||||
"""Token-count a ``prompt`` / ``input`` field that the OpenAI batch
|
||||
schema allows in four shapes:
|
||||
|
|
|
|||
|
|
@ -135,7 +135,7 @@
|
|||
"limit": 30
|
||||
},
|
||||
"PERF401": {
|
||||
"limit": 144
|
||||
"limit": 143
|
||||
},
|
||||
"PERF402": {
|
||||
"limit": 9
|
||||
|
|
@ -315,16 +315,16 @@
|
|||
"limit": 98
|
||||
},
|
||||
"TRY201": {
|
||||
"limit": 424
|
||||
"limit": 422
|
||||
},
|
||||
"TRY203": {
|
||||
"limit": 123
|
||||
"limit": 122
|
||||
},
|
||||
"TRY300": {
|
||||
"limit": 883
|
||||
"limit": 881
|
||||
},
|
||||
"UP006": {
|
||||
"limit": 12147
|
||||
"limit": 12143
|
||||
},
|
||||
"UP007": {
|
||||
"limit": 2526
|
||||
|
|
@ -363,6 +363,6 @@
|
|||
"limit": 105
|
||||
},
|
||||
"UP045": {
|
||||
"limit": 17824
|
||||
"limit": 17823
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -12,8 +12,7 @@ import litellm
|
|||
import pytest
|
||||
|
||||
from litellm.batches.batch_utils import (
|
||||
_batch_cost_calculator,
|
||||
_get_batch_job_cost_from_file_content,
|
||||
_aggregate_batch_cost_usage_models,
|
||||
calculate_batch_cost_and_usage,
|
||||
)
|
||||
from litellm.cost_calculator import batch_cost_calculator
|
||||
|
|
@ -113,28 +112,12 @@ def test_batch_cost_calculator_uses_custom_model_info():
|
|||
), 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."""
|
||||
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)]
|
||||
|
||||
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,
|
||||
cost, _, _ = _aggregate_batch_cost_usage_models(
|
||||
entries=file_content,
|
||||
custom_llm_provider="openai",
|
||||
model_info=CUSTOM_MODEL_INFO,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -913,7 +913,7 @@ async def test_batch_logging_azure_credentials_regression():
|
|||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
from litellm.batches.batch_utils import (
|
||||
_extract_file_access_credentials,
|
||||
_get_batch_output_file_content_as_dictionary,
|
||||
_fetch_batch_output_file_content,
|
||||
_handle_completed_batch,
|
||||
)
|
||||
from litellm.types.llms.openai import Batch, HttpxBinaryResponseContent
|
||||
|
|
@ -996,7 +996,7 @@ async def test_batch_logging_azure_credentials_regression():
|
|||
with patch(
|
||||
"litellm.files.main.afile_content", side_effect=mock_afile_content_tracker
|
||||
):
|
||||
result = await _get_batch_output_file_content_as_dictionary(
|
||||
result = await _fetch_batch_output_file_content(
|
||||
batch=mock_batch,
|
||||
custom_llm_provider="azure",
|
||||
litellm_params=azure_credentials,
|
||||
|
|
@ -1092,7 +1092,7 @@ async def test_batch_logging_azure_credentials_regression():
|
|||
)
|
||||
|
||||
# Call without litellm_params (should still work for OpenAI)
|
||||
result = await _get_batch_output_file_content_as_dictionary(
|
||||
result = await _fetch_batch_output_file_content(
|
||||
batch=mock_batch,
|
||||
custom_llm_provider="openai",
|
||||
litellm_params=None,
|
||||
|
|
|
|||
|
|
@ -19,10 +19,8 @@ import litellm
|
|||
from litellm import create_batch, create_file
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.batches.batch_utils import (
|
||||
_batch_cost_calculator,
|
||||
_aggregate_batch_cost_usage_models,
|
||||
_get_file_content_as_dictionary,
|
||||
_get_batch_job_cost_from_file_content,
|
||||
_get_batch_job_total_usage_from_file_content,
|
||||
_get_batch_job_usage_from_response_body,
|
||||
_get_response_from_batch_job_output_file,
|
||||
_batch_response_was_successful,
|
||||
|
|
@ -139,9 +137,10 @@ def test_get_file_content_as_dictionary(sample_file_content):
|
|||
|
||||
|
||||
def test_get_batch_job_total_usage_from_file_content(sample_file_content_dict):
|
||||
usage = _get_batch_job_total_usage_from_file_content(
|
||||
sample_file_content_dict, custom_llm_provider="openai"
|
||||
)
|
||||
with patch("litellm.completion_cost", return_value=0.0):
|
||||
_, usage, _ = _aggregate_batch_cost_usage_models(
|
||||
entries=sample_file_content_dict, custom_llm_provider="openai"
|
||||
)
|
||||
assert usage.total_tokens == 62 # 30 + 32
|
||||
assert usage.prompt_tokens == 42 # 20 + 22
|
||||
assert usage.completion_tokens == 20 # 10 + 10
|
||||
|
|
@ -157,8 +156,8 @@ async def test_batch_cost_calculator(sample_file_content_dict):
|
|||
so we expect the cost to be 0.5 * 2 = 1.0
|
||||
"""
|
||||
with patch("litellm.completion_cost", return_value=0.5):
|
||||
cost = _batch_cost_calculator(
|
||||
file_content_dictionary=sample_file_content_dict,
|
||||
cost, _, _ = _aggregate_batch_cost_usage_models(
|
||||
entries=sample_file_content_dict,
|
||||
custom_llm_provider="openai",
|
||||
)
|
||||
assert cost == 1.0 # 0.5 * 2 successful responses
|
||||
|
|
@ -278,9 +277,12 @@ async def test_handle_completed_batch_computes_real_cost_from_output_file(
|
|||
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._get_batch_output_file_content_as_dictionary",
|
||||
new=AsyncMock(return_value=sample_file_content_dict),
|
||||
"litellm.batches.batch_utils._fetch_batch_output_file_content",
|
||||
new=AsyncMock(return_value=sample_file_content_bytes),
|
||||
):
|
||||
cost, usage, models = await _handle_completed_batch(
|
||||
batch=batch, custom_llm_provider="openai"
|
||||
|
|
|
|||
|
|
@ -201,29 +201,34 @@ def test_estimate_tokens_never_zero_for_short_rows():
|
|||
|
||||
|
||||
# =========================================================================== #
|
||||
# _get_batch_models_from_file_content (output file)
|
||||
# _aggregate_batch_cost_usage_models: models (output file)
|
||||
# =========================================================================== #
|
||||
|
||||
|
||||
def test_output_models_uses_model_name_override():
|
||||
# model_name short-circuits: content is ignored entirely.
|
||||
assert bu._get_batch_models_from_file_content([_success_row(model="ignored")], model_name="forced-model") == [
|
||||
"forced-model"
|
||||
]
|
||||
def test_output_models_uses_model_name_override(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0)
|
||||
_, _, models = bu._aggregate_batch_cost_usage_models(
|
||||
entries=[_success_row(model="ignored")], custom_llm_provider="openai", model_name="forced-model"
|
||||
)
|
||||
assert models == ["forced-model"]
|
||||
|
||||
|
||||
def test_output_models_collects_from_successful_only():
|
||||
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"),
|
||||
]
|
||||
assert bu._get_batch_models_from_file_content(rows) == ["gpt-4o", "claude-3"]
|
||||
_, _, models = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
|
||||
assert models == ["gpt-4o", "claude-3"]
|
||||
|
||||
|
||||
def test_output_models_skips_successful_without_model():
|
||||
def test_output_models_skips_successful_without_model(monkeypatch):
|
||||
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 0.0)
|
||||
rows = [{"response": {"status_code": 200, "body": {}}}]
|
||||
assert bu._get_batch_models_from_file_content(rows) == []
|
||||
_, _, models = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
|
||||
assert models == []
|
||||
|
||||
|
||||
# =========================================================================== #
|
||||
|
|
@ -379,17 +384,18 @@ def test_count_entry_uses_model_name_fallback(monkeypatch):
|
|||
|
||||
|
||||
# =========================================================================== #
|
||||
# _get_batch_job_total_usage_from_file_content (output usage aggregation)
|
||||
# _aggregate_batch_cost_usage_models: usage (output usage aggregation)
|
||||
# =========================================================================== #
|
||||
|
||||
|
||||
def test_total_usage_sums_successful_only():
|
||||
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
|
||||
]
|
||||
usage = bu._get_batch_job_total_usage_from_file_content(rows)
|
||||
_, usage, _ = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
|
||||
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (
|
||||
30,
|
||||
15,
|
||||
|
|
@ -398,7 +404,9 @@ def test_total_usage_sums_successful_only():
|
|||
|
||||
|
||||
def test_total_usage_empty_is_zero():
|
||||
usage = bu._get_batch_job_total_usage_from_file_content([])
|
||||
cost, usage, models = bu._aggregate_batch_cost_usage_models(entries=[], custom_llm_provider="openai")
|
||||
assert cost == 0.0
|
||||
assert models == []
|
||||
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (
|
||||
0,
|
||||
0,
|
||||
|
|
@ -407,7 +415,7 @@ def test_total_usage_empty_is_zero():
|
|||
|
||||
|
||||
# =========================================================================== #
|
||||
# _get_batch_job_cost_from_file_content (cost maps mocked)
|
||||
# _aggregate_batch_cost_usage_models: cost (cost maps mocked)
|
||||
# =========================================================================== #
|
||||
|
||||
|
||||
|
|
@ -426,7 +434,7 @@ def test_cost_from_content_completion_cost_path(monkeypatch):
|
|||
_success_row(usage=_usage(20, 10)),
|
||||
]
|
||||
|
||||
total = bu._get_batch_job_cost_from_file_content(rows, custom_llm_provider="openai")
|
||||
total, _, _ = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
|
||||
|
||||
assert total == 1.0 # 2 successful * 0.5
|
||||
assert len(calls) == 2 # failed row not costed
|
||||
|
|
@ -442,8 +450,8 @@ def test_cost_from_content_model_info_path(monkeypatch):
|
|||
_success_row(usage=_usage(20, 10)),
|
||||
]
|
||||
|
||||
total = bu._get_batch_job_cost_from_file_content(
|
||||
rows,
|
||||
total, _, _ = 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
|
||||
)
|
||||
|
|
@ -451,32 +459,65 @@ def test_cost_from_content_model_info_path(monkeypatch):
|
|||
assert total == 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))])
|
||||
|
||||
cost, usage, models = bu._aggregate_batch_cost_usage_models(entries=one_shot, custom_llm_provider="openai")
|
||||
|
||||
assert cost == 1.0
|
||||
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (30, 15, 45)
|
||||
assert models == ["gpt-4o", "gpt-4o"]
|
||||
|
||||
|
||||
# =========================================================================== #
|
||||
# _batch_cost_calculator (dispatch: vertex-disable-transform vs generic)
|
||||
# calculate_batch_cost_and_usage (dispatch: vertex-disable-transform vs generic)
|
||||
# =========================================================================== #
|
||||
|
||||
|
||||
def test_batch_cost_calculator_generic_path(monkeypatch):
|
||||
monkeypatch.setattr(bu, "_get_batch_job_cost_from_file_content", lambda **kw: 4.2)
|
||||
assert bu._batch_cost_calculator([], custom_llm_provider="openai", model_name="gpt-4o") == 4.2
|
||||
|
||||
|
||||
def test_batch_cost_calculator_vertex_disable_transform_path(monkeypatch):
|
||||
@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: (9.9, Usage()),
|
||||
lambda content, model: (9.9, Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3)),
|
||||
)
|
||||
# generic path must NOT be taken
|
||||
monkeypatch.setattr(
|
||||
bu,
|
||||
"_get_batch_job_cost_from_file_content",
|
||||
"_aggregate_batch_cost_usage_models",
|
||||
lambda **kw: pytest.fail("generic path should not run"),
|
||||
)
|
||||
|
||||
cost = bu._batch_cost_calculator([], custom_llm_provider="vertex_ai", model_name="gemini-2.0-flash-001")
|
||||
cost, usage, models = await bu.calculate_batch_cost_and_usage(
|
||||
file_content_dictionary=[], custom_llm_provider="vertex_ai", model_name="gemini-2.0-flash-001"
|
||||
)
|
||||
assert cost == 9.9
|
||||
assert usage.total_tokens == 3
|
||||
assert 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"),
|
||||
)
|
||||
|
||||
cost, usage, models = await bu.calculate_batch_cost_and_usage(
|
||||
file_content_dictionary=[], custom_llm_provider="vertex_ai"
|
||||
)
|
||||
assert cost == 0.0
|
||||
assert usage.total_tokens == 0
|
||||
assert models == []
|
||||
|
||||
|
||||
# =========================================================================== #
|
||||
|
|
@ -586,24 +627,19 @@ def test_vertex_cost_error_in_line_is_swallowed(monkeypatch):
|
|||
@pytest.mark.asyncio
|
||||
async def test_calculate_batch_cost_and_usage_orchestration(monkeypatch):
|
||||
rows = [_success_row(model="gpt-4o", usage=_usage(10, 5))]
|
||||
monkeypatch.setattr(bu, "_batch_cost_calculator", lambda **kw: 2.5)
|
||||
monkeypatch.setattr(
|
||||
bu,
|
||||
"_get_batch_job_total_usage_from_file_content",
|
||||
lambda **kw: Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
|
||||
)
|
||||
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 2.5)
|
||||
|
||||
cost, usage, models = await bu.calculate_batch_cost_and_usage(
|
||||
file_content_dictionary=rows, custom_llm_provider="openai"
|
||||
)
|
||||
|
||||
assert cost == 2.5
|
||||
assert usage.total_tokens == 15
|
||||
assert models == ["gpt-4o"] # real _get_batch_models_from_file_content
|
||||
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (10, 5, 15)
|
||||
assert models == ["gpt-4o"]
|
||||
|
||||
|
||||
# =========================================================================== #
|
||||
# _get_batch_output_file_content_as_dictionary (file fetch + credential merge)
|
||||
# _fetch_batch_output_file_content (file fetch + credential merge)
|
||||
# =========================================================================== #
|
||||
|
||||
|
||||
|
|
@ -664,7 +700,7 @@ async def test_output_file_content_vertex_fetches_via_afile_content(monkeypatch)
|
|||
|
||||
monkeypatch.setattr(files_main, "afile_content", fake_afile_content)
|
||||
|
||||
result = await bu._get_batch_output_file_content_as_dictionary(
|
||||
result = await bu._fetch_batch_output_file_content(
|
||||
_batch("gs://litellm-bucket/output/predictions.jsonl"),
|
||||
custom_llm_provider="vertex_ai",
|
||||
litellm_params={
|
||||
|
|
@ -676,7 +712,7 @@ async def test_output_file_content_vertex_fetches_via_afile_content(monkeypatch)
|
|||
},
|
||||
)
|
||||
|
||||
assert result == rows
|
||||
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"
|
||||
|
|
@ -705,7 +741,7 @@ async def test_output_file_content_vertex_unified_file_id_extracts_gcs_uri(monke
|
|||
)
|
||||
encoded_id = base64.urlsafe_b64encode(unified_id.encode()).decode().rstrip("=")
|
||||
|
||||
await bu._get_batch_output_file_content_as_dictionary(_batch(encoded_id), custom_llm_provider="vertex_ai")
|
||||
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"
|
||||
|
|
@ -740,7 +776,7 @@ async def test_handle_completed_vertex_batch_computes_cost_usage_and_models(monk
|
|||
@pytest.mark.asyncio
|
||||
async def test_output_file_content_no_output_file_id_raises():
|
||||
with pytest.raises(ValueError, match="Output file id is None"):
|
||||
await bu._get_batch_output_file_content_as_dictionary(_batch(None), custom_llm_provider="openai")
|
||||
await bu._fetch_batch_output_file_content(_batch(None), custom_llm_provider="openai")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -757,13 +793,13 @@ async def test_output_file_content_fetches_and_parses(monkeypatch):
|
|||
monkeypatch.setattr(files_main, "afile_content", fake_afile_content)
|
||||
monkeypatch.setattr(cu, "_is_base64_encoded_unified_file_id", lambda fid: False)
|
||||
|
||||
result = await bu._get_batch_output_file_content_as_dictionary(
|
||||
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 == [{"a": 1}, {"b": 2}]
|
||||
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"
|
||||
|
|
@ -792,13 +828,13 @@ async def test_output_file_content_unified_file_id_extraction(monkeypatch):
|
|||
lambda fid: "litellm_proxy;llm_output_file_id,real-file-99;rest",
|
||||
)
|
||||
|
||||
await bu._get_batch_output_file_content_as_dictionary(_batch("encoded-blob"), custom_llm_provider="openai")
|
||||
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 -> cost/usage/models)
|
||||
# _handle_completed_batch (async orchestrator: fetch -> single-pass aggregate)
|
||||
# =========================================================================== #
|
||||
|
||||
|
||||
|
|
@ -806,49 +842,48 @@ async def test_output_file_content_unified_file_id_extraction(monkeypatch):
|
|||
async def test_handle_completed_batch_orchestration(monkeypatch):
|
||||
rows = [_success_row(model="gpt-4o", usage=_usage(10, 5))]
|
||||
|
||||
async def fake_get_content(batch, custom_llm_provider, litellm_params=None):
|
||||
return rows
|
||||
async def fake_fetch(batch, custom_llm_provider, litellm_params=None):
|
||||
return _vertex_jsonl(rows)
|
||||
|
||||
monkeypatch.setattr(bu, "_get_batch_output_file_content_as_dictionary", fake_get_content)
|
||||
monkeypatch.setattr(bu, "_batch_cost_calculator", lambda **kw: 3.3)
|
||||
monkeypatch.setattr(
|
||||
bu,
|
||||
"_get_batch_job_total_usage_from_file_content",
|
||||
lambda **kw: Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15),
|
||||
)
|
||||
monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch)
|
||||
monkeypatch.setattr(litellm, "completion_cost", lambda **kw: 3.3)
|
||||
|
||||
cost, usage, models = await bu._handle_completed_batch(_batch("of"), custom_llm_provider="openai")
|
||||
|
||||
assert cost == 3.3
|
||||
assert usage.total_tokens == 15
|
||||
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (10, 5, 15)
|
||||
assert models == ["gpt-4o"]
|
||||
|
||||
|
||||
# =========================================================================== #
|
||||
# Remaining branch: vertex usage disable-transform path.
|
||||
#
|
||||
# NOTE: the error path of _get_batch_job_cost_from_file_content (its `raise e`)
|
||||
# is intentionally NOT tested: the preceding line logs via
|
||||
# `verbose_logger.error("...", e)`, which passes the exception as a logging
|
||||
# format-arg with no placeholder and itself raises TypeError under
|
||||
# logging.raiseExceptions, masking the original error. Asserting that masked
|
||||
# behavior would lock a source bug; left uncovered on purpose.
|
||||
# =========================================================================== #
|
||||
@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)
|
||||
|
||||
def test_total_usage_vertex_disable_transform_path(monkeypatch):
|
||||
monkeypatch.setattr(bu, "_fetch_batch_output_file_content", fake_fetch)
|
||||
monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", True, raising=False)
|
||||
monkeypatch.setattr(
|
||||
bu,
|
||||
"calculate_vertex_ai_batch_cost_and_usage",
|
||||
lambda content, model: (
|
||||
0.0,
|
||||
Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3),
|
||||
),
|
||||
seen: dict = {}
|
||||
|
||||
def fake_vertex_calc(content, model):
|
||||
seen["content"] = content
|
||||
seen["model"] = model
|
||||
return 7.7, Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3)
|
||||
|
||||
monkeypatch.setattr(bu, "calculate_vertex_ai_batch_cost_and_usage", fake_vertex_calc)
|
||||
|
||||
cost, usage, models = await bu._handle_completed_batch(
|
||||
_batch("gs://litellm-bucket/output/predictions.jsonl"),
|
||||
custom_llm_provider="vertex_ai",
|
||||
model_name="gemini-x",
|
||||
)
|
||||
|
||||
usage = bu._get_batch_job_total_usage_from_file_content([], custom_llm_provider="vertex_ai", model_name="gemini-x")
|
||||
assert cost == 7.7
|
||||
assert usage.total_tokens == 3
|
||||
assert 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):
|
||||
|
|
@ -948,46 +983,54 @@ def test_bedrock_cost_uses_deployment_model_name():
|
|||
"recordId": "1",
|
||||
"modelOutput": {"model": "claude-sonnet-4-6", "usage": {"input_tokens": 13, "output_tokens": 5}},
|
||||
}
|
||||
cost = bu._get_batch_job_cost_from_file_content(
|
||||
file_content_dictionary=[row],
|
||||
cost, _, models = bu._aggregate_batch_cost_usage_models(
|
||||
entries=[row],
|
||||
custom_llm_provider="bedrock",
|
||||
model_name="us.anthropic.claude-sonnet-4-6",
|
||||
model_info={},
|
||||
)
|
||||
assert cost > 0
|
||||
assert models == ["us.anthropic.claude-sonnet-4-6"]
|
||||
|
||||
|
||||
def test_anthropic_total_usage_sums_succeeded_only():
|
||||
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)),
|
||||
]
|
||||
usage = bu._get_batch_job_total_usage_from_file_content(rows, custom_llm_provider="anthropic")
|
||||
_, usage, _ = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="anthropic")
|
||||
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (130, 15, 145)
|
||||
|
||||
|
||||
def test_anthropic_total_usage_aggregates_cache_token_details():
|
||||
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)),
|
||||
]
|
||||
usage = bu._get_batch_job_total_usage_from_file_content(rows, custom_llm_provider="anthropic")
|
||||
_, usage, _ = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="anthropic")
|
||||
assert usage.prompt_tokens_details.cached_tokens == 8700
|
||||
assert usage.prompt_tokens_details.cache_creation_tokens == 2300
|
||||
assert usage.cache_read_input_tokens == 8700
|
||||
assert usage.cache_creation_input_tokens == 2300
|
||||
|
||||
|
||||
def test_total_usage_without_cache_tokens_has_no_prompt_details():
|
||||
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}}},
|
||||
}
|
||||
]
|
||||
usage = bu._get_batch_job_total_usage_from_file_content(rows, custom_llm_provider="openai")
|
||||
_, usage, _ = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
|
||||
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (10, 5, 15)
|
||||
assert usage.prompt_tokens_details is None
|
||||
|
||||
|
|
@ -1000,8 +1043,8 @@ def test_anthropic_cost_applies_batch_discount_and_cache_pricing():
|
|||
_anthropic_errored_row(),
|
||||
]
|
||||
|
||||
total = bu._get_batch_job_cost_from_file_content(
|
||||
rows,
|
||||
total, _, _ = bu._aggregate_batch_cost_usage_models(
|
||||
entries=rows,
|
||||
custom_llm_provider="anthropic",
|
||||
model_info=_ANTHROPIC_MODEL_INFO, # type: ignore[arg-type]
|
||||
)
|
||||
|
|
@ -1026,8 +1069,8 @@ def test_anthropic_cost_without_model_info_uses_batch_cost_calculator(monkeypatc
|
|||
lambda **kw: pytest.fail("anthropic rows must not go through completion_cost"),
|
||||
)
|
||||
|
||||
total = bu._get_batch_job_cost_from_file_content(
|
||||
[_anthropic_succeeded_row()], custom_llm_provider="anthropic"
|
||||
total, _, _ = bu._aggregate_batch_cost_usage_models(
|
||||
entries=[_anthropic_succeeded_row()], custom_llm_provider="anthropic"
|
||||
)
|
||||
|
||||
assert total == pytest.approx(0.3)
|
||||
|
|
@ -1036,12 +1079,16 @@ def test_anthropic_cost_without_model_info_uses_batch_cost_calculator(monkeypatc
|
|||
assert seen[0]["usage"].prompt_tokens == 10
|
||||
|
||||
|
||||
def test_anthropic_batch_models_collected_from_succeeded_rows():
|
||||
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(),
|
||||
]
|
||||
assert bu._get_batch_models_from_file_content(rows, None, "anthropic") == ["claude-sonnet-4-5-20250929"]
|
||||
_, _, models = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="anthropic")
|
||||
assert models == ["claude-sonnet-4-5-20250929"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
|
|||
|
|
@ -590,22 +590,20 @@ class TestVertexAIBatchCostCalculation:
|
|||
assert usage.completion_tokens == 0
|
||||
assert usage.total_tokens == 0
|
||||
|
||||
def test_openai_shaped_output_records_nonzero_cost_and_usage(self):
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_shaped_output_records_nonzero_cost_and_usage(self):
|
||||
"""
|
||||
Regression test for the bug where Vertex batch cost/usage was always 0.
|
||||
|
||||
After PR #25627 (transform_file_content_response), the GCS predictions.jsonl
|
||||
is rewritten into OpenAI batch shape before the cost-tracking path sees it.
|
||||
With disable_vertex_batch_output_transformation=False (default), the content
|
||||
is OpenAI-shaped, so _batch_cost_calculator must fall through to the generic
|
||||
path rather than calling calculate_vertex_ai_batch_cost_and_usage (which only
|
||||
reads raw usageMetadata fields).
|
||||
With disable_vertex_batch_output_transformation=False (default), the cost
|
||||
dispatch must fall through to the generic aggregation path rather than
|
||||
calling calculate_vertex_ai_batch_cost_and_usage (which only reads raw
|
||||
usageMetadata fields).
|
||||
"""
|
||||
import litellm
|
||||
from litellm.batches.batch_utils import (
|
||||
_batch_cost_calculator,
|
||||
_get_batch_job_total_usage_from_file_content,
|
||||
)
|
||||
from litellm.batches.batch_utils import calculate_batch_cost_and_usage
|
||||
|
||||
openai_shaped_responses = [
|
||||
{
|
||||
|
|
@ -668,12 +666,7 @@ class TestVertexAIBatchCostCalculation:
|
|||
try:
|
||||
litellm.disable_vertex_batch_output_transformation = False
|
||||
|
||||
cost = _batch_cost_calculator(
|
||||
file_content_dictionary=openai_shaped_responses,
|
||||
custom_llm_provider="vertex_ai",
|
||||
model_name="gemini-2.0-flash-001",
|
||||
)
|
||||
usage = _get_batch_job_total_usage_from_file_content(
|
||||
cost, usage, _ = await calculate_batch_cost_and_usage(
|
||||
file_content_dictionary=openai_shaped_responses,
|
||||
custom_llm_provider="vertex_ai",
|
||||
model_name="gemini-2.0-flash-001",
|
||||
|
|
@ -694,16 +687,14 @@ class TestVertexAIBatchCostCalculation:
|
|||
cost > 0
|
||||
), f"expected non-zero cost for completed Vertex batch, got {cost}"
|
||||
|
||||
def test_raw_vertex_output_still_works_when_transformation_disabled(self):
|
||||
@pytest.mark.asyncio
|
||||
async def test_raw_vertex_output_still_works_when_transformation_disabled(self):
|
||||
"""
|
||||
When disable_vertex_batch_output_transformation=True the GCS file is returned
|
||||
as raw Vertex predictions.jsonl; the specialized reader must be used.
|
||||
"""
|
||||
import litellm
|
||||
from litellm.batches.batch_utils import (
|
||||
_batch_cost_calculator,
|
||||
_get_batch_job_total_usage_from_file_content,
|
||||
)
|
||||
from litellm.batches.batch_utils import calculate_batch_cost_and_usage
|
||||
|
||||
raw_vertex_responses = [
|
||||
{
|
||||
|
|
@ -727,12 +718,7 @@ class TestVertexAIBatchCostCalculation:
|
|||
try:
|
||||
litellm.disable_vertex_batch_output_transformation = True
|
||||
|
||||
cost = _batch_cost_calculator(
|
||||
file_content_dictionary=raw_vertex_responses,
|
||||
custom_llm_provider="vertex_ai",
|
||||
model_name="gemini-2.0-flash-001",
|
||||
)
|
||||
usage = _get_batch_job_total_usage_from_file_content(
|
||||
cost, usage, _ = await calculate_batch_cost_and_usage(
|
||||
file_content_dictionary=raw_vertex_responses,
|
||||
custom_llm_provider="vertex_ai",
|
||||
model_name="gemini-2.0-flash-001",
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
{
|
||||
"LIT001": {
|
||||
"limit": 23287
|
||||
"limit": 23279
|
||||
},
|
||||
"LIT002": {
|
||||
"limit": 27473
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue