fix(batches): read OpenAI batch usage into the batch token counts

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)

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)