Merge 6e3bf73da9 into f2118fc767

litellm/litellm_core_utils/litellm_logging.py

@@ -6603,6 +6603,18 @@ class StandardLoggingPayloadSetup:
 
         return clean_metadata
 
+    @staticmethod
+    def _get_rerank_usage_from_meta(meta: object) -> Usage | None:
+        if not isinstance(meta, dict) or ("billed_units" not in meta and "tokens" not in meta):
+            return None
+        billed_units: Final = meta.get("billed_units")
+        tokens: Final = meta.get("tokens")
+        meta_total: Final = billed_units.get("total_tokens") if isinstance(billed_units, dict) else None
+        meta_input: Final = tokens.get("input_tokens") if isinstance(tokens, dict) else None
+        total_tokens: Final = meta_total if meta_total is not None else (meta_input if meta_input is not None else 0)
+        prompt_tokens: Final = meta_input if meta_input is not None else (meta_total if meta_total is not None else 0)
+        return Usage(prompt_tokens=prompt_tokens, completion_tokens=0, total_tokens=total_tokens)
+
     @staticmethod
     def get_usage_from_response_obj(response_obj: dict | None, combined_usage_object: Usage | None = None) -> Usage:
         ## BASE CASE ##
@@ -6616,6 +6628,10 @@ class StandardLoggingPayloadSetup:
             )
 
         usage: Final = response_obj.get("usage", None) or {}
