scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 156

fix(hosted_vllm/rerank): bill one search unit so input_cost_per_query is tracked

Browse source
This commit is contained in:
Devin AI 2026-07-30 05:03:49 +00:00
parent 4d54324515
commit 1a84e9a241
2 changed files with 17 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -176,7 +176,7 @@ class HostedVLLMRerankConfig(BaseRerankConfig):
def _transform_response(self, response: dict) -> RerankResponse:
# Extract usage information
usage_data = response.get("usage", {})
_billed_units = RerankBilledUnits(total_tokens=usage_data.get("total_tokens", 0))
_billed_units = RerankBilledUnits(search_units=1, total_tokens=usage_data.get("total_tokens", 0))
_tokens = RerankTokens(input_tokens=usage_data.get("total_tokens", 0))
rerank_meta = RerankResponseMeta(billed_units=_billed_units, tokens=_tokens)

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

@ -129,8 +129,24 @@ class TestHostedVLLMRerankTransform:
assert result.results[0]["relevance_score"] == 0.9
assert result.results[0]["document"]["text"] == "doc1 text"
assert result.meta["billed_units"]["total_tokens"] == 42
assert result.meta["billed_units"]["search_units"] == 1
assert result.meta["tokens"]["input_tokens"] == 42
def test_transform_response_reports_search_unit_for_per_query_cost(self):
response_dict = {
"results": [{"index": 0, "relevance_score": 0.9}],
"usage": {"total_tokens": 42},
}
result = self.config._transform_response(response_dict)
assert result.meta["billed_units"]["search_units"] == 1
prompt_cost, completion_cost = self.config.calculate_rerank_cost(
model=self.model,
billed_units=result.meta["billed_units"],
model_info={"input_cost_per_query": 0.002},
)
assert prompt_cost == 0.002
assert completion_cost == 0.0
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="):