mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-01 07:03:02 +00:00
fix(batches): return real batch token counts next to OpenAI's batch usage fields (#47740)
* fix(batches): read OpenAI batch usage into the batch token counts * fix(batches): serve OpenAI's batch usage shape from retrieve and the cost poller * fix(batches): keep the chat usage keys next to OpenAI's batch usage fields Batch usage now carries every chat key it had plus OpenAI's five batch fields with the same counts, so poller-billed batches keep prompt_tokens and cache-write tokens, and older proxies still read the new rows. Usage with missing or null detail counts reads its real input and output counts. * test(batches): check the batch usage counts in the e2e terminal-state test * fix(types): apply nested include and exclude to both batch usage shapes * fix(batches): log no cost for a batch retrieve that is not final yet A retrieve now reads OpenAI's batch usage into the chat counts, so a retrieve that saw the batch completed before its output file arrived priced those counts for the logging callbacks, and the final retrieve priced the whole batch again. A batch that is not final now logs no cost, as it did before. Batch usage dumps also pass exclude_unset and exclude_defaults through to the usage keys. * fix(batches): pass through every usage detail field the provider sends Batch dumps rebuilt input_tokens_details and output_tokens_details from the chat keys, which dropped any detail field past cached_tokens and reasoning_tokens that the merge base passed through. The dump now starts from the provider's stored details and only sets the two counts derived from the chat keys * fix(batches): skip sync loggers on the cost poller's own retrieve and honor exclude_none on provider details * fix(batches): keep sync logger events at $0 on poller-owned retrieves, pass unparsed provider counts through, log batch usage in chat shape * fix(batches): store only chat usage extras on batch spend rows Spend rows for batch retrieves copied OpenAI's batch fields (input_tokens, output_tokens and their details) into metadata.additional_usage_values next to the chat counts they duplicate. Build the row from the batch's chat usage, shared with the logging payload through batch_chat_usage.
This commit is contained in:
Mateo Wang
GitHub
parent
594f862bf3
commit
326f17594e
9 changed files with 575 additions and 23 deletions
|
|
@ -202,6 +202,7 @@ from litellm.types.utils import (
|
|||
TextCompletionResponse,
|
||||
TranscriptionResponse,
|
||||
Usage,
|
||||
batch_chat_usage,
|
||||
)
|
||||
from litellm.types.videos.main import VideoObject
|
||||
from litellm.utils import get_base_model_from_metadata, print_verbose
|
||||
|
|
@ -2091,6 +2092,11 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
elif router_model_id is None and "model_id" in hidden_params: # use model_id if not already set
|
||||
router_model_id = hidden_params["model_id"]
|
||||
|
||||
if isinstance(priced_result, LiteLLMBatch) and (
|
||||
not batch_cost_is_final(priced_result) or self._is_batch_poll_without_logging(priced_result)
|
||||
):
|
||||
return 0.0
|
||||
|
||||
# Fallback: extract router_model_id from litellm_params when not available
|
||||
# from the result object. ResponsesAPIResponse objects (used by /v1/responses
|
||||
# streaming) don't carry _hidden_params["model_id"] like ModelResponse does.
|
||||
|
|
@ -3346,6 +3352,10 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
finally:
|
||||
self._restore_correlation_context()
|
||||
|
||||
def _is_batch_poll_without_logging(self, result: object) -> bool:
|
||||
litellm_metadata: Final = (self.litellm_params or {}).get("litellm_metadata") or {}
|
||||
return isinstance(result, LiteLLMBatch) and litellm_metadata.get("batch_ignore_default_logging", False) is True
|
||||
|
||||
def _success_handler_body(
|
||||
self,
|
||||
result: object = None,
|
||||
|
|
@ -3860,11 +3870,7 @@ class Logging(LiteLLMLoggingBaseClass):
|
|||
|
||||
## CALCULATE COST FOR BATCH JOBS
|
||||
if self.call_type == CallTypes.aretrieve_batch.value and isinstance(result, LiteLLMBatch):
|
||||
litellm_params = self.litellm_params or {}
|
||||
litellm_metadata: Final = litellm_params.get("litellm_metadata") or {}
|
||||
if (
|
||||
litellm_metadata.get("batch_ignore_default_logging", False) is True
|
||||
): # polling job will query these frequently, don't spam db logs
|
||||
if self._is_batch_poll_without_logging(result):
|
||||
return
|
||||
|
||||
batch_cost: Final = kwargs.get("batch_cost", None)
|
||||
|
|
@ -7481,7 +7487,8 @@ def get_standard_logging_object_payload(
|
|||
# Extract usage as a plain dict, avoiding Pydantic round-trip
|
||||
raw_usage_dict: Final = StandardLoggingPayloadSetup.get_usage_as_dict(
|
||||
response_obj=None if is_unbilled_non_inference_call(call_type, metadata, response_obj) else response_obj,
|
||||
combined_usage_object=cast(Usage | None, kwargs.get("combined_usage_object")),
|
||||
combined_usage_object=cast(Usage | None, kwargs.get("combined_usage_object"))
|
||||
or (batch_chat_usage(init_response_obj) if isinstance(init_response_obj, LiteLLMBatch) else None),
|
||||
)
|
||||
usage_dict: Final = (
|
||||
{**raw_usage_dict, "output_image_count": len(init_response_obj.data)}
|
||||
|
|
|
|||
|
|
@ -87,6 +87,7 @@ from litellm.types.utils import (
|
|||
AzureSpillover,
|
||||
CallTypes,
|
||||
CostBreakdown,
|
||||
LiteLLMBatch,
|
||||
LlmProviders,
|
||||
StandardLoggingGuardrailInformation,
|
||||
StandardLoggingKubernetesPodRouting,
|
||||
|
|
@ -96,6 +97,7 @@ from litellm.types.utils import (
|
|||
StandardLoggingPayloadErrorInformation,
|
||||
StandardLoggingVectorStoreRequest,
|
||||
VectorStoreSearchResponse,
|
||||
batch_chat_usage,
|
||||
)
|
||||
from litellm.utils import get_end_user_id_for_cost_tracking
|
||||
|
||||
|
|
@ -626,7 +628,12 @@ def get_logging_payload(
|
|||
usage = _extract_usage_for_ocr_call(response_obj, response_obj_dict)
|
||||
elif not is_unbilled_non_inference_call(call_type, metadata, response_obj_dict):
|
||||
# Use response_obj_dict instead of response_obj to avoid calling .get() on Pydantic models
|
||||
_usage: Final = response_obj_dict.get("usage", None) or {}
|
||||
chat_usage_of_batch: Final = batch_chat_usage(response_obj) if isinstance(response_obj, LiteLLMBatch) else None
|
||||
_usage: Final = (
|
||||
chat_usage_of_batch.model_dump()
|
||||
if chat_usage_of_batch is not None
|
||||
else response_obj_dict.get("usage", None) or {}
|
||||
)
|
||||
if isinstance(_usage, litellm.Usage):
|
||||
usage = dict(_usage)
|
||||
elif isinstance(_usage, dict):
|
||||
|
|
|
|||
|
|
@ -20,6 +20,7 @@ from openai._models import BaseModel as OpenAIObject
|
|||
from openai.types.audio.transcription_create_params import (
|
||||
FileTypes as FileTypes,
|
||||
)
|
||||
from openai.types.batch_usage import BatchUsage, InputTokensDetails, OutputTokensDetails
|
||||
from openai.types.chat.chat_completion import ChatCompletion as ChatCompletion
|
||||
from openai.types.completion_usage import (
|
||||
CompletionTokensDetails,
|
||||
|
|
@ -44,6 +45,7 @@ from pydantic import (
|
|||
JsonValue,
|
||||
PrivateAttr,
|
||||
SkipValidation,
|
||||
ValidationError,
|
||||
field_serializer,
|
||||
field_validator,
|
||||
)
|
||||
|
|
@ -4785,6 +4787,112 @@ class LiteLLMFineTuningJob(FineTuningJob):
|
|||
self._hidden_params = kwargs.get("_hidden_params", {})
|
||||
|
||||
|
||||
class _BatchInputTokensDetails(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
cached_tokens: int | None = None
|
||||
|
||||
|
||||
class _BatchOutputTokensDetails(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
reasoning_tokens: int | None = None
|
||||
|
||||
|
||||
class _BatchTokenCounts(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
input_tokens: int
|
||||
output_tokens: int
|
||||
total_tokens: int | None = None
|
||||
input_tokens_details: _BatchInputTokensDetails | None = None
|
||||
output_tokens_details: _BatchOutputTokensDetails | None = None
|
||||
|
||||
|
||||
def _parse_batch_token_counts(value: object) -> _BatchTokenCounts | None:
|
||||
if not isinstance(value, Mapping):
|
||||
return None
|
||||
try:
|
||||
return _BatchTokenCounts.model_validate(value)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def _prompt_details_with_cached_tokens(
|
||||
details: PromptTokensDetailsWrapper | None, cached_tokens: int | None
|
||||
) -> PromptTokensDetailsWrapper | None:
|
||||
if not cached_tokens:
|
||||
return details
|
||||
return (details or PromptTokensDetailsWrapper()).model_copy(update={"cached_tokens": cached_tokens})
|
||||
|
||||
|
||||
def _completion_details_with_reasoning_tokens(
|
||||
details: CompletionTokensDetailsWrapper | None, reasoning_tokens: int | None
|
||||
) -> CompletionTokensDetailsWrapper | None:
|
||||
if not reasoning_tokens:
|
||||
return details
|
||||
return (details or CompletionTokensDetailsWrapper()).model_copy(update={"reasoning_tokens": reasoning_tokens})
|
||||
|
||||
|
||||
def _usage_with_batch_token_counts(usage: Usage, counts: _BatchTokenCounts) -> Usage:
|
||||
return usage.model_copy(
|
||||
update={
|
||||
"prompt_tokens": counts.input_tokens,
|
||||
"completion_tokens": counts.output_tokens,
|
||||
"total_tokens": (
|
||||
counts.input_tokens + counts.output_tokens if counts.total_tokens is None else counts.total_tokens
|
||||
),
|
||||
"prompt_tokens_details": _prompt_details_with_cached_tokens(
|
||||
usage.prompt_tokens_details,
|
||||
counts.input_tokens_details.cached_tokens if counts.input_tokens_details else None,
|
||||
),
|
||||
"completion_tokens_details": _completion_details_with_reasoning_tokens(
|
||||
usage.completion_tokens_details,
|
||||
counts.output_tokens_details.reasoning_tokens if counts.output_tokens_details else None,
|
||||
),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class _ProviderBatchUsageDetails(BaseModel):
|
||||
model_config = ConfigDict(frozen=True, from_attributes=True)
|
||||
|
||||
input_tokens_details: Mapping[str, JsonValue] | None = None
|
||||
output_tokens_details: Mapping[str, JsonValue] | None = None
|
||||
|
||||
|
||||
def _provider_batch_usage_details(usage: Usage) -> _ProviderBatchUsageDetails:
|
||||
try:
|
||||
return _ProviderBatchUsageDetails.model_validate(usage)
|
||||
except ValidationError:
|
||||
return _ProviderBatchUsageDetails()
|
||||
|
||||
|
||||
def _has_unparsed_batch_token_counts(usage: Usage) -> bool:
|
||||
provider_fields: Final = usage.model_extra or {}
|
||||
sent_batch_counts: Final = "input_tokens" in provider_fields or "output_tokens" in provider_fields
|
||||
return sent_batch_counts and _parse_batch_token_counts(provider_fields) is None
|
||||
|
||||
|
||||
def _batch_usage_from_chat_usage(usage: Usage) -> BatchUsage:
|
||||
provider_details: Final = _provider_batch_usage_details(usage)
|
||||
cached_tokens: Final = usage.prompt_tokens_details.cached_tokens if usage.prompt_tokens_details else None
|
||||
reasoning_tokens: Final = (
|
||||
usage.completion_tokens_details.reasoning_tokens if usage.completion_tokens_details else None
|
||||
)
|
||||
return BatchUsage(
|
||||
input_tokens=usage.prompt_tokens,
|
||||
input_tokens_details=InputTokensDetails.model_validate(
|
||||
{**(provider_details.input_tokens_details or {}), "cached_tokens": cached_tokens or 0}
|
||||
),
|
||||
output_tokens=usage.completion_tokens,
|
||||
output_tokens_details=OutputTokensDetails.model_validate(
|
||||
{**(provider_details.output_tokens_details or {}), "reasoning_tokens": reasoning_tokens or 0}
|
||||
),
|
||||
total_tokens=usage.total_tokens,
|
||||
)
|
||||
|
||||
|
||||
class LiteLLMBatch(Batch):
|
||||
_hidden_params: dict = {}
|
||||
|
||||
|
|
@ -4798,6 +4906,40 @@ class LiteLLMBatch(Batch):
|
|||
|
||||
usage: Usage | None = None
|
||||
|
||||
@field_validator("usage", mode="before")
|
||||
@classmethod
|
||||
def _usage_from_batch_token_counts(cls, value: object) -> object:
|
||||
counts: Final = _parse_batch_token_counts(value)
|
||||
return value if counts is None else _usage_with_batch_token_counts(Usage.model_validate(value), counts)
|
||||
|
||||
@field_serializer("usage")
|
||||
def _serialize_usage_with_batch_fields(
|
||||
self, usage: Usage | None, info: FieldSerializationInfo
|
||||
) -> Mapping[str, object] | None:
|
||||
if not isinstance(usage, Usage):
|
||||
return usage
|
||||
chat_fields: Final = usage.model_dump(
|
||||
mode=info.mode,
|
||||
include=info.include,
|
||||
exclude=info.exclude,
|
||||
exclude_unset=info.exclude_unset,
|
||||
exclude_defaults=info.exclude_defaults,
|
||||
exclude_none=info.exclude_none,
|
||||
)
|
||||
if _has_unparsed_batch_token_counts(usage):
|
||||
return chat_fields
|
||||
return {
|
||||
**chat_fields,
|
||||
**_batch_usage_from_chat_usage(usage).model_dump(
|
||||
mode=info.mode,
|
||||
include=info.include,
|
||||
exclude=info.exclude,
|
||||
exclude_unset=info.exclude_unset,
|
||||
exclude_defaults=info.exclude_defaults,
|
||||
exclude_none=info.exclude_none,
|
||||
),
|
||||
}
|
||||
|
||||
def __contains__(self, key) -> bool:
|
||||
# Define custom behavior for the 'in' operator
|
||||
return hasattr(self, key)
|
||||
|
|
@ -4818,6 +4960,15 @@ class LiteLLMBatch(Batch):
|
|||
return self.dict()
|
||||
|
||||
|
||||
_PROVIDER_ONLY_BATCH_USAGE_FIELDS: Final = frozenset(BatchUsage.model_fields) - frozenset(Usage.model_fields)
|
||||
|
||||
|
||||
def batch_chat_usage(batch: LiteLLMBatch) -> Usage | None:
|
||||
if batch.usage is None:
|
||||
return None
|
||||
return Usage.model_validate(batch.usage.model_dump(exclude=set(_PROVIDER_ONLY_BATCH_USAGE_FIELDS)))
|
||||
|
||||
|
||||
class LiteLLMRealtimeStreamLoggingObject(LiteLLMPydanticObjectBase):
|
||||
# Events are already well-formed provider dicts. Validating them against the
|
||||
# OpenAIRealtimeEvents union makes Pydantic try every member per event, which
|
||||
|
|
|
|||
|
|
@ -15,7 +15,7 @@ from __future__ import annotations
|
|||
from dataclasses import dataclass
|
||||
from typing import Final, Literal
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
from pydantic import BaseModel, ConfigDict, Field
|
||||
|
||||
from e2e_metadata import step
|
||||
from proxy_client import ProxyClient
|
||||
|
|
@ -77,6 +77,16 @@ class BatchErrorList(BaseModel):
|
|||
data: list[BatchErrorItem] = []
|
||||
|
||||
|
||||
class BatchUsage(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
prompt_tokens: int | None = None
|
||||
completion_tokens: int | None = None
|
||||
total_tokens: int | None = None
|
||||
input_tokens: int | None = None
|
||||
output_tokens: int | None = None
|
||||
|
||||
|
||||
class BatchObject(BaseModel):
|
||||
id: str
|
||||
object: str | None = None
|
||||
|
|
@ -90,6 +100,7 @@ class BatchObject(BaseModel):
|
|||
completion_window: str | None = None
|
||||
created_at: int | None = None
|
||||
model: str | None = None
|
||||
usage: BatchUsage | None = None
|
||||
|
||||
|
||||
class BatchList(BaseModel):
|
||||
|
|
|
|||
|
|
@ -1893,6 +1893,13 @@ class TestBatchTerminalState:
|
|||
f"listed-completed marker retrieved as {fetched.status!r}"
|
||||
)
|
||||
assert fetched.output_file_id, "completed batch has no output_file_id"
|
||||
usage: Final = fetched.usage
|
||||
assert usage is not None and usage.input_tokens and usage.output_tokens, (
|
||||
f"completed batch has no batch token counts: {usage!r}"
|
||||
)
|
||||
assert (usage.prompt_tokens, usage.completion_tokens) == (usage.input_tokens, usage.output_tokens), (
|
||||
f"chat token counts disagree with the batch token counts: {usage!r}"
|
||||
)
|
||||
|
||||
downloaded = client.proxy.transport.download(
|
||||
f"/v1/files/{fetched.output_file_id}/content",
|
||||
|
|
|
|||
|
|
@ -1629,22 +1629,33 @@ class TestRetrieveBatchPricesOnlyFinalBatches:
|
|||
litellm_call_id="batch-call-2",
|
||||
function_id="f",
|
||||
)
|
||||
obj.custom_llm_provider = "openai"
|
||||
obj.update_environment_variables(
|
||||
model="gpt-5.6-luna", user=None, optional_params={}, litellm_params={}, custom_llm_provider="openai"
|
||||
)
|
||||
return obj
|
||||
|
||||
@staticmethod
|
||||
def _batch(status: str, output_file_id: str | None):
|
||||
from litellm.types.utils import LiteLLMBatch
|
||||
|
||||
return LiteLLMBatch(
|
||||
id="batch_6a9c99e185588190877d391f8b9d7f8a",
|
||||
completion_window="24h",
|
||||
created_at=1,
|
||||
endpoint="/v1/chat/completions",
|
||||
input_file_id="file-in",
|
||||
object="batch",
|
||||
status="validating",
|
||||
output_file_id=output_file_id,
|
||||
return LiteLLMBatch.model_validate(
|
||||
{
|
||||
"id": "batch_6a9c99e185588190877d391f8b9d7f8a",
|
||||
"completion_window": "24h",
|
||||
"created_at": 1,
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"input_file_id": "file-in",
|
||||
"object": "batch",
|
||||
"status": "validating",
|
||||
"output_file_id": output_file_id,
|
||||
"usage": {
|
||||
"input_tokens": 29,
|
||||
"input_tokens_details": {"cached_tokens": 0},
|
||||
"output_tokens": 48,
|
||||
"output_tokens_details": {"reasoning_tokens": 0},
|
||||
"total_tokens": 77,
|
||||
},
|
||||
}
|
||||
).model_copy(update={"status": status})
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -1658,11 +1669,12 @@ class TestRetrieveBatchPricesOnlyFinalBatches:
|
|||
handle_completed_batch = AsyncMock()
|
||||
monkeypatch.setattr(logging_module, "handle_completed_batch", handle_completed_batch)
|
||||
batch = self._batch(status, output_file_id)
|
||||
logging_obj = self._logging_obj()
|
||||
|
||||
await self._logging_obj()._async_success_handler_body(result=batch, start_time=None, end_time=None)
|
||||
await logging_obj._async_success_handler_body(result=batch, start_time=None, end_time=None)
|
||||
|
||||
handle_completed_batch.assert_not_awaited()
|
||||
assert "response_cost" not in batch._hidden_params
|
||||
assert logging_obj.model_call_details["response_cost"] == 0.0
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_completed_batch_with_output_is_priced(self, monkeypatch) -> None:
|
||||
|
|
@ -1689,6 +1701,80 @@ class TestRetrieveBatchPricesOnlyFinalBatches:
|
|||
assert batch.usage is not None
|
||||
assert batch.usage.total_tokens == 35
|
||||
|
||||
@pytest.mark.parametrize("batch_ignore_default_logging", [True, False])
|
||||
def test_sync_logger_gets_a_free_event_for_the_cost_poller_retrieve(self, batch_ignore_default_logging: bool) -> None:
|
||||
sync_events: list[object] = []
|
||||
|
||||
class SyncOnlyLogger(CustomLogger):
|
||||
def log_success_event(
|
||||
self,
|
||||
kwargs: dict[str, object],
|
||||
response_obj: object,
|
||||
start_time: dt_object,
|
||||
end_time: dt_object,
|
||||
) -> None:
|
||||
sync_events.append(kwargs["response_cost"])
|
||||
|
||||
logging_obj = LitellmLogging(
|
||||
model="gpt-5.6-luna",
|
||||
messages=[{"role": "user", "content": "Hey"}],
|
||||
stream=False,
|
||||
call_type="aretrieve_batch",
|
||||
start_time=time.time(),
|
||||
litellm_call_id="batch-call-3",
|
||||
function_id="f",
|
||||
dynamic_success_callbacks=[SyncOnlyLogger()],
|
||||
)
|
||||
logging_obj.update_environment_variables(
|
||||
model="gpt-5.6-luna",
|
||||
user=None,
|
||||
optional_params={},
|
||||
litellm_params={"litellm_metadata": {"batch_ignore_default_logging": batch_ignore_default_logging}},
|
||||
custom_llm_provider="openai",
|
||||
)
|
||||
|
||||
logging_obj.success_handler(
|
||||
result=self._batch("completed", "file-out"),
|
||||
start_time=dt_object.now(),
|
||||
end_time=dt_object.now(),
|
||||
for_async_request=True,
|
||||
)
|
||||
|
||||
assert len(sync_events) == 1
|
||||
assert isinstance(sync_events[0], float)
|
||||
assert (sync_events[0] == 0.0) is batch_ignore_default_logging
|
||||
|
||||
def test_logged_usage_object_keeps_the_batch_chat_usage(self) -> None:
|
||||
from litellm.litellm_core_utils.litellm_logging import get_standard_logging_object_payload
|
||||
from litellm.types.utils import Usage
|
||||
|
||||
aggregate: Final = Usage(
|
||||
prompt_tokens=29, completion_tokens=48, total_tokens=77, reasoning_tokens=3, cache_creation_input_tokens=300
|
||||
)
|
||||
billed: Final = self._batch("completed", "file-out").model_copy(update={"usage": aggregate})
|
||||
provider_read: Final = self._batch("completed", "file-out")
|
||||
now: Final = dt_object(2026, 1, 1, 12, 0, 0)
|
||||
|
||||
usage_objects: Final = tuple(
|
||||
payload["metadata"]["usage_object"]
|
||||
for payload in (
|
||||
get_standard_logging_object_payload(
|
||||
kwargs={"litellm_params": {}},
|
||||
init_response_obj=batch,
|
||||
start_time=now,
|
||||
end_time=now,
|
||||
logging_obj=self._logging_obj(),
|
||||
status="success",
|
||||
)
|
||||
for batch in (billed, provider_read)
|
||||
)
|
||||
if payload is not None
|
||||
)
|
||||
|
||||
assert usage_objects[0] == aggregate.model_dump()
|
||||
assert usage_objects[1]["prompt_tokens"] == 29
|
||||
assert "input_tokens" not in usage_objects[1]
|
||||
|
||||
|
||||
class TestAnthropicPassthroughCustomPricing:
|
||||
"""Verify the Anthropic pass-through handler forwards custom pricing."""
|
||||
|
|
|
|||
|
|
@ -80,10 +80,11 @@ def _batch_cost_result(
|
|||
"""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
|
||||
from litellm.types.utils import Usage
|
||||
|
||||
return BatchCostUsageResult(
|
||||
cost=cost,
|
||||
usage=usage,
|
||||
usage=Usage(**usage),
|
||||
models=models,
|
||||
successful_requests=successful_requests,
|
||||
failed_requests=failed_requests,
|
||||
|
|
@ -3261,7 +3262,7 @@ class TestMultiPodBatchCostClaim:
|
|||
patch(
|
||||
"litellm.batches.batch_utils.calculate_batch_cost_and_usage",
|
||||
new_callable=AsyncMock,
|
||||
return_value=_batch_cost_result(0.01, {"prompt_tokens": 10, "completion_tokens": 5}, ["gpt-4"]),
|
||||
return_value=_batch_cost_result(0.01, {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, ["gpt-4"]),
|
||||
),
|
||||
patch(
|
||||
"litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider",
|
||||
|
|
@ -3676,7 +3677,17 @@ class TestMultiPodBatchCostClaim:
|
|||
|
||||
assert journal == ["fetch", "claim", "bill"]
|
||||
assert row.status == "complete"
|
||||
assert json.loads(row.file_object)["usage"] == {"prompt_tokens": 10, "completion_tokens": 5}
|
||||
assert json.loads(row.file_object)["usage"] == {
|
||||
"prompt_tokens": 10,
|
||||
"completion_tokens": 5,
|
||||
"total_tokens": 15,
|
||||
"prompt_tokens_details": None,
|
||||
"completion_tokens_details": None,
|
||||
"input_tokens": 10,
|
||||
"input_tokens_details": {"cached_tokens": 0},
|
||||
"output_tokens": 5,
|
||||
"output_tokens_details": {"reasoning_tokens": 0},
|
||||
}
|
||||
logging_obj.async_success_handler.assert_awaited_once()
|
||||
assert row.batch_processed is True
|
||||
|
||||
|
|
|
|||
|
|
@ -69,6 +69,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import (
|
|||
from litellm.types.router import GenericLiteLLMParams
|
||||
from litellm.types.utils import (
|
||||
BACKGROUND_RESPONSE_COST_POLL_CALL_ORIGIN,
|
||||
LiteLLMBatch,
|
||||
StandardLoggingHiddenParams,
|
||||
StandardLoggingMetadata,
|
||||
StandardLoggingModelInformation,
|
||||
|
|
@ -222,6 +223,64 @@ def test_get_logging_payload_maps_openai_cached_tokens_to_cache_read_input_token
|
|||
assert additional_usage_values["prompt_tokens_details"]["cached_tokens"] == 123
|
||||
|
||||
|
||||
_PROVIDER_READ_BATCH: Final = LiteLLMBatch.model_validate(
|
||||
{
|
||||
"id": "batch_spend_row",
|
||||
"completion_window": "24h",
|
||||
"created_at": 1,
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"input_file_id": "file-in",
|
||||
"object": "batch",
|
||||
"status": "completed",
|
||||
"output_file_id": "file-out",
|
||||
"usage": {
|
||||
"input_tokens": 29,
|
||||
"input_tokens_details": {"cached_tokens": 8},
|
||||
"output_tokens": 46,
|
||||
"output_tokens_details": {"reasoning_tokens": 0},
|
||||
"total_tokens": 75,
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"batch",
|
||||
[
|
||||
_PROVIDER_READ_BATCH,
|
||||
_PROVIDER_READ_BATCH.model_copy(
|
||||
update={
|
||||
"usage": litellm.Usage(
|
||||
prompt_tokens=29,
|
||||
completion_tokens=46,
|
||||
total_tokens=75,
|
||||
prompt_tokens_details={"cached_tokens": 8},
|
||||
)
|
||||
}
|
||||
),
|
||||
],
|
||||
ids=["provider_read", "poller_billed"],
|
||||
)
|
||||
def test_batch_spend_row_stores_the_same_usage_extras_as_a_chat_call_with_its_counts(batch: LiteLLMBatch) -> None:
|
||||
payload: Final = get_logging_payload(
|
||||
kwargs={
|
||||
"model": "gpt-4o-mini",
|
||||
"call_type": "aretrieve_batch",
|
||||
"litellm_params": {"metadata": {"user_api_key": "test-key"}},
|
||||
},
|
||||
response_obj=batch,
|
||||
start_time=datetime.datetime.now(timezone.utc),
|
||||
end_time=datetime.datetime.now(timezone.utc),
|
||||
)
|
||||
|
||||
assert (payload["prompt_tokens"], payload["completion_tokens"], payload["total_tokens"]) == (29, 46, 75)
|
||||
assert json.loads(payload["metadata"])["additional_usage_values"] == _get_additional_usage_values_for_usage(
|
||||
litellm.Usage(
|
||||
prompt_tokens=29, completion_tokens=46, total_tokens=75, prompt_tokens_details={"cached_tokens": 8}
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class _HashingCache(litellm.Cache):
|
||||
def __init__(self) -> None:
|
||||
pass
|
||||
|
|
|
|||
|
|
@ -25,11 +25,13 @@ from litellm.types.utils import (
|
|||
HiddenParams,
|
||||
ImageObject,
|
||||
ImageResponse,
|
||||
LiteLLMBatch,
|
||||
ModelResponse,
|
||||
ModelResponseStream,
|
||||
TextCompletionResponse,
|
||||
all_litellm_params,
|
||||
is_decisions_model_mode,
|
||||
Usage,
|
||||
text_tokens_without_nested_reasoning,
|
||||
)
|
||||
|
||||
|
|
@ -1007,3 +1009,214 @@ def test_text_tokens_without_nested_reasoning_clamps(
|
|||
)
|
||||
== expected_text_tokens
|
||||
)
|
||||
|
||||
|
||||
_BATCH_FIELDS: Final = {
|
||||
"id": "batch_abc",
|
||||
"object": "batch",
|
||||
"endpoint": "/v1/responses",
|
||||
"input_file_id": "file-abc",
|
||||
"completion_window": "24h",
|
||||
"status": "completed",
|
||||
"created_at": 1,
|
||||
}
|
||||
_PROVIDER_BATCH_USAGE: Final = {
|
||||
"input_tokens": 21,
|
||||
"input_tokens_details": {"cached_tokens": 5},
|
||||
"output_tokens": 543,
|
||||
"output_tokens_details": {"reasoning_tokens": 448},
|
||||
"total_tokens": 564,
|
||||
}
|
||||
|
||||
|
||||
def _batch_token_counts(batch: LiteLLMBatch) -> tuple[int, int, int, int | None, int | None]:
|
||||
assert batch.usage is not None
|
||||
assert batch.usage.prompt_tokens_details is not None
|
||||
assert batch.usage.completion_tokens_details is not None
|
||||
return (
|
||||
batch.usage.prompt_tokens,
|
||||
batch.usage.completion_tokens,
|
||||
batch.usage.total_tokens,
|
||||
batch.usage.prompt_tokens_details.cached_tokens,
|
||||
batch.usage.completion_tokens_details.reasoning_tokens,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"usage",
|
||||
(
|
||||
pytest.param(_PROVIDER_BATCH_USAGE, id="provider_batch_usage"),
|
||||
pytest.param(
|
||||
{"prompt_tokens": 0, "completion_tokens": 0, **_PROVIDER_BATCH_USAGE},
|
||||
id="stored_row_with_zeroed_chat_fields",
|
||||
),
|
||||
),
|
||||
)
|
||||
def test_batch_usage_reads_token_counts_from_the_provider_batch_usage(usage: dict[str, object]) -> None:
|
||||
batch: Final = LiteLLMBatch.model_validate({**_BATCH_FIELDS, "usage": usage})
|
||||
|
||||
assert _batch_token_counts(batch) == (21, 543, 564, 5, 448)
|
||||
assert batch.usage is not None
|
||||
assert (batch.usage.input_tokens, batch.usage.output_tokens) == (21, 543)
|
||||
|
||||
|
||||
def test_batch_usage_keeps_chat_shaped_usage_as_is() -> None:
|
||||
batch: Final = LiteLLMBatch.model_validate(
|
||||
{**_BATCH_FIELDS, "usage": {"prompt_tokens": 28, "completion_tokens": 8, "total_tokens": 36}}
|
||||
)
|
||||
|
||||
assert batch.usage is not None
|
||||
assert (batch.usage.prompt_tokens, batch.usage.completion_tokens, batch.usage.total_tokens) == (28, 8, 36)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"usage",
|
||||
(
|
||||
pytest.param({"input_tokens": 10, "output_tokens": 5}, id="no_details_or_total"),
|
||||
pytest.param(
|
||||
{"input_tokens": 10, "input_tokens_details": None, "output_tokens": 5, "output_tokens_details": None},
|
||||
id="null_details",
|
||||
),
|
||||
pytest.param(
|
||||
{
|
||||
"input_tokens": 10,
|
||||
"input_tokens_details": {"cached_tokens": None},
|
||||
"output_tokens": 5,
|
||||
"output_tokens_details": {"reasoning_tokens": None},
|
||||
"total_tokens": 15,
|
||||
},
|
||||
id="null_detail_counts",
|
||||
),
|
||||
),
|
||||
)
|
||||
def test_batch_usage_reads_token_counts_from_partial_provider_batch_usage(usage: dict[str, object]) -> None:
|
||||
batch: Final = LiteLLMBatch.model_validate({**_BATCH_FIELDS, "usage": usage})
|
||||
dumped_usage: Final = json.loads(batch.model_dump_json())["usage"]
|
||||
|
||||
assert batch.usage is not None
|
||||
assert (batch.usage.prompt_tokens, batch.usage.completion_tokens, batch.usage.total_tokens) == (10, 5, 15)
|
||||
assert (dumped_usage["input_tokens"], dumped_usage["output_tokens"], dumped_usage["total_tokens"]) == (10, 5, 15)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"usage",
|
||||
(
|
||||
pytest.param(_PROVIDER_BATCH_USAGE, id="provider_batch_usage"),
|
||||
pytest.param(
|
||||
Usage(
|
||||
prompt_tokens=21,
|
||||
completion_tokens=543,
|
||||
total_tokens=564,
|
||||
prompt_tokens_details={"cached_tokens": 5},
|
||||
completion_tokens_details={"reasoning_tokens": 448},
|
||||
),
|
||||
id="billed_chat_usage",
|
||||
),
|
||||
),
|
||||
)
|
||||
def test_batch_dump_carries_the_chat_and_provider_batch_usage_fields(usage: dict[str, object] | Usage) -> None:
|
||||
batch: Final = LiteLLMBatch.model_validate({**_BATCH_FIELDS, "usage": usage})
|
||||
dumped: Final = batch.model_dump_json()
|
||||
dumped_usage: Final = json.loads(dumped)["usage"]
|
||||
|
||||
assert {key: dumped_usage[key] for key in _PROVIDER_BATCH_USAGE} == _PROVIDER_BATCH_USAGE
|
||||
assert (dumped_usage["prompt_tokens"], dumped_usage["completion_tokens"]) == (21, 543)
|
||||
assert batch.model_dump()["usage"] == dumped_usage
|
||||
assert _batch_token_counts(LiteLLMBatch.model_validate_json(dumped)) == (21, 543, 564, 5, 448)
|
||||
|
||||
|
||||
def test_batch_dump_keeps_cache_write_tokens_from_billed_usage() -> None:
|
||||
batch: Final = LiteLLMBatch.model_validate(_BATCH_FIELDS).model_copy(
|
||||
update={"usage": Usage(prompt_tokens=30, completion_tokens=4, total_tokens=34, cache_creation_input_tokens=7)}
|
||||
)
|
||||
dumped: Final = batch.model_dump_json()
|
||||
dumped_usage: Final = json.loads(dumped)["usage"]
|
||||
|
||||
assert (dumped_usage["prompt_tokens"], dumped_usage["input_tokens"]) == (30, 30)
|
||||
assert (dumped_usage["completion_tokens"], dumped_usage["output_tokens"]) == (4, 4)
|
||||
assert dumped_usage["prompt_tokens_details"]["cache_write_tokens"] == 7
|
||||
assert None not in batch.model_dump(exclude_none=True)["usage"]["prompt_tokens_details"].values()
|
||||
assert json.loads(LiteLLMBatch.model_validate_json(dumped).model_dump_json())["usage"] == dumped_usage
|
||||
|
||||
|
||||
def test_batch_json_keeps_usage_that_skipped_validation() -> None:
|
||||
batch: Final = LiteLLMBatch.model_validate(_BATCH_FIELDS).model_copy(update={"usage": {"prompt_tokens": 3}})
|
||||
|
||||
assert json.loads(batch.model_dump_json(warnings=False))["usage"] == {"prompt_tokens": 3}
|
||||
|
||||
|
||||
def test_batch_dump_applies_nested_include_and_exclude_to_both_usage_shapes() -> None:
|
||||
batch: Final = LiteLLMBatch.model_validate({**_BATCH_FIELDS, "usage": _PROVIDER_BATCH_USAGE})
|
||||
|
||||
excluded: Final = batch.model_dump(exclude={"usage": {"prompt_tokens_details", "input_tokens_details"}})["usage"]
|
||||
included: Final = batch.model_dump(include={"usage": {"prompt_tokens", "input_tokens"}})["usage"]
|
||||
|
||||
assert {"prompt_tokens_details", "input_tokens_details"}.isdisjoint(excluded)
|
||||
assert (excluded["prompt_tokens"], excluded["input_tokens"]) == (21, 21)
|
||||
assert included == {"prompt_tokens": 21, "input_tokens": 21}
|
||||
|
||||
|
||||
def test_batch_dump_applies_exclude_defaults_to_the_chat_usage_keys() -> None:
|
||||
batch: Final = LiteLLMBatch.model_validate(
|
||||
{**_BATCH_FIELDS, "usage": Usage(prompt_tokens=10, completion_tokens=5, total_tokens=15)}
|
||||
)
|
||||
|
||||
python_usage: Final = batch.model_dump(exclude_defaults=True)["usage"]
|
||||
json_usage: Final = json.loads(batch.model_dump_json(exclude_defaults=True, warnings=False))["usage"]
|
||||
|
||||
assert python_usage == json_usage
|
||||
assert {"prompt_tokens_details", "completion_tokens_details"}.isdisjoint(python_usage)
|
||||
assert (python_usage["prompt_tokens"], python_usage["input_tokens"]) == (10, 10)
|
||||
|
||||
|
||||
def test_batch_dump_forwards_every_provider_usage_detail_field() -> None:
|
||||
batch: Final = LiteLLMBatch.model_validate(
|
||||
{
|
||||
**_BATCH_FIELDS,
|
||||
"usage": {
|
||||
**_PROVIDER_BATCH_USAGE,
|
||||
"input_tokens_details": {"cached_tokens": 5, "audio_tokens": 3},
|
||||
"output_tokens_details": {"reasoning_tokens": 448, "audio_tokens": 2},
|
||||
},
|
||||
}
|
||||
)
|
||||
|
||||
python_usage: Final = batch.model_dump()["usage"]
|
||||
json_usage: Final = json.loads(batch.model_dump_json(warnings=False))["usage"]
|
||||
|
||||
assert python_usage == json_usage
|
||||
assert python_usage["input_tokens_details"] == {"cached_tokens": 5, "audio_tokens": 3}
|
||||
assert python_usage["output_tokens_details"] == {"reasoning_tokens": 448, "audio_tokens": 2}
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"usage",
|
||||
[
|
||||
{"input_tokens": 29, "output_tokens": None, "total_tokens": 29},
|
||||
{"input_tokens": 29, "input_tokens_details": {"cached_tokens": 8}, "total_tokens": 29},
|
||||
],
|
||||
)
|
||||
def test_batch_dump_serves_unparsed_provider_usage_unchanged(usage: dict[str, object]) -> None:
|
||||
batch: Final = LiteLLMBatch.model_validate({**_BATCH_FIELDS, "usage": usage})
|
||||
|
||||
python_usage: Final = batch.model_dump()["usage"]
|
||||
json_usage: Final = json.loads(batch.model_dump_json(warnings=False))["usage"]
|
||||
|
||||
assert python_usage == json_usage
|
||||
assert {key: python_usage[key] for key in usage} == usage
|
||||
|
||||
|
||||
def test_batch_dump_applies_exclude_none_to_provider_usage_detail_fields() -> None:
|
||||
batch: Final = LiteLLMBatch.model_validate(
|
||||
{
|
||||
**_BATCH_FIELDS,
|
||||
"usage": {**_PROVIDER_BATCH_USAGE, "input_tokens_details": {"cached_tokens": 5, "audio_tokens": None}},
|
||||
}
|
||||
)
|
||||
|
||||
python_usage: Final = batch.model_dump(exclude_none=True)["usage"]
|
||||
json_usage: Final = json.loads(batch.model_dump_json(exclude_none=True, warnings=False))["usage"]
|
||||
|
||||
assert python_usage == json_usage
|
||||
assert python_usage["input_tokens_details"] == {"cached_tokens": 5}
|
||||
assert (python_usage["prompt_tokens"], python_usage["input_tokens"]) == (21, 21)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue