test_bedrock_get_base_model

Author: ishaan-jaff

tests/llm_translation/test_bedrock_completion.py

@@ -1265,7 +1265,9 @@ def test_bedrock_cross_region_inference(model):
     ],
 )
 def test_bedrock_get_base_model(model, expected_base_model):
-    assert litellm.AmazonConverseConfig()._get_base_model(model) == expected_base_model
+    from litellm.llms.bedrock.common_utils import BedrockModelInfo
+
+    assert BedrockModelInfo.get_base_model(model) == expected_base_model
 
 
 from litellm.litellm_core_utils.prompt_templates.factory import (
@@ -1982,9 +1984,18 @@ def test_bedrock_mapped_converse_models():
 
 
 def test_bedrock_base_model_helper():
+    from litellm.llms.bedrock.common_utils import BedrockModelInfo
+
     model = "us.amazon.nova-pro-v1:0"
-    litellm.AmazonConverseConfig()._get_base_model(model)
-    assert model == "us.amazon.nova-pro-v1:0"
+    base_model = BedrockModelInfo.get_base_model(model)
+    assert base_model == "us.amazon.nova-pro-v1:0"
+
+    assert (
+        BedrockModelInfo.get_base_model(
+            "invoke/anthropic.claude-3-5-sonnet-20241022-v2:0"
+        )
+        == "anthropic.claude-3-5-sonnet-20241022-v2:0"
+    )
 
 
 @pytest.mark.parametrize(