From 1a84e9a24182809ff6282530913ced6cb93b5d2b Mon Sep 17 00:00:00 2001
From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Date: Thu, 30 Jul 2026 05:03:49 +0000
Subject: [PATCH] fix(hosted_vllm/rerank): bill one search unit so
 input_cost_per_query is tracked

---
 .../llms/hosted_vllm/rerank/transformation.py    |  2 +-
 .../test_hosted_vllm_rerank_transformation.py    | 16 ++++++++++++++++
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/litellm/llms/hosted_vllm/rerank/transformation.py b/litellm/llms/hosted_vllm/rerank/transformation.py
index 77504eba04a..04c27c5ba28 100644
--- a/litellm/llms/hosted_vllm/rerank/transformation.py
+++ b/litellm/llms/hosted_vllm/rerank/transformation.py
@@ -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)
 
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..e5cff6d2ae5 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
@@ -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="):