From bc190f201ec29d38009feece7b8c277f36092785 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Thu, 20 Feb 2025 15:32:49 -0800
Subject: [PATCH] fix(cost_calculator.py): bedrock/common_utils.py

get base model from model w/ arn -> handles rerank model
---
 litellm/cost_calculator.py                    |  1 +
 litellm/llms/bedrock/common_utils.py          | 19 +++++++++++++++++++
 .../llm_translation/base_rerank_unit_tests.py |  2 ++
 3 files changed, 22 insertions(+)

diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py
index c335a3d5f10..f8ae47d04a9 100644
--- a/litellm/cost_calculator.py
+++ b/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
diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py
index 8a534f6eac4..54be3598974 100644
--- a/litellm/llms/bedrock/common_utils.py
+++ b/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)[
diff --git a/tests/llm_translation/base_rerank_unit_tests.py b/tests/llm_translation/base_rerank_unit_tests.py
index 707de7d86ce..4ebf509f7a0 100644
--- a/tests/llm_translation/base_rerank_unit_tests.py
+++ b/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: