fix(cost_calculator.py): bedrock/common_utils.py

get base model from model w/ arn -> handles rerank model

litellm/cost_calculator.py

@@ -838,6 +838,7 @@ def rerank_cost(
     Returns
     - float or None: cost of response OR none if error.
     """
+
     default_num_queries = 1
     _, custom_llm_provider, _, _ = litellm.get_llm_provider(
         model=model, custom_llm_provider=custom_llm_provider

litellm/llms/bedrock/common_utils.py

@@ -318,6 +318,23 @@ class BedrockModelInfo(BaseLLMModelInfo):
     global_config = AmazonBedrockGlobalConfig()
     all_global_regions = global_config.get_all_regions()
 
+    @staticmethod
+    def extract_model_name_from_arn(model: str) -> str:
+        """
+        Extract the model name from an AWS Bedrock ARN.
+        Returns the string after the last '/' if 'arn' is in the input string.
+
+        Args:
+            arn (str): The ARN string to parse
+
+        Returns:
+            str: The extracted model name if 'arn' is in the string,
+                otherwise returns the original string
+        """
+        if "arn" in model.lower():
+            return model.split("/")[-1]
+        return model
+
     @staticmethod
     def get_base_model(model: str) -> str:
         """
@@ -335,6 +352,8 @@ class BedrockModelInfo(BaseLLMModelInfo):
         if model.startswith("invoke/"):
             model = model.split("/", 1)[1]
 
+        model = BedrockModelInfo.extract_model_name_from_arn(model)
+
         potential_region = model.split(".", 1)[0]
 
         alt_potential_region = model.split("/", 1)[

tests/llm_translation/base_rerank_unit_tests.py

@@ -80,6 +80,8 @@ class BaseLLMRerankTest(ABC):
     @pytest.mark.parametrize("sync_mode", [True, False])
     async def test_basic_rerank(self, sync_mode):
         litellm.set_verbose = True
+        os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
+        litellm.model_cost = litellm.get_model_cost_map(url="")
         rerank_call_args = self.get_base_rerank_call_args()
         custom_llm_provider = self.get_custom_llm_provider()
         if sync_mode is True: