test(bedrock): add test for [1m] context window suffix stripping in cost lookup

This commit is contained in:
Ishaan Jaffer 2026-04-01 12:17:08 -07:00
parent 440c3387df
commit a73a944abc

View file

@ -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"
)