From 7df89930b467667ad8faab172b947725f5639094 Mon Sep 17 00:00:00 2001 From: Alexsander Hamir Date: Tue, 10 Feb 2026 12:06:49 -0800 Subject: [PATCH] fix: resolve bedrock base model to cost map key in test_get_model_info_bedrock_models get_base_model() returns model id without 'bedrock/' prefix; cost map keys use 'bedrock/'. Resolve base_model to actual key (try base_model then bedrock/base_model) and skip when not in map to fix KeyError for moonshotai.kimi-k2-thinking. --- tests/local_testing/test_get_model_info.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/tests/local_testing/test_get_model_info.py b/tests/local_testing/test_get_model_info.py index b84fc22af01..d46a087eb73 100644 --- a/tests/local_testing/test_get_model_info.py +++ b/tests/local_testing/test_get_model_info.py @@ -311,7 +311,15 @@ def test_get_model_info_bedrock_models(): for commitment in potential_commitments: k = k.replace(f"{commitment}/", "") base_model = BedrockModelInfo.get_base_model(k) - base_model_info = litellm.model_cost[base_model] + # get_base_model() returns model id without "bedrock/" prefix; cost map keys use "bedrock/" + base_model_key = ( + base_model + if base_model in litellm.model_cost + else f"bedrock/{base_model}" + ) + if base_model_key not in litellm.model_cost: + continue + base_model_info = litellm.model_cost[base_model_key] for base_model_key, base_model_value in base_model_info.items(): if "invoke/" in k: continue