diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 43e60e49e3b..fd746d62e76 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -1981,16 +1981,20 @@ def get_model_from_request( if vertex_match: model = vertex_match.group(1) - if model is None and route.lower().startswith("/bedrock"): - from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( - _extract_model_from_bedrock_endpoint, - ) - + if route.lower().startswith("/bedrock"): bedrock_endpoint = re.sub(r"^/bedrock/", "", route, flags=re.IGNORECASE) - try: - model = _extract_model_from_bedrock_endpoint(bedrock_endpoint) - except ValueError: - model = None + is_bedrock_count_tokens_route = ( + "count_tokens" in bedrock_endpoint.lower() or "count-tokens" in bedrock_endpoint.lower() + ) + if not is_bedrock_count_tokens_route: + from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( + _extract_model_from_bedrock_endpoint, + ) + + try: + model = _extract_model_from_bedrock_endpoint(bedrock_endpoint) + except ValueError: + pass return model diff --git a/tests/test_litellm/proxy/auth/test_auth_utils.py b/tests/test_litellm/proxy/auth/test_auth_utils.py index f3426534c28..61d1a308b9d 100644 --- a/tests/test_litellm/proxy/auth/test_auth_utils.py +++ b/tests/test_litellm/proxy/auth/test_auth_utils.py @@ -478,6 +478,46 @@ def test_get_model_from_request_bedrock_unparseable_endpoint_returns_none(): ) +def test_get_model_from_request_bedrock_url_model_overrides_body_model(): + assert ( + get_model_from_request( + request_data={"model": "us.anthropic.claude-sonnet-4-6"}, + route="/bedrock/model/us.anthropic.claude-opus-4-6-v1/converse", + ) + == "us.anthropic.claude-opus-4-6-v1" + ) + + +def test_get_model_from_request_bedrock_invoke_url_model_overrides_body_model(): + assert ( + get_model_from_request( + request_data={"model": "us.anthropic.claude-sonnet-4-6"}, + route="/bedrock/model/us.anthropic.claude-opus-4-6-v1/invoke", + ) + == "us.anthropic.claude-opus-4-6-v1" + ) + + +def test_get_model_from_request_bedrock_count_tokens_uses_body_model(): + assert ( + get_model_from_request( + request_data={"model": "us.anthropic.claude-sonnet-4-6"}, + route="/bedrock/v1/messages/count_tokens", + ) + == "us.anthropic.claude-sonnet-4-6" + ) + + +def test_get_model_from_request_bedrock_unparseable_endpoint_keeps_body_model(): + assert ( + get_model_from_request( + request_data={"model": "us.anthropic.claude-sonnet-4-6"}, + route="/bedrock/agents/some-agent-route", + ) + == "us.anthropic.claude-sonnet-4-6" + ) + + def test_get_model_from_request_includes_file_endpoint_header_model(): assert ( get_model_from_request(