From 64eb016b3d149e17852b615dc27fa91b32adde2d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 20 Oct 2026 18:52:22 -0700 Subject: [PATCH] fix(batches): read OpenAI batch usage into the batch token counts --- litellm/types/utils.py | 33 +++++++++++++++++ tests/unit/types/test_types_utils.py | 54 ++++++++++++++++++++++++++++ 2 files changed, 87 insertions(+) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 5939a350970..2a4c04e3c01 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -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 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, ) @@ -4720,6 +4722,31 @@ class LiteLLMFineTuningJob(FineTuningJob): self._hidden_params = kwargs.get("_hidden_params", {}) +def _parse_batch_usage(value: object) -> BatchUsage | None: + if isinstance(value, BatchUsage): + return value + if not isinstance(value, Mapping): + return None + try: + return BatchUsage.model_validate(value) + except ValidationError: + return None + + +def _chat_usage_from_batch_usage(batch_usage: BatchUsage) -> Usage: + return Usage( + prompt_tokens=batch_usage.input_tokens, + completion_tokens=batch_usage.output_tokens, + total_tokens=batch_usage.total_tokens, + prompt_tokens_details={"cached_tokens": batch_usage.input_tokens_details.cached_tokens}, + completion_tokens_details={"reasoning_tokens": batch_usage.output_tokens_details.reasoning_tokens}, + input_tokens=batch_usage.input_tokens, + input_tokens_details=batch_usage.input_tokens_details.model_dump(), + output_tokens=batch_usage.output_tokens, + output_tokens_details=batch_usage.output_tokens_details.model_dump(), + ) + + class LiteLLMBatch(Batch): _hidden_params: dict = {} @@ -4733,6 +4760,12 @@ class LiteLLMBatch(Batch): usage: Usage | None = None + @field_validator("usage", mode="before") + @classmethod + def _usage_from_batch_usage(cls, value: object) -> object: + batch_usage: Final = _parse_batch_usage(value) + return value if batch_usage is None else _chat_usage_from_batch_usage(batch_usage) + def __contains__(self, key) -> bool: # Define custom behavior for the 'in' operator return hasattr(self, key) diff --git a/tests/unit/types/test_types_utils.py b/tests/unit/types/test_types_utils.py index d85e48540f9..9095f0721f4 100644 --- a/tests/unit/types/test_types_utils.py +++ b/tests/unit/types/test_types_utils.py @@ -25,6 +25,7 @@ from litellm.types.utils import ( HiddenParams, ImageObject, ImageResponse, + LiteLLMBatch, ModelResponse, ModelResponseStream, TextCompletionResponse, @@ -1007,3 +1008,56 @@ 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, +} + + +@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.usage is not None + assert batch.usage.prompt_tokens_details is not None + assert batch.usage.completion_tokens_details is not None + assert ( + 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, + ) == (21, 543, 564, 5, 448) + 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)