diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 0b6a012284a..a054e367660 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -769,15 +769,15 @@ def is_request_body_safe( ) for target in iter_request_fallback_targets(request_body): if isinstance(target, dict): + target_model = target.get("model") _check_banned_params( target, general_settings, llm_router, - model, + target_model if isinstance(target_model, str) else model, manages_deployments=manages_deployments, calls_an_llm=calls_an_llm, ) - target_model = target.get("model") if isinstance(target_model, str): _reject_url_valued_fallback_target(target_model) elif isinstance(target, str): diff --git a/tests/unit/proxy/auth/test_auth_utils.py b/tests/unit/proxy/auth/test_auth_utils.py index 1396ee3919d..f19fdf810d8 100644 --- a/tests/unit/proxy/auth/test_auth_utils.py +++ b/tests/unit/proxy/auth/test_auth_utils.py @@ -5088,6 +5088,46 @@ def test_model_id_stays_allowed_where_it_does_not_pick_the_bedrock_model( ) +@pytest.mark.parametrize( + "fallback_model, model_id, accepted", + [("claude", "us.amazon.nova-micro-v1:0", False), ("sm-2", "my-inference-component", True)], + ids=["bedrock fallback rejected", "sagemaker fallback allowed"], +) +def test_fallback_target_model_id_is_judged_by_the_fallback_model( + fallback_model: str, model_id: str, accepted: bool +) -> None: + router: Final = _router_with( + ("sm", "sagemaker_chat/my-endpoint"), + ("sm-2", "sagemaker_chat/other-endpoint"), + ("claude", "bedrock/us.anthropic.claude-sonnet-4-6"), + ) + request_body: Final = { + "model": "sm", + "messages": [{"role": "user", "content": "hi"}], + "fallbacks": [{"model": fallback_model, "model_id": model_id}], + } + if accepted: + assert ( + is_request_body_safe( + request_body=request_body, + general_settings={}, + llm_router=router, + model="sm", + route="/v1/chat/completions", + ) + is True + ) + return + with pytest.raises(ValueError, match="model_id"): + is_request_body_safe( + request_body=request_body, + general_settings={}, + llm_router=router, + model="sm", + route="/v1/chat/completions", + ) + + def test_model_id_is_rejected_when_the_model_group_mixes_in_a_bedrock_deployment() -> None: router: Final = _router_with(("sm", "sagemaker_chat/my-endpoint"), ("sm", "bedrock/us.anthropic.claude-sonnet-4-6")) with pytest.raises(ValueError, match="model_id"):