fix(rerank): preserve token usage in spend logs (#49308)

litellm/litellm_core_utils/litellm_logging.py

@@ -165,7 +165,7 @@ from litellm.types.llms.openai import (
 )
 from litellm.types.mcp import MCPPostCallResponseObject
 from litellm.types.prompts.init_prompts import PromptSpec
-from litellm.types.rerank import RerankResponse
+from litellm.types.rerank import RerankResponse, get_rerank_usage
 from litellm.types.router import RouterRateLimitError, SilentExperimentFlag
 from litellm.types.utils import (
     DEPLOYMENT_SCOPED_PRICING_FIELDS,
@@ -355,6 +355,7 @@ def _get_provider_request_id(original_exception: Exception) -> str | None:
 # Cache custom pricing keys as frozenset for O(1) lookups instead of looping through 49 keys
 _CUSTOM_PRICING_KEYS: Final[frozenset[str]] = frozenset(CustomPricingLiteLLMParams.model_fields.keys())
 _METADATA_MAPPING: Final = TypeAdapter(Mapping[str, object])
+_RESPONSE_MAPPING: Final = TypeAdapter(Mapping[str, object])
 _MODEL_INFO_CUSTOM_PRICING_KEYS: Final[frozenset[str]] = _CUSTOM_PRICING_KEYS | DEPLOYMENT_SCOPED_PRICING_FIELDS
 _UNSERIALIZABLE_METADATA_KEYS: Final[frozenset[str]] = frozenset(
     ("user_api_key_auth", "user_api_key_budget_reservation")
@@ -6817,7 +6818,11 @@ class StandardLoggingPayloadSetup:
                 total_tokens=0,
             )
 
-        usage: Final = response_obj.get("usage", None) or {}
+        usage: Final = response_obj.get("usage", None)
+        if usage is None:
+            rerank_usage: Final = get_rerank_usage(_RESPONSE_MAPPING.validate_python(response_obj))
+            if rerank_usage is not None:
+                return Usage(**rerank_usage)
         if usage is None or (not isinstance(usage, dict) and not isinstance(usage, Usage)):
             return Usage(
                 prompt_tokens=0,
@@ -6853,7 +6858,7 @@ class StandardLoggingPayloadSetup:
             return _empty
         _raw: Final = response_obj.get("usage", None)
         if _raw is None:
-            return _empty
+            return dict(get_rerank_usage(_RESPONSE_MAPPING.validate_python(response_obj)) or _empty)
         if isinstance(_raw, ResponseAPIUsage):
             return ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(_raw).model_dump()
         if isinstance(_raw, dict):

litellm/llms/bedrock/rerank/transformation.py

@@ -6,6 +6,8 @@ Why separate file? Make it easy to see how transformation works
 
 from typing import Final
 
+from pydantic import TypeAdapter
+
 from litellm._uuid import uuid
 from litellm.types.llms.bedrock import (
     BedrockRerankBedrockRerankingConfiguration,
@@ -90,7 +92,9 @@ class BedrockRerankConfig:
         example input:
         {"results":[{"index":0,"relevanceScore":0.6847912669181824},{"index":1,"relevanceScore":0.5980774760246277}]}
         """
-        _billed_units = RerankBilledUnits(**response.get("usage", {"search_units": 1}))  # by default 1 search unit
+        _billed_units: Final = TypeAdapter(RerankBilledUnits).validate_python(
+            response.get("usage", {"search_units": 1})
+        )
         _tokens: Final = RerankTokens(**response.get("usage", {}))
         rerank_meta: Final = RerankResponseMeta(billed_units=_billed_units, tokens=_tokens)
 

litellm/llms/hosted_vllm/rerank/transformation.py

@@ -244,5 +244,7 @@ class HostedVLLMRerankConfig(BaseRerankConfig):
         return RerankResponse(
             id=response.get("id") or str(uuid.uuid4()),
             results=rerank_results,
-            meta=rerank_meta,
+            meta=TypeAdapter(RerankResponseMeta).validate_python(response["meta"])
+            if "meta" in response and response["meta"] is not None
+            else rerank_meta,
         )

litellm/llms/together_ai/rerank/transformation.py

@@ -6,6 +6,8 @@ Why separate file? Make it easy to see how transformation works
 
 from typing import Final
 
+from pydantic import TypeAdapter
+
 from litellm._uuid import uuid
 from litellm.types.rerank import (
     RerankBilledUnits,
@@ -19,7 +21,7 @@ from litellm.types.rerank import (
 
 class TogetherAIRerankConfig:
     def _transform_response(self, response: dict) -> RerankResponse:
-        _billed_units: Final = RerankBilledUnits(**response.get("usage", {}))
+        _billed_units: Final = TypeAdapter(RerankBilledUnits).validate_python(response.get("usage", {}))
         _tokens: Final = RerankTokens(**response.get("usage", {}))
         rerank_meta: Final = RerankResponseMeta(billed_units=_billed_units, tokens=_tokens)
 

litellm/types/rerank.py

@@ -4,10 +4,10 @@ https://docs.cohere.com/reference/rerank
 
 """
 
-from collections.abc import Sequence
-from typing import Literal, overload
+from collections.abc import Mapping, Sequence
+from typing import Final, Literal, overload
 
-from pydantic import ConfigDict, PrivateAttr
+from pydantic import ConfigDict, PrivateAttr, TypeAdapter, ValidationError
 from typing_extensions import ReadOnly, Required, TypedDict
 
 from litellm.types.llms.base import LiteLLMBaseModel
@@ -57,6 +57,8 @@ class OptionalRerankParams(TypedDict, total=False):
 class RerankBilledUnits(TypedDict, total=False):
     search_units: int | None
     total_tokens: int | None
+    input_tokens: ReadOnly[int | None]
+    output_tokens: ReadOnly[int | None]
 
 
 class RerankTokens(TypedDict, total=False):
@@ -65,11 +67,47 @@ class RerankTokens(TypedDict, total=False):
 
 
 class RerankResponseMeta(TypedDict, total=False):
-    api_version: dict | None
+    api_version: dict[str, object] | None
     billed_units: RerankBilledUnits | None
     tokens: RerankTokens | None
 
 
+class RerankUsage(TypedDict):
+    prompt_tokens: ReadOnly[int]
+    completion_tokens: ReadOnly[int]
+    total_tokens: ReadOnly[int]
+
+
+_RERANK_META: Final = TypeAdapter(RerankResponseMeta)
+
+
+def get_rerank_usage(response: Mapping[str, object]) -> RerankUsage | None:
+    if "results" not in response or response.get("meta") is None:
+        return None
+    try:
+        meta: Final = _RERANK_META.validate_python(response.get("meta"))
+    except ValidationError:
+        return None
+    tokens: Final = meta.get("tokens") or RerankTokens()
+    billed: Final = meta.get("billed_units") or RerankBilledUnits()
+    input_tokens: Final = next(
+        (count for count in (tokens.get("input_tokens"), billed.get("input_tokens")) if count is not None), None
+    )
+    output_tokens: Final = next(
+        (count for count in (tokens.get("output_tokens"), billed.get("output_tokens")) if count is not None), None
+    )
+    total_tokens: Final = billed.get("total_tokens")
+    if input_tokens is None and output_tokens is None and total_tokens is None:
+        return None
+    completion_tokens: Final = output_tokens or 0
+    prompt_tokens: Final = input_tokens if input_tokens is not None else max((total_tokens or 0) - completion_tokens, 0)
+    return RerankUsage(
+        prompt_tokens=prompt_tokens,
+        completion_tokens=completion_tokens,
+        total_tokens=total_tokens if total_tokens is not None else prompt_tokens + completion_tokens,
+    )
+
+
 class RerankResponseDocument(TypedDict):
     text: str
 

tests/unit/litellm_core_utils/test_litellm_logging.py

@@ -59,6 +59,7 @@ from litellm.proxy.hooks.max_iterations_limiter import PROXY_MaxIterationsHandle
 from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3
 from litellm.proxy.utils import InternalUsageCache
 from litellm.types.llms.openai import ResponseAPIUsage, ResponseCompletedEvent, ResponsesAPIResponse
+from litellm.types.rerank import RerankResponse
 from litellm.types.utils import (
     CallTypes,
     CostBreakdown,
@@ -4514,6 +4515,105 @@ def test_logging_preserves_missing_usage_without_accepting_request_metadata(
     assert (payload["prompt_tokens"], payload["completion_tokens"], payload["total_tokens"]) == (0, 0, 0)
 
 
+@pytest.mark.parametrize(
+    ("meta", "expected"),
+    [
+        ({"tokens": {"input_tokens": 7, "output_tokens": 3}}, (7, 3, 10)),
+        ({"tokens": {}, "billed_units": {"total_tokens": 42}}, (42, 0, 42)),
+        ({"tokens": {"output_tokens": 3}, "billed_units": {"total_tokens": 12}}, (9, 3, 12)),
+        ({"tokens": {"input_tokens": 7, "output_tokens": 3}, "billed_units": {"total_tokens": 0}}, (7, 3, 0)),
+        ({"billed_units": {"input_tokens": 7, "output_tokens": 3}}, (7, 3, 10)),
+        (
+            {
+                "tokens": {"input_tokens": 7, "output_tokens": 3},
+                "billed_units": {"input_tokens": 90, "output_tokens": 10, "total_tokens": 12},
+            },
+            (7, 3, 12),
+        ),
+        (
+            {
+                "tokens": {"input_tokens": 0, "output_tokens": 0},
+                "billed_units": {"input_tokens": 90, "output_tokens": 10, "total_tokens": 12},
+            },
+            (0, 0, 12),
+        ),
+        (
+            {
+                "tokens": {"input_tokens": None, "output_tokens": None},
+                "billed_units": {"input_tokens": 7, "output_tokens": 3, "total_tokens": None},
+            },
+            (7, 3, 10),
+        ),
+        ({"tokens": {"input_tokens": 0}}, (0, 0, 0)),
+        (None, None),
+        ({}, None),
+        ({"tokens": None, "billed_units": None}, None),
+        ({"billed_units": {"search_units": 3}}, None),
+    ],
+)
+def test_rerank_metadata_reaches_standard_and_spend_logs(
+    logging_obj: Logging, meta: dict[str, object] | None, expected: tuple[int, int, int] | None
+) -> None:
+    from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload
+
+    now: Final = datetime_unit_test(2026, 1, 1, 12, 0, 0)
+    native: Final = {"id": "rerank-usage", "results": [{"index": 0, "relevance_score": 0.9}], "meta": meta}
+    response: Final = RerankResponse.model_validate(native)
+    counts: Final = expected if expected is not None else (0, 0, 0)
+    expected_usage: Final = dict(zip(("prompt_tokens", "completion_tokens", "total_tokens"), counts))
+    assert response.model_dump() == native
+    usage: Final = StandardLoggingPayloadSetup.get_usage_from_response_obj(response.model_dump())
+    assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == counts
+    assert StandardLoggingPayloadSetup.get_usage_as_dict(response.model_dump()) == (
+        expected_usage if expected is not None else {}
+    )
+    kwargs: Final = {"model": "rerank-model", "call_type": "arerank", "litellm_params": {"metadata": {}}}
+    standard: Final = get_standard_logging_object_payload(
+        kwargs=kwargs,
+        init_response_obj=response,
+        start_time=now,
+        end_time=now,
+        logging_obj=logging_obj,
+        status="success",
+    )
+    assert standard is not None
+    assert (standard["prompt_tokens"], standard["completion_tokens"], standard["total_tokens"]) == counts
+    assert standard["metadata"]["usage_object"] == (expected_usage if expected is not None else {})
+    spend: Final = get_logging_payload(
+        kwargs={**kwargs, "standard_logging_object": standard}, response_obj=response, start_time=now, end_time=now
+    )
+    assert (spend["prompt_tokens"], spend["completion_tokens"], spend["total_tokens"]) == counts
+
+
+@pytest.mark.parametrize("combined", [None, Usage(prompt_tokens=11, completion_tokens=2, total_tokens=13)])
+@pytest.mark.parametrize("explicit_usage", [{}, {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}])
+def test_rerank_metadata_does_not_override_explicit_or_combined_usage(
+    combined: Usage | None, explicit_usage: dict[str, int]
+) -> None:
+    response: Final = {"results": [], "meta": {"tokens": {"input_tokens": 42}}, "usage": explicit_usage}
+    expected: Final = combined.model_dump() if combined is not None else explicit_usage
+    assert StandardLoggingPayloadSetup.get_usage_as_dict(response, combined) == expected
+    usage: Final = StandardLoggingPayloadSetup.get_usage_from_response_obj(response, combined)
+    assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (
+        expected.get("prompt_tokens", 0),
+        expected.get("completion_tokens", 0),
+        expected.get("total_tokens", 0),
+    )
+
+
+@pytest.mark.parametrize(
+    "response",
+    [
+        {"meta": {"tokens": {"input_tokens": 42}}},
+        {"results": [], "meta": {"tokens": {"input_tokens": "invalid"}}},
+    ],
+)
+def test_rerank_metadata_ignores_unrelated_and_invalid_responses(response: dict[str, object]) -> None:
+    assert StandardLoggingPayloadSetup.get_usage_as_dict(response) == {}
+    usage: Final = StandardLoggingPayloadSetup.get_usage_from_response_obj(response)
+    assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (0, 0, 0)
+
+
 def test_append_system_prompt_messages():
     """
     Test append_system_prompt_messages prepends system message from kwargs to messages list.

tests/unit/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py

@@ -148,6 +148,41 @@ class TestHostedVLLMRerankTransform:
             self.config._transform_response(response_dict)
 
 
+@pytest.mark.parametrize(
+    "meta",
+    [
+        {
+            "api_version": {"version": "2"},
+            "tokens": {"input_tokens": 7, "output_tokens": 3},
+            "billed_units": {"total_tokens": 12},
+        },
+        {"tokens": {"input_tokens": 0}, "billed_units": {"input_tokens": 0, "output_tokens": 0}},
+        {"tokens": None, "billed_units": None},
+        {},
+    ],
+)
+@pytest.mark.parametrize("legacy_usage", [{}, {"usage": {"total_tokens": 99}}])
+def test_hosted_vllm_rerank_preserves_native_metadata_over_legacy_usage(
+    meta: dict[str, object], legacy_usage: dict[str, object]
+) -> None:
+    response: Final = HostedVLLMRerankConfig()._transform_response(
+        {"id": "rerank-native", "results": [{"index": 0, "relevance_score": 0.9}], "meta": meta, **legacy_usage}
+    )
+    assert response.meta == meta
+    assert response.model_dump() == {
+        "id": "rerank-native",
+        "results": [{"index": 0, "relevance_score": 0.9}],
+        "meta": meta,
+    }
+
+
+def test_hosted_vllm_rerank_null_metadata_uses_legacy_usage() -> None:
+    response: Final = HostedVLLMRerankConfig()._transform_response(
+        {"results": [], "meta": None, "usage": {"total_tokens": 42}}
+    )
+    assert response.meta == {"billed_units": {"total_tokens": 42}, "tokens": {"input_tokens": 42}}
+
+
 class TestGetOptionalRerankParamsInstruction:
     """`instruction` is threaded through get_optional_rerank_params only when set."""
 

tests/unit/llms/jina_ai/rerank/test_transformation.py

@@ -46,7 +46,10 @@ def test_transform_rerank_response_maps_results_and_usage():
         ({"results": [], "usage": {}}, {"billed_units": {}, "tokens": {}}),
         (
             {"results": [], "usage": {"total_tokens": 12, "input_tokens": 7, "output_tokens": 3, "unknown": 1}},
-            {"billed_units": {"total_tokens": 12}, "tokens": {"input_tokens": 7, "output_tokens": 3}},
+            {
+                "billed_units": {"total_tokens": 12, "input_tokens": 7, "output_tokens": 3},
+                "tokens": {"input_tokens": 7, "output_tokens": 3},
+            },
         ),
     ],
 )