diff --git a/litellm/llms/bedrock/passthrough/transformation.py b/litellm/llms/bedrock/passthrough/transformation.py index 5791bfb801..568fe94171 100644 --- a/litellm/llms/bedrock/passthrough/transformation.py +++ b/litellm/llms/bedrock/passthrough/transformation.py @@ -34,11 +34,12 @@ class BedrockPassthroughConfig( litellm_params: dict, ) -> Tuple["URL", str]: optional_params = litellm_params.copy() + model_id = optional_params.get("model_id", None) aws_region_name = self._get_aws_region_name( optional_params=optional_params, model=model, - model_id=None, + model_id=model_id, ) aws_bedrock_runtime_endpoint = optional_params.get("aws_bedrock_runtime_endpoint") @@ -49,6 +50,12 @@ class BedrockPassthroughConfig( endpoint_type="runtime", ) + # If model_id is provided (e.g., Application Inference Profile ARN), use it in the endpoint + # instead of the translated model name + if model_id is not None: + # Replace the model name in the endpoint with the model_id + import re + endpoint = re.sub(r'model/[^/]+/', f'model/{model_id}/', endpoint) return self.format_url(endpoint, endpoint_url, request_query_params or {}), endpoint_url def sign_request( diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index 3df3037ed5..df4737cec8 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -216,6 +216,11 @@ def llm_passthrough_route( ) litellm_params_dict = get_litellm_params(**kwargs) + + # Add model_id to litellm_params if present in kwargs (for Bedrock Application Inference Profiles) + if "model_id" in kwargs: + litellm_params_dict["model_id"] = kwargs["model_id"] + litellm_logging_obj.update_environment_variables( model=model, litellm_params=litellm_params_dict, diff --git a/tests/test_litellm/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py b/tests/test_litellm/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py index 7cb1ee2b54..07253a8e09 100644 --- a/tests/test_litellm/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py +++ b/tests/test_litellm/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py @@ -175,3 +175,132 @@ def test_format_url_handles_trailing_slash_normalization(): assert str(result_with_slash) == "http://proxy.com/bedrockproxy/model/test/invoke" +def test_bedrock_passthrough_with_application_inference_profile(): + """ + Test get_complete_url with Application Inference Profile ARN as model_id. + + This test verifies the fix for GitHub issue #18761 where Bedrock passthrough + was not working with Application Inference Profiles. The model_id (ARN) should + replace the translated model name in the endpoint URL. + """ + config = BedrockPassthroughConfig() + + model = "anthropic.claude-sonnet-4-20250514-v1:0" + model_id = "arn:aws:bedrock:eu-west-1:123456789:application-inference-profile/abcdefgh1234" + endpoint = f"model/{model}/invoke" + + with patch.object(config, '_get_aws_region_name', return_value="eu-west-1"), \ + patch.object(config, 'get_runtime_endpoint', return_value=( + "https://bedrock-runtime.eu-west-1.amazonaws.com", + "https://bedrock-runtime.eu-west-1.amazonaws.com" + )): + + url, api_base = config.get_complete_url( + api_base=None, + api_key=None, + model=model, + endpoint=endpoint, + request_query_params=None, + litellm_params={"model_id": model_id, "aws_region_name": "eu-west-1"} + ) + + # Verify that the URL contains the model_id (ARN) instead of the model name + url_str = str(url) + assert model_id in url_str, f"Expected model_id ARN in URL, but got: {url_str}" + assert model not in url_str, f"Model name should be replaced by model_id, but got: {url_str}" + assert "/invoke" in url_str, "Expected /invoke action in URL" + + # Verify the complete URL structure + expected_url = f"https://bedrock-runtime.eu-west-1.amazonaws.com/model/{model_id}/invoke" + assert url_str == expected_url, f"Expected {expected_url}, but got: {url_str}" + + +def test_bedrock_passthrough_with_inference_profile_converse_endpoint(): + """Test Application Inference Profile with converse endpoint""" + config = BedrockPassthroughConfig() + + model = "anthropic.claude-sonnet-4-20250514-v1:0" + model_id = "arn:aws:bedrock:us-east-1:123456789:application-inference-profile/xyz123" + endpoint = f"model/{model}/converse" + + with patch.object(config, '_get_aws_region_name', return_value="us-east-1"), \ + patch.object(config, 'get_runtime_endpoint', return_value=( + "https://bedrock-runtime.us-east-1.amazonaws.com", + "https://bedrock-runtime.us-east-1.amazonaws.com" + )): + + url, api_base = config.get_complete_url( + api_base=None, + api_key=None, + model=model, + endpoint=endpoint, + request_query_params=None, + litellm_params={"model_id": model_id} + ) + + url_str = str(url) + assert model_id in url_str + assert "/converse" in url_str + assert model not in url_str + + +def test_bedrock_passthrough_without_model_id_backward_compatibility(): + """ + Test that passthrough still works without model_id (backward compatibility). + + When model_id is not provided, the system should use the model name as before. + """ + config = BedrockPassthroughConfig() + + model = "anthropic.claude-3-sonnet" + endpoint = f"model/{model}/invoke" + + with patch.object(config, '_get_aws_region_name', return_value="us-east-1"), \ + patch.object(config, 'get_runtime_endpoint', return_value=( + "https://bedrock-runtime.us-east-1.amazonaws.com", + "https://bedrock-runtime.us-east-1.amazonaws.com" + )): + + url, api_base = config.get_complete_url( + api_base=None, + api_key=None, + model=model, + endpoint=endpoint, + request_query_params=None, + litellm_params={} # No model_id provided + ) + + # Verify that the URL contains the model name (not replaced) + url_str = str(url) + assert model in url_str, f"Expected model name in URL when model_id not provided, but got: {url_str}" + expected_url = f"https://bedrock-runtime.us-east-1.amazonaws.com/model/{model}/invoke" + assert url_str == expected_url + + +def test_bedrock_passthrough_region_extraction_from_inference_profile_arn(): + """Test that AWS region is correctly extracted from Application Inference Profile ARN""" + config = BedrockPassthroughConfig() + + model = "anthropic.claude-sonnet-4-20250514-v1:0" + # ARN contains us-west-2 region + model_id = "arn:aws:bedrock:us-west-2:123456789:application-inference-profile/test123" + endpoint = f"model/{model}/invoke" + + # Don't provide aws_region_name in litellm_params to test ARN extraction + with patch.object(config, 'get_runtime_endpoint', return_value=( + "https://bedrock-runtime.us-west-2.amazonaws.com", + "https://bedrock-runtime.us-west-2.amazonaws.com" + )): + + url, api_base = config.get_complete_url( + api_base=None, + api_key=None, + model=model, + endpoint=endpoint, + request_query_params=None, + litellm_params={"model_id": model_id} # Region should be extracted from ARN + ) + + # Verify that the region from ARN is used in the base URL + assert "us-west-2" in api_base, f"Expected region 'us-west-2' from ARN in base URL, but got: {api_base}" +