+        if not usage:
+            rerank_usage: Final = StandardLoggingPayloadSetup._get_rerank_usage_from_meta(response_obj.get("meta"))
+            if rerank_usage is not None:
+                return rerank_usage
         if usage is None or (not isinstance(usage, dict) and not isinstance(usage, Usage)):
             return Usage(
                 prompt_tokens=0,
@@ -6651,6 +6667,9 @@ class StandardLoggingPayloadSetup:
             return _empty
         _raw: Final = response_obj.get("usage", None)
         if _raw is None:
+            rerank_usage: Final = StandardLoggingPayloadSetup._get_rerank_usage_from_meta(response_obj.get("meta"))
+            if rerank_usage is not None:
+                return rerank_usage.model_dump()
             return _empty
         if isinstance(_raw, ResponseAPIUsage):
             return ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(_raw).model_dump()

litellm/llms/hosted_vllm/rerank/transformation.py

@@ -203,10 +203,21 @@ class HostedVLLMRerankConfig(BaseRerankConfig):
         return HostedVLLMRerankError(message=error_message, status_code=status_code, headers=headers)
 
     def _transform_response(self, response: dict) -> RerankResponse:
-        # Extract usage information
-        usage_data: Final = response.get("usage", {})
-        _billed_units: Final = RerankBilledUnits(total_tokens=usage_data.get("total_tokens", 0))
-        _tokens: Final = RerankTokens(input_tokens=usage_data.get("total_tokens", 0))
+        raw_meta: Final = response.get("meta")
+        usage_data: Final = response.get("usage")
+        meta_billed_units: Final = raw_meta.get("billed_units") if isinstance(raw_meta, dict) else None
+        meta_tokens: Final = raw_meta.get("tokens") if isinstance(raw_meta, dict) else None
+        usage_total_tokens: Final = usage_data.get("total_tokens", 0) if isinstance(usage_data, dict) else 0
+        meta_total: Final = meta_billed_units.get("total_tokens") if isinstance(meta_billed_units, dict) else None
+        meta_input: Final = meta_tokens.get("input_tokens") if isinstance(meta_tokens, dict) else None
+        total_tokens: Final = (
+            meta_total if meta_total is not None else (meta_input if meta_input is not None else usage_total_tokens)
+        )
+        input_tokens: Final = (
+            meta_input if meta_input is not None else (meta_total if meta_total is not None else usage_total_tokens)
+        )
+        _billed_units: Final = RerankBilledUnits(total_tokens=total_tokens)
+        _tokens: Final = RerankTokens(input_tokens=input_tokens)
         rerank_meta: Final = RerankResponseMeta(billed_units=_billed_units, tokens=_tokens)
 
         # Extract results

tests/unit/litellm_core_utils/test_litellm_logging.py

@@ -7323,6 +7323,97 @@ def test_pre_call_does_not_pin_request_in_module_state(logging_obj):
     assert litellm.error_logs == {}
 
 
+@pytest.mark.parametrize(
+    "meta,expected",
+    [
+        (None, None),
+        ({}, None),
+        ({"billed_units": {"total_tokens": 42}, "tokens": {"input_tokens": 42}}, (42, 42)),
+        ({"tokens": {"input_tokens": 38}}, (38, 38)),
+        ({"billed_units": {"total_tokens": 38}}, (38, 38)),
+    ],
+)
+def test_get_rerank_usage_from_meta(meta, expected):
+    """Direct unit test of the helper both get_usage_from_response_obj and
+    get_usage_as_dict delegate to for the rerank `meta` shape."""
+    from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
+
+    result = StandardLoggingPayloadSetup._get_rerank_usage_from_meta(meta)
+
+    if expected is None:
+        assert result is None
+    else:
+        prompt_tokens, total_tokens = expected
+        assert result is not None
+        assert result.prompt_tokens == prompt_tokens
+        assert result.completion_tokens == 0
+        assert result.total_tokens == total_tokens
+
+
+def test_get_usage_from_response_obj_rerank_meta():
+    """
+    RerankResponse has no top-level `usage` field - token usage lives under
+    `meta.billed_units`/`meta.tokens` (Cohere-style API), which is also what
+    hosted_vllm/other rerank servers return. Without this, /ui/usage always
+    showed 0 tokens for every rerank call.
+    """
+    from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
+
+    response_obj = {
+        "id": "rerank-d618748e0f5543e8ba09ee7dd131ac59",
+        "results": [],
+        "meta": {
+            "billed_units": {"total_tokens": 42},
+            "tokens": {"input_tokens": 42},
+        },
+    }
+
+    usage = StandardLoggingPayloadSetup.get_usage_from_response_obj(response_obj)
+
+    assert usage.prompt_tokens == 42
+    assert usage.completion_tokens == 0
+    assert usage.total_tokens == 42
+
+
+def test_get_usage_as_dict_rerank_meta():
+    """
+    get_usage_as_dict() is the function actually called by
+    get_standard_logging_object_payload() (the hot path used to build spend
+    logs / /ui/usage). It must handle the rerank `meta.billed_units`/`meta.tokens`
+    shape the same way get_usage_from_response_obj() does, or /ui/usage keeps
+    showing 0 tokens for rerank calls even after fixing the other function.
+    """
+    from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
+
+    response_obj = {
+        "id": "rerank-d618748e0f5543e8ba09ee7dd131ac59",
+        "results": [],
+        "meta": {
+            "billed_units": {"total_tokens": 42},
+            "tokens": {"input_tokens": 42},
+        },
+    }
+
+    usage_dict = StandardLoggingPayloadSetup.get_usage_as_dict(response_obj)
+
+    assert usage_dict["prompt_tokens"] == 42
+    assert usage_dict["completion_tokens"] == 0
+    assert usage_dict["total_tokens"] == 42
+
+
+def test_get_usage_as_dict_rerank_meta_no_meta():
+    """A rerank response with no usable `meta` (and no `usage`) falls back to zero."""
+    from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
+
+    response_obj = {"id": "rerank-abc", "results": [], "meta": None}
+
+    usage_dict = StandardLoggingPayloadSetup.get_usage_as_dict(response_obj)
+
+    assert usage_dict["prompt_tokens"] == 0
+    assert usage_dict["completion_tokens"] == 0
+    assert usage_dict["total_tokens"] == 0
+
+
 def test_handle_anthropic_messages_response_logging_preserves_fast_mode_speed():
     """/v1/messages non-streaming rebuilds usage by re-transforming the raw Anthropic
     response. Anthropic's fast-mode multiplier is applied off ``usage.speed``, which the
@@ -7662,6 +7753,31 @@ async def test_restore_correlation_context_works_across_asyncio_task_boundary(mo
         session_id_var.set("")
 
 
+@pytest.mark.parametrize(
+    "meta",
+    [
+        {"tokens": {"input_tokens": 38}},
+        {"billed_units": {"total_tokens": 38}},
+    ],
+)
+def test_get_usage_as_dict_rerank_meta_partial_fields_cross_fill(meta):
+    """
+    Some rerank servers report only `tokens.input_tokens` or only
+    `billed_units.total_tokens`, not both. Since rerank has no completion step,
+    either field alone should populate both prompt_tokens and total_tokens
+    instead of leaving one of them at zero.
+    """
+    from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
+
+    response_obj = {"id": "rerank-abc", "results": [], "meta": meta}
+
+    usage_dict = StandardLoggingPayloadSetup.get_usage_as_dict(response_obj)
+
+    assert usage_dict["prompt_tokens"] == 38
+    assert usage_dict["completion_tokens"] == 0
+    assert usage_dict["total_tokens"] == 38
+
+
 class TestNonInferenceCallTypesAreNotBilled:
     """A retrieved response replays the usage of the call that created it, so pricing a read
     of it double bills the same tokens. Regression tests for LIT-5602."""

tests/unit/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py

@@ -133,6 +133,47 @@ class TestHostedVLLMRerankTransform:
         assert result.meta["billed_units"]["total_tokens"] == 42
         assert result.meta["tokens"]["input_tokens"] == 42
 
+    def test_transform_response_with_meta(self):
+        """Some vLLM rerank servers (Cohere-compatible API) return usage under
+        `meta.billed_units`/`meta.tokens` instead of a top-level `usage` key."""
+        response_dict = {
+            "id": "rerank-d618748e0f5543e8ba09ee7dd131ac59",
+            "results": [
+                {"index": 0, "relevance_score": 0.9, "document": {"text": "doc1 text"}},
+                {"index": 1, "relevance_score": 0.7, "document": {"text": "doc2 text"}},
+            ],
+            "meta": {
+                "billed_units": {"total_tokens": 42},
+                "tokens": {"input_tokens": 42},
+            },
+        }
+        result = self.config._transform_response(response_dict)
+        assert result.id == "rerank-d618748e0f5543e8ba09ee7dd131ac59"
+        assert result.meta["billed_units"]["total_tokens"] == 42
+        assert result.meta["tokens"]["input_tokens"] == 42
+
+    @pytest.mark.parametrize(
+        "meta",
+        [
+            {"tokens": {"input_tokens": 38}},
+            {"billed_units": {"total_tokens": 38}},
+        ],
+    )
+    def test_transform_response_with_meta_partial_fields_cross_fill(self, meta):
+        """Some servers report only one of `tokens`/`billed_units`. Since rerank
+        has no completion step, either field alone should populate both
+        `meta.tokens.input_tokens` and `meta.billed_units.total_tokens`."""
+        response_dict = {
+            "id": "rerank-abc",
+            "results": [
+                {"index": 0, "relevance_score": 0.9, "document": {"text": "doc1 text"}},
+            ],
+            "meta": meta,
+        }
+        result = self.config._transform_response(response_dict)
+        assert result.meta["billed_units"]["total_tokens"] == 38
+        assert result.meta["tokens"]["input_tokens"] == 38
+
     def test_transform_response_missing_results(self):
         response_dict = {"id": "abc123", "usage": {"total_tokens": 10}}
         with pytest.raises(ValueError, match="No results found in the response="):