feat(bedrock/rerank/transformation.py): include search units for bedrock rerank result

Resolves https://github.com/BerriAI/litellm/issues/7258#issuecomment-2671557137

litellm/llms/bedrock/rerank/transformation.py

@@ -91,7 +91,9 @@ class BedrockRerankConfig:
         example input:
         {"results":[{"index":0,"relevanceScore":0.6847912669181824},{"index":1,"relevanceScore":0.5980774760246277}]}
         """
-        _billed_units = RerankBilledUnits(**response.get("usage", {}))
+        _billed_units = RerankBilledUnits(
+            **response.get("usage", {"search_units": 1})
+        )  # by default 1 search unit
         _tokens = RerankTokens(**response.get("usage", {}))
         rerank_meta = RerankResponseMeta(billed_units=_billed_units, tokens=_tokens)
 

tests/litellm/llms/bedrock/rerank/transformation.py

@@ -0,0 +1,14 @@
+import json
+import os
+import sys
+
+import pytest
+from fastapi.testclient import TestClient
+
+sys.path.insert(
+    0, os.path.abspath("../../..")
+)  # Adds the parent directory to the system path
+from unittest.mock import MagicMock, patch
+
+from litellm import rerank
+from litellm.llms.custom_httpx.http_handler import HTTPHandler

tests/llm_translation/base_rerank_unit_tests.py

@@ -52,13 +52,13 @@ def assert_response_shape(response, custom_llm_provider):
             response.meta["api_version"]["version"],
             expected_api_version_shape["version"],
         )
-        assert isinstance(
-            response.meta["billed_units"], expected_meta_shape["billed_units"]
-        )
-        assert isinstance(
-            response.meta["billed_units"]["search_units"],
-            expected_billed_units_shape["search_units"],
-        )
+    assert isinstance(
+        response.meta["billed_units"], expected_meta_shape["billed_units"]
+    )
+    assert isinstance(
+        response.meta["billed_units"]["search_units"],
+        expected_billed_units_shape["search_units"],
+    )
 
 
 class BaseLLMRerankTest(ABC):