scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-07 02:59:05 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 103

fix(hosted_vllm): read rerank usage from both meta and usage

Browse source 
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`.
This commit is contained in:
spions 2026-08-07 13:17:49 +03:00
parent ad9f6ce23b
commit beb462c85c
2 changed files with 29 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

13
litellm/llms/hosted_vllm/rerank/transformation.py
View file

@ -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

19
tests/test_litellm/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py
View file

@ -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="):