From a73a944abc77dd1ac1da2b7d265a6153862feedf Mon Sep 17 00:00:00 2001 From: Ishaan Jaffer Date: Wed, 1 Apr 2026 12:17:08 -0700 Subject: [PATCH] test(bedrock): add test for [1m] context window suffix stripping in cost lookup --- .../llms/bedrock/test_bedrock_common_utils.py | 40 +++++++++++++++++-- 1 file changed, 36 insertions(+), 4 deletions(-) diff --git a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py index b804f549ecb..e1c1ea04e65 100644 --- a/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py +++ b/tests/test_litellm/llms/bedrock/test_bedrock_common_utils.py @@ -29,25 +29,25 @@ def test_govcloud_cross_region_inference_prefix(): Test that GovCloud models with cross-region inference prefix (us-gov.) are parsed correctly """ bedrock_model_info = BedrockModelInfo - + # Test us-gov prefix is stripped correctly for Claude models base_model = bedrock_model_info.get_base_model( model="bedrock/us-gov.anthropic.claude-3-5-sonnet-20240620-v1:0" ) assert base_model == "anthropic.claude-3-5-sonnet-20240620-v1:0" - + # Test us-gov prefix is stripped correctly for different Claude versions base_model = bedrock_model_info.get_base_model( model="bedrock/us-gov.anthropic.claude-sonnet-4-5-20250929-v1:0" ) assert base_model == "anthropic.claude-sonnet-4-5-20250929-v1:0" - + # Test us-gov prefix is stripped correctly for Haiku models base_model = bedrock_model_info.get_base_model( model="bedrock/us-gov.anthropic.claude-3-haiku-20240307-v1:0" ) assert base_model == "anthropic.claude-3-haiku-20240307-v1:0" - + # Test us-gov prefix is stripped correctly for Meta models base_model = bedrock_model_info.get_base_model( model="bedrock/us-gov.meta.llama3-8b-instruct-v1:0" @@ -55,3 +55,35 @@ def test_govcloud_cross_region_inference_prefix(): assert base_model == "meta.llama3-8b-instruct-v1:0" +def test_context_window_suffix_stripped_for_cost_lookup(): + """ + Test that [1m], [200k] etc. context window suffixes are stripped from + Bedrock model names before cost lookup. + + Models configured like `bedrock/us.anthropic.claude-opus-4-6-v1[1m]` + should resolve to the base model name so pricing can be found. + """ + from litellm.llms.bedrock.common_utils import get_bedrock_base_model + + assert ( + get_bedrock_base_model("us.anthropic.claude-opus-4-6-v1[1m]") + == "anthropic.claude-opus-4-6-v1" + ) + assert ( + get_bedrock_base_model("us.anthropic.claude-sonnet-4-6[1m]") + == "anthropic.claude-sonnet-4-6" + ) + assert ( + get_bedrock_base_model("global.anthropic.claude-opus-4-5-20251101-v1:0[1m]") + == "anthropic.claude-opus-4-5-20251101-v1:0" + ) + # Ensure models without suffix are unaffected + assert ( + get_bedrock_base_model("us.anthropic.claude-opus-4-6-v1") + == "anthropic.claude-opus-4-6-v1" + ) + # Ensure :51k throughput suffix still works + assert ( + get_bedrock_base_model("anthropic.claude-3-5-sonnet-20241022-v2:0:51k") + == "anthropic.claude-3-5-sonnet-20241022-v2:0" + )