mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
feat(batches): aggregate reasoning tokens and per-line pass/fail counts
Batch retrieval already computed cost/usage on completion, but silently dropped reasoning tokens and never counted per-line success/failure. Adds BatchCostUsageResult (replacing bare cost/usage/models tuples) with successful_requests/failed_requests, and threads reasoning_tokens through the aggregated Usage. Both surface on SpendLogs the same way batch_models already does.
This commit is contained in:
mubashir1osmani
parent
973329e986
commit
141ada1118
12 changed files with 406 additions and 211 deletions
|
|
@ -651,16 +651,14 @@ class CheckBatchCost:
|
|||
# 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, # type: ignore[arg-type]
|
||||
)
|
||||
batch_result = 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, # type: ignore[arg-type]
|
||||
)
|
||||
logging_obj = LiteLLMLogging(
|
||||
model=batch_models[0],
|
||||
model=batch_result.models[0],
|
||||
messages=[{"role": "user", "content": "<retrieve_batch>"}],
|
||||
stream=False,
|
||||
call_type="aretrieve_batch",
|
||||
|
|
@ -684,9 +682,11 @@ class CheckBatchCost:
|
|||
|
||||
await logging_obj.async_success_handler(
|
||||
result=response,
|
||||
batch_cost=batch_cost,
|
||||
batch_usage=batch_usage,
|
||||
batch_models=batch_models,
|
||||
batch_cost=batch_result.cost,
|
||||
batch_usage=batch_result.usage,
|
||||
batch_models=batch_result.models,
|
||||
batch_successful_requests=batch_result.successful_requests,
|
||||
batch_failed_requests=batch_result.failed_requests,
|
||||
)
|
||||
|
||||
# Record batch duration (completed_at - created_at)
|
||||
|
|
|
|||
|
|
@ -12,12 +12,23 @@ from litellm.types.utils import CallTypes, ModelInfo, Usage
|
|||
from litellm.utils import token_counter
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class BatchCostUsageResult:
|
||||
"""Aggregate cost, usage, and per-line pass/fail counts for a completed batch."""
|
||||
|
||||
cost: float
|
||||
usage: Usage
|
||||
models: list[str]
|
||||
successful_requests: int
|
||||
failed_requests: int
|
||||
|
||||
|
||||
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: str | None = None,
|
||||
model_info: ModelInfo | None = None,
|
||||
) -> tuple[float, Usage, list[str]]:
|
||||
) -> BatchCostUsageResult:
|
||||
"""
|
||||
Calculate the cost and usage of a batch.
|
||||
|
||||
|
|
@ -32,8 +43,7 @@ async def calculate_batch_cost_and_usage(
|
|||
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 calculate_vertex_ai_batch_cost_and_usage(file_content_dictionary, model_name)
|
||||
|
||||
return _aggregate_batch_cost_usage_models(
|
||||
entries=file_content_dictionary,
|
||||
|
|
@ -48,7 +58,7 @@ async def _handle_completed_batch(
|
|||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"],
|
||||
model_name: str | None = None,
|
||||
litellm_params: dict | None = None,
|
||||
) -> tuple[float, Usage, list[str]]:
|
||||
) -> BatchCostUsageResult:
|
||||
"""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.
|
||||
|
|
@ -66,10 +76,7 @@ async def _handle_completed_batch(
|
|||
and model_name
|
||||
and getattr(litellm, "disable_vertex_batch_output_transformation", False)
|
||||
):
|
||||
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]
|
||||
return calculate_vertex_ai_batch_cost_and_usage(_get_file_content_as_dictionary(file_content), model_name)
|
||||
|
||||
return _aggregate_batch_cost_usage_models(
|
||||
entries=_iter_batch_input_entries(file_content),
|
||||
|
|
@ -86,19 +93,24 @@ class _BatchOutputLineStats:
|
|||
total_tokens: int
|
||||
cache_read_tokens: int
|
||||
cache_creation_tokens: int
|
||||
reasoning_tokens: int
|
||||
model: str | None
|
||||
|
||||
|
||||
def _iter_successful_output_line_stats(
|
||||
def _classify_output_line_stats(
|
||||
entries: Iterable[dict],
|
||||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
|
||||
model_name: str | None,
|
||||
model_info: ModelInfo | None,
|
||||
) -> Iterator[_BatchOutputLineStats]:
|
||||
) -> Iterator[_BatchOutputLineStats | None]:
|
||||
"""Classify every output line in a single pass: yields stats for a
|
||||
successful line, ``None`` for a failed one (per ``_batch_response_was_successful``).
|
||||
Counting failures this way avoids a second pass over a potentially huge output file."""
|
||||
from litellm.cost_calculator import batch_cost_calculator
|
||||
|
||||
for entry in entries:
|
||||
if not _batch_response_was_successful(entry, custom_llm_provider):
|
||||
yield None
|
||||
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)
|
||||
|
|
@ -123,6 +135,7 @@ def _iter_successful_output_line_stats(
|
|||
custom_llm_provider=custom_llm_provider,
|
||||
call_type=CallTypes.aretrieve_batch.value,
|
||||
)
|
||||
reasoning_tokens = usage.completion_tokens_details.reasoning_tokens if usage.completion_tokens_details else None
|
||||
yield _BatchOutputLineStats(
|
||||
cost=line_cost,
|
||||
prompt_tokens=usage.prompt_tokens,
|
||||
|
|
@ -130,6 +143,7 @@ def _iter_successful_output_line_stats(
|
|||
total_tokens=usage.total_tokens,
|
||||
cache_read_tokens=prompt_details["cache_hit_tokens"],
|
||||
cache_creation_tokens=prompt_details["cache_creation_tokens"],
|
||||
reasoning_tokens=reasoning_tokens or 0,
|
||||
model=response_model,
|
||||
)
|
||||
|
||||
|
|
@ -139,10 +153,14 @@ def _aggregate_batch_cost_usage_models(
|
|||
custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic", "bedrock"],
|
||||
model_name: str | None = None,
|
||||
model_info: ModelInfo | None = 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: Final = tuple(_iter_successful_output_line_stats(entries, custom_llm_provider, model_name, model_info))
|
||||
) -> BatchCostUsageResult:
|
||||
"""Aggregate cost, usage, models, and pass/fail counts from batch output
|
||||
entries in a single pass, holding one small stats record per line instead
|
||||
of the parsed file."""
|
||||
all_results: Final = tuple(_classify_output_line_stats(entries, custom_llm_provider, model_name, model_info))
|
||||
line_stats: Final = tuple(stats for stats in all_results if stats is not None)
|
||||
successful_requests: Final = len(line_stats)
|
||||
failed_requests: Final = len(all_results) - successful_requests
|
||||
|
||||
cache_token_params: Final = {
|
||||
key: tokens
|
||||
|
|
@ -156,18 +174,32 @@ def _aggregate_batch_cost_usage_models(
|
|||
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),
|
||||
reasoning_tokens=sum(stats.reasoning_tokens for stats in line_stats),
|
||||
**cache_token_params,
|
||||
)
|
||||
batch_models: Final = [model_name] if model_name else [stats.model for stats in line_stats if stats.model]
|
||||
total_cost: Final = 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
|
||||
verbose_logger.debug(
|
||||
"batch output aggregate: cost=%s usage=%s models=%s successful=%d failed=%d",
|
||||
total_cost,
|
||||
batch_usage,
|
||||
batch_models,
|
||||
successful_requests,
|
||||
failed_requests,
|
||||
)
|
||||
return BatchCostUsageResult(
|
||||
cost=total_cost,
|
||||
usage=batch_usage,
|
||||
models=batch_models,
|
||||
successful_requests=successful_requests,
|
||||
failed_requests=failed_requests,
|
||||
)
|
||||
|
||||
|
||||
def calculate_vertex_ai_batch_cost_and_usage(
|
||||
vertex_ai_batch_responses: list[dict],
|
||||
model_name: str | None = None,
|
||||
) -> tuple[float, Usage]:
|
||||
) -> BatchCostUsageResult:
|
||||
"""
|
||||
Calculate both cost and usage from raw Vertex AI batch responses.
|
||||
|
||||
|
|
@ -178,6 +210,10 @@ def calculate_vertex_ai_batch_cost_and_usage(
|
|||
{"request": ..., "response": {"candidates": [...], "usageMetadata": {...}}}
|
||||
|
||||
usageMetadata contains promptTokenCount, candidatesTokenCount, totalTokenCount.
|
||||
|
||||
A row with no ``response`` is counted as failed - the same signal already
|
||||
used to skip it from cost/usage aggregation, since Vertex batch prediction
|
||||
output doesn't establish a distinct error shape in this (non-default) path.
|
||||
"""
|
||||
from litellm.cost_calculator import batch_cost_calculator
|
||||
|
||||
|
|
@ -185,12 +221,16 @@ def calculate_vertex_ai_batch_cost_and_usage(
|
|||
total_tokens = 0
|
||||
prompt_tokens = 0
|
||||
completion_tokens = 0
|
||||
successful_requests = 0
|
||||
failed_requests = 0
|
||||
actual_model_name: Final = model_name or "gemini-2.0-flash-001"
|
||||
|
||||
for response in vertex_ai_batch_responses:
|
||||
response_body = response.get("response")
|
||||
if response_body is None:
|
||||
failed_requests += 1
|
||||
continue
|
||||
successful_requests += 1
|
||||
|
||||
usage_metadata = response_body.get("usageMetadata", {})
|
||||
_prompt = usage_metadata.get("promptTokenCount", 0) or 0
|
||||
|
|
@ -218,17 +258,25 @@ def calculate_vertex_ai_batch_cost_and_usage(
|
|||
total_tokens += _total
|
||||
|
||||
verbose_logger.info(
|
||||
"vertex_ai batch cost: cost=%s, prompt=%d, completion=%d, total=%d",
|
||||
"vertex_ai batch cost: cost=%s, prompt=%d, completion=%d, total=%d, successful=%d, failed=%d",
|
||||
total_cost,
|
||||
prompt_tokens,
|
||||
completion_tokens,
|
||||
total_tokens,
|
||||
successful_requests,
|
||||
failed_requests,
|
||||
)
|
||||
|
||||
return total_cost, Usage(
|
||||
total_tokens=total_tokens,
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=completion_tokens,
|
||||
return BatchCostUsageResult(
|
||||
cost=total_cost,
|
||||
usage=Usage(
|
||||
total_tokens=total_tokens,
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=completion_tokens,
|
||||
),
|
||||
models=[actual_model_name],
|
||||
successful_requests=successful_requests,
|
||||
failed_requests=failed_requests,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -2574,6 +2574,8 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
batch_cost: Final = kwargs.get("batch_cost", None)
|
||||
batch_usage = kwargs.get("batch_usage", None)
|
||||
batch_models = kwargs.get("batch_models", None)
|
||||
batch_successful_requests = kwargs.get("batch_successful_requests", None)
|
||||
batch_failed_requests = kwargs.get("batch_failed_requests", None)
|
||||
has_explicit_batch_data: Final = all(x is not None for x in (batch_cost, batch_usage, batch_models))
|
||||
|
||||
should_compute_batch_data: Final = (
|
||||
|
|
@ -2582,22 +2584,22 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
if has_explicit_batch_data:
|
||||
result._hidden_params["response_cost"] = batch_cost
|
||||
result._hidden_params["batch_models"] = batch_models
|
||||
result._hidden_params["batch_successful_requests"] = batch_successful_requests
|
||||
result._hidden_params["batch_failed_requests"] = batch_failed_requests
|
||||
result.usage = batch_usage
|
||||
|
||||
elif should_compute_batch_data:
|
||||
(
|
||||
response_cost,
|
||||
batch_usage,
|
||||
batch_models,
|
||||
) = await _handle_completed_batch(
|
||||
batch_result = await _handle_completed_batch(
|
||||
batch=result,
|
||||
custom_llm_provider=self.custom_llm_provider,
|
||||
litellm_params=self.litellm_params,
|
||||
)
|
||||
|
||||
result._hidden_params["response_cost"] = response_cost
|
||||
result._hidden_params["batch_models"] = batch_models
|
||||
result.usage = batch_usage
|
||||
result._hidden_params["response_cost"] = batch_result.cost
|
||||
result._hidden_params["batch_models"] = batch_result.models
|
||||
result._hidden_params["batch_successful_requests"] = batch_result.successful_requests
|
||||
result._hidden_params["batch_failed_requests"] = batch_result.failed_requests
|
||||
result.usage = batch_result.usage
|
||||
|
||||
start_time, end_time, result = self._success_handler_helper_fn(
|
||||
start_time=start_time,
|
||||
|
|
@ -5062,6 +5064,8 @@ class StandardLoggingPayloadSetup:
|
|||
additional_headers=None,
|
||||
litellm_overhead_time_ms=None,
|
||||
batch_models=None,
|
||||
batch_successful_requests=None,
|
||||
batch_failed_requests=None,
|
||||
litellm_model_name=None,
|
||||
usage_object=None,
|
||||
)
|
||||
|
|
@ -5451,6 +5455,8 @@ def _extract_response_obj_and_hidden_params(
|
|||
response_cost=None,
|
||||
litellm_overhead_time_ms=None,
|
||||
batch_models=None,
|
||||
batch_successful_requests=None,
|
||||
batch_failed_requests=None,
|
||||
litellm_model_name=None,
|
||||
usage_object=None,
|
||||
)
|
||||
|
|
@ -5819,6 +5825,8 @@ def create_dummy_standard_logging_payload() -> StandardLoggingPayload:
|
|||
additional_headers=None,
|
||||
litellm_overhead_time_ms=None,
|
||||
batch_models=None,
|
||||
batch_successful_requests=None,
|
||||
batch_failed_requests=None,
|
||||
litellm_model_name=None,
|
||||
usage_object=None,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -3462,6 +3462,8 @@ class SpendLogsMetadata(TypedDict):
|
|||
status: StandardLoggingPayloadStatus
|
||||
proxy_server_request: str | None
|
||||
batch_models: list[str] | None
|
||||
batch_successful_requests: int | None
|
||||
batch_failed_requests: int | None
|
||||
error_information: StandardLoggingPayloadErrorInformation | None
|
||||
usage_object: dict | None
|
||||
model_map_information: StandardLoggingModelInformation | None
|
||||
|
|
|
|||
|
|
@ -74,6 +74,8 @@ def _get_spend_logs_metadata(
|
|||
metadata: dict | None,
|
||||
applied_guardrails: list[str] | None = None,
|
||||
batch_models: list[str] | None = None,
|
||||
batch_successful_requests: int | None = None,
|
||||
batch_failed_requests: int | None = None,
|
||||
mcp_tool_call_metadata: StandardLoggingMCPToolCall | None = None,
|
||||
vector_store_request_metadata: list[StandardLoggingVectorStoreRequest] | None = None,
|
||||
guardrail_information: list[StandardLoggingGuardrailInformation] | None = None,
|
||||
|
|
@ -102,6 +104,8 @@ def _get_spend_logs_metadata(
|
|||
error_information=None,
|
||||
proxy_server_request=None,
|
||||
batch_models=None,
|
||||
batch_successful_requests=None,
|
||||
batch_failed_requests=None,
|
||||
mcp_tool_call_metadata=None,
|
||||
vector_store_request_metadata=None,
|
||||
model_map_information=None,
|
||||
|
|
@ -128,6 +132,8 @@ def _get_spend_logs_metadata(
|
|||
clean_metadata["user_api_key"] = _hash_api_key_for_spend_log(raw_user_api_key)
|
||||
clean_metadata["applied_guardrails"] = applied_guardrails
|
||||
clean_metadata["batch_models"] = batch_models
|
||||
clean_metadata["batch_successful_requests"] = batch_successful_requests
|
||||
clean_metadata["batch_failed_requests"] = batch_failed_requests
|
||||
clean_metadata["mcp_tool_call_metadata"] = mcp_tool_call_metadata
|
||||
clean_metadata["vector_store_request_metadata"] = _get_vector_store_request_for_spend_logs_payload(
|
||||
vector_store_request_metadata
|
||||
|
|
@ -310,6 +316,16 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs
|
|||
if standard_logging_payload is not None
|
||||
else None
|
||||
),
|
||||
batch_successful_requests=(
|
||||
standard_logging_payload.get("hidden_params", {}).get("batch_successful_requests", None)
|
||||
if standard_logging_payload is not None
|
||||
else None
|
||||
),
|
||||
batch_failed_requests=(
|
||||
standard_logging_payload.get("hidden_params", {}).get("batch_failed_requests", None)
|
||||
if standard_logging_payload is not None
|
||||
else None
|
||||
),
|
||||
mcp_tool_call_metadata=(
|
||||
standard_logging_payload["metadata"].get("mcp_tool_call_metadata", None)
|
||||
if standard_logging_payload is not None
|
||||
|
|
|
|||
|
|
@ -39,7 +39,7 @@ from pydantic import (
|
|||
field_serializer,
|
||||
field_validator,
|
||||
)
|
||||
from typing_extensions import Required, TypedDict
|
||||
from typing_extensions import ReadOnly, Required, TypedDict
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm._uuid import uuid
|
||||
|
|
@ -2880,6 +2880,8 @@ class StandardLoggingHiddenParams(TypedDict):
|
|||
litellm_overhead_time_ms: float | None
|
||||
additional_headers: StandardLoggingAdditionalHeaders | None
|
||||
batch_models: list[str] | None
|
||||
batch_successful_requests: ReadOnly[int | None]
|
||||
batch_failed_requests: ReadOnly[int | None]
|
||||
litellm_model_name: str | None # the model name sent to the provider by litellm
|
||||
usage_object: dict | None
|
||||
|
||||
|
|
|
|||
|
|
@ -116,16 +116,16 @@ 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, _, _ = _aggregate_batch_cost_usage_models(
|
||||
result = _aggregate_batch_cost_usage_models(
|
||||
entries=file_content,
|
||||
custom_llm_provider="openai",
|
||||
model_info=CUSTOM_MODEL_INFO,
|
||||
)
|
||||
|
||||
expected = (10 * 0.00125) + (5 * 0.005)
|
||||
assert cost == pytest.approx(
|
||||
assert result.cost == pytest.approx(
|
||||
expected
|
||||
), f"Expected total cost {expected}, got {cost}"
|
||||
), f"Expected total cost {expected}, got {result.cost}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("data_residency", ["eu", "us"])
|
||||
|
|
@ -164,15 +164,15 @@ 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(
|
||||
result = await calculate_batch_cost_and_usage(
|
||||
file_content_dictionary=file_content,
|
||||
custom_llm_provider="openai",
|
||||
model_info=CUSTOM_MODEL_INFO,
|
||||
)
|
||||
|
||||
expected = (10 * 0.00125) + (5 * 0.005)
|
||||
assert batch_cost == pytest.approx(
|
||||
assert result.cost == pytest.approx(
|
||||
expected
|
||||
), f"Expected total cost {expected}, got {batch_cost}"
|
||||
assert batch_usage.prompt_tokens == 10
|
||||
assert batch_usage.completion_tokens == 5
|
||||
), f"Expected total cost {expected}, got {result.cost}"
|
||||
assert result.usage.prompt_tokens == 10
|
||||
assert result.usage.completion_tokens == 5
|
||||
|
|
|
|||
|
|
@ -1027,7 +1027,7 @@ async def test_batch_logging_azure_credentials_regression():
|
|||
with patch(
|
||||
"litellm.files.main.afile_content", side_effect=mock_afile_content_tracker
|
||||
):
|
||||
cost, usage, models = await _handle_completed_batch(
|
||||
result = await _handle_completed_batch(
|
||||
batch=mock_batch,
|
||||
custom_llm_provider="azure",
|
||||
litellm_params=azure_credentials,
|
||||
|
|
@ -1039,13 +1039,13 @@ async def test_batch_logging_azure_credentials_regression():
|
|||
], "REGRESSION: Credentials not passed through _handle_completed_batch"
|
||||
|
||||
# Verify cost and usage were calculated
|
||||
assert cost > 0, "Cost should be calculated"
|
||||
assert usage.total_tokens == 40, "Usage should be calculated correctly"
|
||||
assert result.cost > 0, "Cost should be calculated"
|
||||
assert result.usage.total_tokens == 40, "Usage should be calculated correctly"
|
||||
|
||||
print(" ✓ Credentials passed through full flow")
|
||||
print(f" ✓ Cost: {cost}")
|
||||
print(f" ✓ Usage: {usage.total_tokens} tokens")
|
||||
print(f" ✓ Models: {models}")
|
||||
print(f" ✓ Cost: {result.cost}")
|
||||
print(f" ✓ Usage: {result.usage.total_tokens} tokens")
|
||||
print(f" ✓ Models: {result.models}")
|
||||
|
||||
# Test 4: Verify error prevention
|
||||
print("\n4. Testing 'Missing credentials' error prevention...")
|
||||
|
|
@ -1064,7 +1064,7 @@ async def test_batch_logging_azure_credentials_regression():
|
|||
"litellm.files.main.afile_content", side_effect=mock_afile_content_tracker
|
||||
):
|
||||
try:
|
||||
cost, usage, models = await _handle_completed_batch(
|
||||
result = await _handle_completed_batch(
|
||||
batch=mock_batch,
|
||||
custom_llm_provider="azure",
|
||||
litellm_params=azure_credentials,
|
||||
|
|
|
|||
|
|
@ -138,12 +138,12 @@ def test_get_file_content_as_dictionary(sample_file_content):
|
|||
|
||||
def test_get_batch_job_total_usage_from_file_content(sample_file_content_dict):
|
||||
with patch("litellm.completion_cost", return_value=0.0):
|
||||
_, usage, _ = _aggregate_batch_cost_usage_models(
|
||||
result = _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
|
||||
assert result.usage.total_tokens == 62 # 30 + 32
|
||||
assert result.usage.prompt_tokens == 42 # 20 + 22
|
||||
assert result.usage.completion_tokens == 20 # 10 + 10
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -156,11 +156,11 @@ 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, _, _ = _aggregate_batch_cost_usage_models(
|
||||
result = _aggregate_batch_cost_usage_models(
|
||||
entries=sample_file_content_dict,
|
||||
custom_llm_provider="openai",
|
||||
)
|
||||
assert cost == 1.0 # 0.5 * 2 successful responses
|
||||
assert result.cost == 1.0 # 0.5 * 2 successful responses
|
||||
|
||||
|
||||
def test_get_response_from_batch_job_output_file(sample_file_content_dict):
|
||||
|
|
@ -226,6 +226,8 @@ async def test_batch_retrieve_cost_tracking_with_completed_batch_no_explicit_cos
|
|||
logging_obj.custom_llm_provider = "openai"
|
||||
|
||||
# Mock _handle_completed_batch to return cost data
|
||||
from litellm.batches.batch_utils import BatchCostUsageResult
|
||||
|
||||
expected_cost = 0.05
|
||||
expected_usage = litellm.Usage(
|
||||
prompt_tokens=100,
|
||||
|
|
@ -236,7 +238,15 @@ async def test_batch_retrieve_cost_tracking_with_completed_batch_no_explicit_cos
|
|||
|
||||
with patch(
|
||||
"litellm.litellm_core_utils.litellm_logging._handle_completed_batch",
|
||||
new=AsyncMock(return_value=(expected_cost, expected_usage, expected_models)),
|
||||
new=AsyncMock(
|
||||
return_value=BatchCostUsageResult(
|
||||
cost=expected_cost,
|
||||
usage=expected_usage,
|
||||
models=expected_models,
|
||||
successful_requests=10,
|
||||
failed_requests=0,
|
||||
)
|
||||
),
|
||||
) as mock_handle_batch:
|
||||
# Call async_success_handler
|
||||
await logging_obj.async_success_handler(
|
||||
|
|
@ -251,6 +261,8 @@ async def test_batch_retrieve_cost_tracking_with_completed_batch_no_explicit_cos
|
|||
# Verify cost and usage were set on the batch result
|
||||
assert mock_batch._hidden_params["response_cost"] == expected_cost
|
||||
assert mock_batch._hidden_params["batch_models"] == expected_models
|
||||
assert mock_batch._hidden_params["batch_successful_requests"] == 10
|
||||
assert mock_batch._hidden_params["batch_failed_requests"] == 0
|
||||
assert mock_batch.usage == expected_usage
|
||||
|
||||
|
||||
|
|
@ -284,7 +296,7 @@ async def test_handle_completed_batch_computes_real_cost_from_output_file(
|
|||
"litellm.batches.batch_utils._fetch_batch_output_file_content",
|
||||
new=AsyncMock(return_value=sample_file_content_bytes),
|
||||
):
|
||||
cost, usage, models = await _handle_completed_batch(
|
||||
result = await _handle_completed_batch(
|
||||
batch=batch, custom_llm_provider="openai"
|
||||
)
|
||||
|
||||
|
|
@ -294,16 +306,18 @@ async def test_handle_completed_batch_computes_real_cost_from_output_file(
|
|||
+ 20 * pricing["output_cost_per_token_batches"]
|
||||
)
|
||||
|
||||
assert cost == pytest.approx(expected_cost)
|
||||
assert cost > 0
|
||||
assert result.cost == pytest.approx(expected_cost)
|
||||
assert result.cost > 0
|
||||
assert (
|
||||
cost
|
||||
result.cost
|
||||
< 42 * pricing["input_cost_per_token"] + 20 * pricing["output_cost_per_token"]
|
||||
)
|
||||
assert usage.prompt_tokens == 42
|
||||
assert usage.completion_tokens == 20
|
||||
assert usage.total_tokens == 62
|
||||
assert models == ["gpt-4o-mini-2024-07-18", "gpt-4o-mini-2024-07-18"]
|
||||
assert result.usage.prompt_tokens == 42
|
||||
assert result.usage.completion_tokens == 20
|
||||
assert result.usage.total_tokens == 62
|
||||
assert result.models == ["gpt-4o-mini-2024-07-18", "gpt-4o-mini-2024-07-18"]
|
||||
assert result.successful_requests == 2
|
||||
assert result.failed_requests == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -542,9 +556,19 @@ async def test_batch_retrieve_cost_tracking_with_partial_explicit_data():
|
|||
)
|
||||
expected_models = ["gpt-5-mini"]
|
||||
|
||||
from litellm.batches.batch_utils import BatchCostUsageResult
|
||||
|
||||
with patch(
|
||||
"litellm.litellm_core_utils.litellm_logging._handle_completed_batch",
|
||||
new=AsyncMock(return_value=(expected_cost, expected_usage, expected_models)),
|
||||
new=AsyncMock(
|
||||
return_value=BatchCostUsageResult(
|
||||
cost=expected_cost,
|
||||
usage=expected_usage,
|
||||
models=expected_models,
|
||||
successful_requests=8,
|
||||
failed_requests=0,
|
||||
)
|
||||
),
|
||||
) as mock_handle_batch:
|
||||
# Call async_success_handler with partial explicit data
|
||||
await logging_obj.async_success_handler(
|
||||
|
|
@ -560,4 +584,6 @@ async def test_batch_retrieve_cost_tracking_with_partial_explicit_data():
|
|||
# Verify computed cost data was used (not partial explicit data)
|
||||
assert mock_batch._hidden_params["response_cost"] == expected_cost
|
||||
assert mock_batch._hidden_params["batch_models"] == expected_models
|
||||
assert mock_batch._hidden_params["batch_successful_requests"] == 8
|
||||
assert mock_batch._hidden_params["batch_failed_requests"] == 0
|
||||
assert mock_batch.usage == expected_usage
|
||||
|
|
|
|||
|
|
@ -13,6 +13,20 @@ import pytest
|
|||
_IS_B64 = "litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id"
|
||||
|
||||
|
||||
def _batch_cost_result(cost, usage, models, successful_requests=1, failed_requests=0):
|
||||
"""Build the BatchCostUsageResult calculate_batch_cost_and_usage now returns,
|
||||
for mocking it in tests that only care about cost/usage/models."""
|
||||
from litellm.batches.batch_utils import BatchCostUsageResult
|
||||
|
||||
return BatchCostUsageResult(
|
||||
cost=cost,
|
||||
usage=usage,
|
||||
models=models,
|
||||
successful_requests=successful_requests,
|
||||
failed_requests=failed_requests,
|
||||
)
|
||||
|
||||
|
||||
def _unmanaged_vertex_file_object(
|
||||
input_file_id="gs://bucket/litellm-vertex-files/publishers/google/models/gemini-2.5-flash/abc.jsonl",
|
||||
status="validating",
|
||||
|
|
@ -321,7 +335,7 @@ class TestCheckBatchCost:
|
|||
patch(
|
||||
"litellm.batches.batch_utils.calculate_batch_cost_and_usage",
|
||||
new_callable=AsyncMock,
|
||||
return_value=(
|
||||
return_value=_batch_cost_result(
|
||||
0.01,
|
||||
{"prompt_tokens": 10, "completion_tokens": 5},
|
||||
["gpt-4"],
|
||||
|
|
@ -426,7 +440,7 @@ class TestCheckBatchCost:
|
|||
patch(
|
||||
"litellm.batches.batch_utils.calculate_batch_cost_and_usage",
|
||||
new_callable=AsyncMock,
|
||||
return_value=(0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["claude-haiku-4-5"]),
|
||||
return_value=_batch_cost_result(0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["claude-haiku-4-5"]),
|
||||
),
|
||||
patch(
|
||||
"litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider",
|
||||
|
|
@ -526,7 +540,7 @@ class TestCheckBatchCost:
|
|||
patch(
|
||||
"litellm.batches.batch_utils.calculate_batch_cost_and_usage",
|
||||
new_callable=AsyncMock,
|
||||
return_value=(
|
||||
return_value=_batch_cost_result(
|
||||
0.01,
|
||||
{"prompt_tokens": 10, "completion_tokens": 5},
|
||||
["gpt-4"],
|
||||
|
|
@ -656,7 +670,7 @@ class TestCheckBatchCost:
|
|||
patch(
|
||||
"litellm.batches.batch_utils.calculate_batch_cost_and_usage",
|
||||
new_callable=AsyncMock,
|
||||
return_value=(
|
||||
return_value=_batch_cost_result(
|
||||
0.01,
|
||||
{"prompt_tokens": 10, "completion_tokens": 5},
|
||||
["gpt-4"],
|
||||
|
|
@ -1142,7 +1156,7 @@ class TestCheckBatchCost:
|
|||
patch(
|
||||
"litellm.batches.batch_utils.calculate_batch_cost_and_usage",
|
||||
new_callable=AsyncMock,
|
||||
return_value=(
|
||||
return_value=_batch_cost_result(
|
||||
0.01,
|
||||
{"prompt_tokens": 10, "completion_tokens": 5},
|
||||
["gpt-4"],
|
||||
|
|
@ -1271,7 +1285,7 @@ class TestCheckBatchCost:
|
|||
patch(
|
||||
"litellm.batches.batch_utils.calculate_batch_cost_and_usage",
|
||||
new_callable=AsyncMock,
|
||||
return_value=(
|
||||
return_value=_batch_cost_result(
|
||||
0.01,
|
||||
{"prompt_tokens": 10, "completion_tokens": 5},
|
||||
["gpt-4"],
|
||||
|
|
@ -1529,7 +1543,7 @@ class TestUnmanagedVertexRouting:
|
|||
patch(
|
||||
"litellm.batches.batch_utils.calculate_batch_cost_and_usage",
|
||||
new_callable=AsyncMock,
|
||||
return_value=(
|
||||
return_value=_batch_cost_result(
|
||||
0.01,
|
||||
{"prompt_tokens": 10, "completion_tokens": 5},
|
||||
["gemini-2.5-flash"],
|
||||
|
|
@ -1759,7 +1773,7 @@ class TestUnmanagedBedrockRouting:
|
|||
patch(
|
||||
"litellm.batches.batch_utils.calculate_batch_cost_and_usage",
|
||||
new_callable=AsyncMock,
|
||||
return_value=(
|
||||
return_value=_batch_cost_result(
|
||||
0.02,
|
||||
{"prompt_tokens": 10, "completion_tokens": 5},
|
||||
["claude-sonnet-4"],
|
||||
|
|
@ -1951,7 +1965,7 @@ class TestManagedOutputFileIdEncodesPublicModelGroup:
|
|||
patch(
|
||||
"litellm.batches.batch_utils.calculate_batch_cost_and_usage",
|
||||
new_callable=AsyncMock,
|
||||
return_value=(0.01, {"prompt_tokens": 10}, ["gpt-5.5"]),
|
||||
return_value=_batch_cost_result(0.01, {"prompt_tokens": 10}, ["gpt-5.5"]),
|
||||
),
|
||||
patch("litellm.litellm_core_utils.litellm_logging.Logging") as logging_cls,
|
||||
):
|
||||
|
|
|
|||
|
|
@ -210,10 +210,10 @@ def test_estimate_tokens_never_zero_for_short_rows():
|
|||
|
||||
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(
|
||||
result = bu._aggregate_batch_cost_usage_models(
|
||||
entries=[_success_row(model="ignored")], custom_llm_provider="openai", model_name="forced-model"
|
||||
)
|
||||
assert models == ["forced-model"]
|
||||
assert result.models == ["forced-model"]
|
||||
|
||||
|
||||
def test_output_models_collects_from_successful_only(monkeypatch):
|
||||
|
|
@ -223,15 +223,15 @@ def test_output_models_collects_from_successful_only(monkeypatch):
|
|||
_failed_row(model="should-be-skipped"),
|
||||
_success_row(model="claude-3"),
|
||||
]
|
||||
_, _, models = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
|
||||
assert models == ["gpt-4o", "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": {}}}]
|
||||
_, _, models = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
|
||||
assert models == []
|
||||
result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
|
||||
assert result.models == []
|
||||
|
||||
|
||||
# =========================================================================== #
|
||||
|
|
@ -398,8 +398,8 @@ def test_total_usage_sums_successful_only(monkeypatch):
|
|||
_failed_row(), # excluded
|
||||
_success_row(usage=_usage(20, 10)), # 30
|
||||
]
|
||||
_, usage, _ = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
|
||||
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (
|
||||
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,
|
||||
|
|
@ -417,7 +417,7 @@ def test_total_usage_and_cost_normalize_mixed_responses_and_chat():
|
|||
)
|
||||
chat_row = _success_row(usage=_usage(10, 5))
|
||||
|
||||
cost, usage, _ = bu._aggregate_batch_cost_usage_models(
|
||||
result = bu._aggregate_batch_cost_usage_models(
|
||||
entries=[responses_row, chat_row],
|
||||
custom_llm_provider="openai",
|
||||
model_info={
|
||||
|
|
@ -426,22 +426,79 @@ def test_total_usage_and_cost_normalize_mixed_responses_and_chat():
|
|||
},
|
||||
)
|
||||
|
||||
assert usage.prompt_tokens == 30
|
||||
assert usage.completion_tokens == 12
|
||||
assert usage.total_tokens == 42
|
||||
assert usage.cache_read_input_tokens == 3
|
||||
assert cost == pytest.approx((30 * 0.00125) + (12 * 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():
|
||||
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) == (
|
||||
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,
|
||||
)
|
||||
|
||||
|
||||
# =========================================================================== #
|
||||
|
|
@ -464,10 +521,12 @@ def test_cost_from_content_completion_cost_path(monkeypatch):
|
|||
_success_row(usage=_usage(20, 10)),
|
||||
]
|
||||
|
||||
total, _, _ = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
|
||||
result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="openai")
|
||||
|
||||
assert total == 1.0 # 2 successful * 0.5
|
||||
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_cost_from_content_model_info_path(monkeypatch):
|
||||
|
|
@ -480,13 +539,13 @@ def test_cost_from_content_model_info_path(monkeypatch):
|
|||
_success_row(usage=_usage(20, 10)),
|
||||
]
|
||||
|
||||
total, _, _ = bu._aggregate_batch_cost_usage_models(
|
||||
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 total == pytest.approx(0.6) # 2 * (0.1 + 0.2)
|
||||
assert result.cost == pytest.approx(0.6) # 2 * (0.1 + 0.2)
|
||||
|
||||
|
||||
def test_aggregate_consumes_entries_in_a_single_pass(monkeypatch):
|
||||
|
|
@ -496,11 +555,13 @@ def test_aggregate_consumes_entries_in_a_single_pass(monkeypatch):
|
|||
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")
|
||||
result = 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"]
|
||||
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
|
||||
|
||||
|
||||
# =========================================================================== #
|
||||
|
|
@ -514,7 +575,13 @@ async def test_calculate_vertex_disable_transform_path(monkeypatch):
|
|||
monkeypatch.setattr(
|
||||
bu,
|
||||
"calculate_vertex_ai_batch_cost_and_usage",
|
||||
lambda content, model: (9.9, Usage(prompt_tokens=1, completion_tokens=2, total_tokens=3)),
|
||||
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(
|
||||
|
|
@ -523,12 +590,12 @@ async def test_calculate_vertex_disable_transform_path(monkeypatch):
|
|||
lambda **kw: pytest.fail("generic path should not run"),
|
||||
)
|
||||
|
||||
cost, usage, models = await bu.calculate_batch_cost_and_usage(
|
||||
result = 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"]
|
||||
assert result.cost == 9.9
|
||||
assert result.usage.total_tokens == 3
|
||||
assert result.models == ["gemini-2.0-flash-001"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -542,12 +609,12 @@ async def test_calculate_vertex_disable_transform_needs_model_name(monkeypatch):
|
|||
lambda content, model: pytest.fail("raw vertex path should not run"),
|
||||
)
|
||||
|
||||
cost, usage, models = await bu.calculate_batch_cost_and_usage(
|
||||
result = 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 == []
|
||||
assert result.cost == 0.0
|
||||
assert result.usage.total_tokens == 0
|
||||
assert result.models == []
|
||||
|
||||
|
||||
# =========================================================================== #
|
||||
|
|
@ -580,14 +647,16 @@ def test_vertex_cost_and_usage_aggregation(monkeypatch):
|
|||
},
|
||||
]
|
||||
|
||||
cost, usage = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x")
|
||||
result = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x")
|
||||
|
||||
assert cost == pytest.approx(0.6) # 2 * (0.1 + 0.2)
|
||||
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (
|
||||
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):
|
||||
|
|
@ -607,10 +676,12 @@ def test_vertex_cost_skips_none_response_body(monkeypatch):
|
|||
},
|
||||
]
|
||||
|
||||
cost, usage = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x")
|
||||
result = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x")
|
||||
|
||||
assert cost == pytest.approx(1.0) # only one line costed
|
||||
assert usage.total_tokens == 10
|
||||
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):
|
||||
|
|
@ -620,8 +691,8 @@ def test_vertex_usage_total_token_fallback(monkeypatch):
|
|||
monkeypatch.setattr(cc, "batch_cost_calculator", lambda **kw: (0.0, 0.0))
|
||||
responses = [{"response": {"usageMetadata": {"promptTokenCount": 8, "candidatesTokenCount": 4}}}]
|
||||
|
||||
_, usage = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x")
|
||||
assert usage.total_tokens == 12
|
||||
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):
|
||||
|
|
@ -644,9 +715,9 @@ def test_vertex_cost_error_in_line_is_swallowed(monkeypatch):
|
|||
}
|
||||
]
|
||||
|
||||
cost, usage = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x")
|
||||
assert cost == 0.0
|
||||
assert usage.total_tokens == 10
|
||||
result = bu.calculate_vertex_ai_batch_cost_and_usage(responses, "gemini-x")
|
||||
assert result.cost == 0.0
|
||||
assert result.usage.total_tokens == 10
|
||||
|
||||
|
||||
# =========================================================================== #
|
||||
|
|
@ -659,13 +730,11 @@ 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)
|
||||
|
||||
cost, usage, models = await bu.calculate_batch_cost_and_usage(
|
||||
file_content_dictionary=rows, custom_llm_provider="openai"
|
||||
)
|
||||
result = await bu.calculate_batch_cost_and_usage(file_content_dictionary=rows, custom_llm_provider="openai")
|
||||
|
||||
assert cost == 2.5
|
||||
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (10, 5, 15)
|
||||
assert models == ["gpt-4o"]
|
||||
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"]
|
||||
|
||||
|
||||
# =========================================================================== #
|
||||
|
|
@ -883,16 +952,18 @@ async def test_handle_completed_vertex_batch_computes_cost_usage_and_models(monk
|
|||
|
||||
monkeypatch.setattr(files_main, "afile_content", fake_afile_content)
|
||||
|
||||
cost, usage, models = await bu._handle_completed_batch(
|
||||
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"},
|
||||
)
|
||||
|
||||
assert cost > 0
|
||||
assert cost == pytest.approx(30 * 7.5e-07 + 15 * 3.75e-06)
|
||||
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (30, 15, 45)
|
||||
assert models == ["gemini-3.6-flash", "gemini-3.6-flash"]
|
||||
assert result.cost > 0
|
||||
assert result.cost == pytest.approx(30 * 7.5e-07 + 15 * 3.75e-06)
|
||||
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
|
||||
|
|
@ -970,11 +1041,11 @@ async def test_handle_completed_batch_orchestration(monkeypatch):
|
|||
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")
|
||||
result = await bu._handle_completed_batch(_batch("of"), custom_llm_provider="openai")
|
||||
|
||||
assert cost == 3.3
|
||||
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (10, 5, 15)
|
||||
assert models == ["gpt-4o"]
|
||||
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
|
||||
|
|
@ -991,19 +1062,25 @@ async def test_handle_completed_batch_vertex_disable_transform_path(monkeypatch)
|
|||
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)
|
||||
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)
|
||||
|
||||
cost, usage, models = await bu._handle_completed_batch(
|
||||
result = await bu._handle_completed_batch(
|
||||
_batch("gs://litellm-bucket/output/predictions.jsonl"),
|
||||
custom_llm_provider="vertex_ai",
|
||||
model_name="gemini-x",
|
||||
)
|
||||
|
||||
assert cost == 7.7
|
||||
assert usage.total_tokens == 3
|
||||
assert models == ["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"
|
||||
|
||||
|
|
@ -1105,14 +1182,14 @@ def test_bedrock_cost_uses_deployment_model_name():
|
|||
"recordId": "1",
|
||||
"modelOutput": {"model": "claude-sonnet-4-6", "usage": {"input_tokens": 13, "output_tokens": 5}},
|
||||
}
|
||||
cost, _, models = bu._aggregate_batch_cost_usage_models(
|
||||
result = 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"]
|
||||
assert result.cost > 0
|
||||
assert result.models == ["us.anthropic.claude-sonnet-4-6"]
|
||||
|
||||
|
||||
def test_anthropic_total_usage_sums_succeeded_only(monkeypatch):
|
||||
|
|
@ -1124,8 +1201,10 @@ def test_anthropic_total_usage_sums_succeeded_only(monkeypatch):
|
|||
_anthropic_errored_row(),
|
||||
_anthropic_succeeded_row(usage=_anthropic_usage(20, 10, cache_read=100)),
|
||||
]
|
||||
_, 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)
|
||||
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):
|
||||
|
|
@ -1137,11 +1216,11 @@ def test_anthropic_total_usage_aggregates_cache_token_details(monkeypatch):
|
|||
_anthropic_errored_row(),
|
||||
_anthropic_succeeded_row(usage=_anthropic_usage(50, 20, cache_creation=300, cache_read=700)),
|
||||
]
|
||||
_, 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
|
||||
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):
|
||||
|
|
@ -1152,9 +1231,9 @@ def test_total_usage_without_cache_tokens_has_no_prompt_details(monkeypatch):
|
|||
"response": {"status_code": 200, "body": {"model": "gpt-5.2", "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}}},
|
||||
}
|
||||
]
|
||||
_, 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
|
||||
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():
|
||||
|
|
@ -1165,14 +1244,14 @@ def test_anthropic_cost_applies_batch_discount_and_cache_pricing():
|
|||
_anthropic_errored_row(),
|
||||
]
|
||||
|
||||
total, _, _ = bu._aggregate_batch_cost_usage_models(
|
||||
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 total == pytest.approx(expected_half_price)
|
||||
assert result.cost == pytest.approx(expected_half_price)
|
||||
|
||||
|
||||
def test_anthropic_cost_without_model_info_uses_batch_cost_calculator(monkeypatch):
|
||||
|
|
@ -1191,11 +1270,9 @@ 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._aggregate_batch_cost_usage_models(
|
||||
entries=[_anthropic_succeeded_row()], custom_llm_provider="anthropic"
|
||||
)
|
||||
result = bu._aggregate_batch_cost_usage_models(entries=[_anthropic_succeeded_row()], custom_llm_provider="anthropic")
|
||||
|
||||
assert total == pytest.approx(0.3)
|
||||
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
|
||||
|
|
@ -1209,8 +1286,8 @@ def test_anthropic_batch_models_collected_from_succeeded_rows(monkeypatch):
|
|||
_anthropic_succeeded_row(model="claude-sonnet-4-5-20250929"),
|
||||
_anthropic_errored_row(),
|
||||
]
|
||||
_, _, models = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="anthropic")
|
||||
assert models == ["claude-sonnet-4-5-20250929"]
|
||||
result = bu._aggregate_batch_cost_usage_models(entries=rows, custom_llm_provider="anthropic")
|
||||
assert result.models == ["claude-sonnet-4-5-20250929"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -1220,16 +1297,16 @@ async def test_calculate_batch_cost_and_usage_anthropic_end_to_end():
|
|||
_anthropic_errored_row(),
|
||||
]
|
||||
|
||||
cost, usage, models = await bu.calculate_batch_cost_and_usage(
|
||||
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 cost == pytest.approx(1000 * 3e-6 / 2 + 8000 * 3e-7 / 2 + 2000 * 3.75e-6 / 2 + 200 * 15e-6 / 2)
|
||||
assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (11000, 200, 11200)
|
||||
assert models == ["claude-sonnet-4-5"]
|
||||
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():
|
||||
|
|
|
|||
|
|
@ -478,14 +478,14 @@ class TestVertexAIBatchPassthroughHandler:
|
|||
}
|
||||
]
|
||||
|
||||
total_cost, usage = calculate_vertex_ai_batch_cost_and_usage(
|
||||
result = calculate_vertex_ai_batch_cost_and_usage(
|
||||
vertex_ai_batch_responses, model_name="gemini-2.0-flash-001"
|
||||
)
|
||||
|
||||
assert usage.total_tokens == 15
|
||||
assert usage.prompt_tokens == 10
|
||||
assert usage.completion_tokens == 5
|
||||
assert total_cost > 0, "batch_cost_calculator should return a non-zero cost"
|
||||
assert result.usage.total_tokens == 15
|
||||
assert result.usage.prompt_tokens == 10
|
||||
assert result.usage.completion_tokens == 5
|
||||
assert result.cost > 0, "batch_cost_calculator should return a non-zero cost"
|
||||
|
||||
def test_batch_response_transformation(self):
|
||||
"""Test transformation of Vertex AI batch responses to OpenAI format"""
|
||||
|
|
@ -664,14 +664,14 @@ class TestVertexAIBatchCostCalculation:
|
|||
},
|
||||
]
|
||||
|
||||
total_cost, usage = calculate_vertex_ai_batch_cost_and_usage(
|
||||
result = calculate_vertex_ai_batch_cost_and_usage(
|
||||
responses, model_name="gemini-2.0-flash-001"
|
||||
)
|
||||
|
||||
assert usage.prompt_tokens == 18
|
||||
assert usage.completion_tokens == 8
|
||||
assert usage.total_tokens == 26
|
||||
assert total_cost > 0, "batch_cost_calculator should return a non-zero cost"
|
||||
assert result.usage.prompt_tokens == 18
|
||||
assert result.usage.completion_tokens == 8
|
||||
assert result.usage.total_tokens == 26
|
||||
assert result.cost > 0, "batch_cost_calculator should return a non-zero cost"
|
||||
|
||||
def test_should_skip_responses_with_null_response_body(self):
|
||||
"""Failed lines (response: None) are skipped without error."""
|
||||
|
|
@ -699,27 +699,29 @@ class TestVertexAIBatchCostCalculation:
|
|||
},
|
||||
]
|
||||
|
||||
total_cost, usage = calculate_vertex_ai_batch_cost_and_usage(
|
||||
result = calculate_vertex_ai_batch_cost_and_usage(
|
||||
responses, model_name="gemini-2.0-flash-001"
|
||||
)
|
||||
|
||||
assert usage.prompt_tokens == 18
|
||||
assert usage.completion_tokens == 8
|
||||
assert usage.total_tokens == 26
|
||||
assert total_cost > 0
|
||||
assert result.usage.prompt_tokens == 18
|
||||
assert result.usage.completion_tokens == 8
|
||||
assert result.usage.total_tokens == 26
|
||||
assert result.cost > 0
|
||||
assert result.successful_requests == 2
|
||||
assert result.failed_requests == 1
|
||||
|
||||
def test_should_return_zeros_for_empty_response_list(self):
|
||||
"""Empty input → zero cost and zero usage."""
|
||||
from litellm.batches.batch_utils import calculate_vertex_ai_batch_cost_and_usage
|
||||
|
||||
total_cost, usage = calculate_vertex_ai_batch_cost_and_usage(
|
||||
result = calculate_vertex_ai_batch_cost_and_usage(
|
||||
[], model_name="gemini-2.0-flash-001"
|
||||
)
|
||||
|
||||
assert total_cost == 0.0
|
||||
assert usage.total_tokens == 0
|
||||
assert usage.prompt_tokens == 0
|
||||
assert usage.completion_tokens == 0
|
||||
assert result.cost == 0.0
|
||||
assert result.usage.total_tokens == 0
|
||||
assert result.usage.prompt_tokens == 0
|
||||
assert result.usage.completion_tokens == 0
|
||||
|
||||
def test_should_handle_missing_usage_metadata_gracefully(self):
|
||||
"""Response without usageMetadata → 0 tokens, 0 cost for that line."""
|
||||
|
|
@ -729,13 +731,13 @@ class TestVertexAIBatchCostCalculation:
|
|||
{"response": {"candidates": [{"content": {"parts": [{"text": "hi"}]}}]}},
|
||||
]
|
||||
|
||||
total_cost, usage = calculate_vertex_ai_batch_cost_and_usage(
|
||||
result = calculate_vertex_ai_batch_cost_and_usage(
|
||||
responses, model_name="gemini-2.0-flash-001"
|
||||
)
|
||||
|
||||
assert usage.prompt_tokens == 0
|
||||
assert usage.completion_tokens == 0
|
||||
assert usage.total_tokens == 0
|
||||
assert result.usage.prompt_tokens == 0
|
||||
assert result.usage.completion_tokens == 0
|
||||
assert result.usage.total_tokens == 0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_shaped_output_records_nonzero_cost_and_usage(self):
|
||||
|
|
@ -813,7 +815,7 @@ class TestVertexAIBatchCostCalculation:
|
|||
try:
|
||||
litellm.disable_vertex_batch_output_transformation = False
|
||||
|
||||
cost, usage, _ = await calculate_batch_cost_and_usage(
|
||||
result = await calculate_batch_cost_and_usage(
|
||||
file_content_dictionary=openai_shaped_responses,
|
||||
custom_llm_provider="vertex_ai",
|
||||
model_name="gemini-2.0-flash-001",
|
||||
|
|
@ -822,17 +824,17 @@ class TestVertexAIBatchCostCalculation:
|
|||
litellm.disable_vertex_batch_output_transformation = original_flag
|
||||
|
||||
assert (
|
||||
usage.prompt_tokens == 18
|
||||
), f"expected 18 prompt tokens, got {usage.prompt_tokens}"
|
||||
result.usage.prompt_tokens == 18
|
||||
), f"expected 18 prompt tokens, got {result.usage.prompt_tokens}"
|
||||
assert (
|
||||
usage.completion_tokens == 8
|
||||
), f"expected 8 completion tokens, got {usage.completion_tokens}"
|
||||
result.usage.completion_tokens == 8
|
||||
), f"expected 8 completion tokens, got {result.usage.completion_tokens}"
|
||||
assert (
|
||||
usage.total_tokens == 26
|
||||
), f"expected 26 total tokens, got {usage.total_tokens}"
|
||||
result.usage.total_tokens == 26
|
||||
), f"expected 26 total tokens, got {result.usage.total_tokens}"
|
||||
assert (
|
||||
cost > 0
|
||||
), f"expected non-zero cost for completed Vertex batch, got {cost}"
|
||||
result.cost > 0
|
||||
), f"expected non-zero cost for completed Vertex batch, got {result.cost}"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_raw_vertex_output_still_works_when_transformation_disabled(self):
|
||||
|
|
@ -865,7 +867,7 @@ class TestVertexAIBatchCostCalculation:
|
|||
try:
|
||||
litellm.disable_vertex_batch_output_transformation = True
|
||||
|
||||
cost, usage, _ = await calculate_batch_cost_and_usage(
|
||||
result = await calculate_batch_cost_and_usage(
|
||||
file_content_dictionary=raw_vertex_responses,
|
||||
custom_llm_provider="vertex_ai",
|
||||
model_name="gemini-2.0-flash-001",
|
||||
|
|
@ -873,7 +875,7 @@ class TestVertexAIBatchCostCalculation:
|
|||
finally:
|
||||
litellm.disable_vertex_batch_output_transformation = original_flag
|
||||
|
||||
assert usage.prompt_tokens == 10
|
||||
assert usage.completion_tokens == 5
|
||||
assert usage.total_tokens == 15
|
||||
assert cost > 0, "raw Vertex shape should also produce non-zero cost"
|
||||
assert result.usage.prompt_tokens == 10
|
||||
assert result.usage.completion_tokens == 5
|
||||
assert result.usage.total_tokens == 15
|
||||
assert result.cost > 0, "raw Vertex shape should also produce non-zero cost"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue