Skip to content

Commit 230db7e

Browse files
authored
Merge pull request #18100 from BerriAI/litellm_bedrock_qwen_arn_fix
fix: Add qwen 2 and qwen 3 in get_bedrock_model_id
2 parents 25f2213 + f83b821 commit 230db7e

File tree

2 files changed

+77
-0
lines changed

2 files changed

+77
-0
lines changed

litellm/llms/bedrock/base_aws_llm.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -357,6 +357,14 @@ def get_bedrock_model_id(
357357
model_id = BaseAWSLLM._get_model_id_from_model_with_spec(
358358
model_id, spec="openai"
359359
)
360+
elif provider == "qwen2" and "qwen2/" in model_id:
361+
model_id = BaseAWSLLM._get_model_id_from_model_with_spec(
362+
model_id, spec="qwen2"
363+
)
364+
elif provider == "qwen3" and "qwen3/" in model_id:
365+
model_id = BaseAWSLLM._get_model_id_from_model_with_spec(
366+
model_id, spec="qwen3"
367+
)
360368
return model_id
361369

362370
@staticmethod

tests/test_litellm/llms/bedrock/chat/invoke_transformations/test_amazon_qwen2_transformation.py

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -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

Comments
 (0)