From beb462c85cb916a21bc623a740ab9c82b3e14b05 Mon Sep 17 00:00:00 2001
From: spions <spionsbbs@gmail.com>
Date: Fri, 7 Aug 2026 13:17:49 +0300
Subject: [PATCH] fix(hosted_vllm): read rerank usage from both `meta` and
 `usage`

HostedVLLMRerankConfig._transform_response() only read token counts
from a top-level `usage` key. vLLM's /rerank endpoint (Cohere-compatible
API) actually returns them under `meta.billed_units`/`meta.tokens`, so
usage was silently zeroed for hosted_vllm rerank models. Now checks
both shapes, keeping backward compatibility with servers that use `usage`.
---
 .../llms/hosted_vllm/rerank/transformation.py | 13 ++++++++++---
 .../test_hosted_vllm_rerank_transformation.py | 19 +++++++++++++++++++
 2 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/litellm/llms/hosted_vllm/rerank/transformation.py b/litellm/llms/hosted_vllm/rerank/transformation.py
index 74c13b450f5..b467bcc5317 100644
--- a/litellm/llms/hosted_vllm/rerank/transformation.py
+++ b/litellm/llms/hosted_vllm/rerank/transformation.py
@@ -172,10 +172,17 @@ class HostedVLLMRerankConfig(BaseRerankConfig):
         return HostedVLLMRerankError(message=error_message, status_code=status_code, headers=headers)
 
     def _transform_response(self, response: dict) -> RerankResponse:
-        # Extract usage information
+        # Extract usage information - some servers (vLLM/Cohere-style) return a
+        # top-level `meta` object; others (OpenAI/TEI-style) return `usage`.
+        # Check both so either shape is picked up.
+        raw_meta: Final = response.get("meta") or {}
         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))
+        total_tokens: Final = raw_meta.get("billed_units", {}).get("total_tokens") or usage_data.get(
+            "total_tokens", 0
+        )
+        input_tokens: Final = raw_meta.get("tokens", {}).get("input_tokens") or usage_data.get("total_tokens", 0)
+        _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
diff --git a/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py b/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py
index 6425e815db0..989efe4c105 100644
--- a/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py
+++ b/tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py
@@ -131,6 +131,25 @@ 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
+
     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="):