@@ -290,3 +290,72 @@ def test_qwen2_provider_detection():
290290 assert config is not None
291291 assert isinstance (config , AmazonQwen2Config )
292292
293+
294+ def test_qwen2_model_id_extraction_with_arn ():
295+ """Test that model ID is correctly extracted from bedrock/qwen2/arn... paths"""
296+ from litellm .llms .bedrock .base_aws_llm import BaseAWSLLM
297+
298+ # Test case: bedrock/qwen2/arn:aws:bedrock:us-east-1:123456789012:imported-model/test-qwen2
299+ # The qwen2/ prefix should be stripped, leaving only the ARN for encoding
300+ model = "qwen2/arn:aws:bedrock:us-east-1:123456789012:imported-model/test-qwen2"
301+ provider = "qwen2"
302+
303+ result = BaseAWSLLM .get_bedrock_model_id (
304+ optional_params = {},
305+ provider = provider ,
306+ model = model
307+ )
308+
309+ # The result should NOT contain "qwen2/" - it should be stripped
310+ assert "qwen2/" not in result
311+ # The result should be URL-encoded ARN
312+ assert "arn%3Aaws%3Abedrock" in result or "arn:aws:bedrock" in result
313+
314+
315+ def test_qwen2_model_id_extraction_without_qwen2_prefix ():
316+ """Test that model ID extraction doesn't strip qwen2/ when provider is not qwen2"""
317+ from litellm .llms .bedrock .base_aws_llm import BaseAWSLLM
318+
319+ # Test case: just a model name without qwen2/ prefix
320+ model = "arn:aws:bedrock:us-east-1:123456789012:imported-model/test-qwen2"
321+ provider = "qwen2"
322+
323+ result = BaseAWSLLM .get_bedrock_model_id (
324+ optional_params = {},
325+ provider = provider ,
326+ model = model
327+ )
328+
329+ # Result should be encoded ARN
330+ assert "arn" in result .lower () or "aws" in result .lower ()
331+
332+
333+ def test_qwen2_get_bedrock_model_id_with_various_formats ():
334+ """Test get_bedrock_model_id with various Qwen2 model path formats"""
335+ from litellm .llms .bedrock .base_aws_llm import BaseAWSLLM
336+
337+ test_cases = [
338+ {
339+ "model" : "qwen2/arn:aws:bedrock:us-east-1:123456789012:imported-model/test-qwen2" ,
340+ "provider" : "qwen2" ,
341+ "should_not_contain" : "qwen2/" ,
342+ "description" : "Qwen2 imported model ARN"
343+ },
344+ {
345+ "model" : "bedrock/qwen2/arn:aws:bedrock:us-east-1:123456789012:imported-model/test-qwen2" ,
346+ "provider" : "qwen2" ,
347+ "should_not_contain" : "qwen2/" ,
348+ "description" : "Bedrock prefixed Qwen2 ARN"
349+ }
350+ ]
351+
352+ for test_case in test_cases :
353+ result = BaseAWSLLM .get_bedrock_model_id (
354+ optional_params = {},
355+ provider = test_case ["provider" ],
356+ model = test_case ["model" ]
357+ )
358+
359+ assert test_case ["should_not_contain" ] not in result , \
360+ f"Failed for { test_case ['description' ]} : { test_case ['should_not_contain' ]} found in { result } "
361+
0 commit comments