mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-11 17:05:43 +00:00
Fix : model id encoding for bedrock passthrough
This commit is contained in:
@@ -24,6 +24,37 @@ class BedrockPassthroughConfig(
|
||||
def is_streaming_request(self, endpoint: str, request_data: dict) -> bool:
|
||||
return "stream" in endpoint
|
||||
|
||||
def _encode_model_id_for_endpoint(self, model_id: str) -> str:
|
||||
"""
|
||||
Encode model_id (especially ARNs) for use in Bedrock endpoints.
|
||||
|
||||
ARNs contain special characters like colons and slashes that need to be
|
||||
properly URL-encoded when used in HTTP request paths. For example:
|
||||
arn:aws:bedrock:us-east-1:123:application-inference-profile/abc123
|
||||
becomes:
|
||||
arn:aws:bedrock:us-east-1:123:application-inference-profile%2Fabc123
|
||||
|
||||
Args:
|
||||
model_id: The model ID or ARN to encode
|
||||
|
||||
Returns:
|
||||
The encoded model_id suitable for use in endpoint URLs
|
||||
"""
|
||||
from litellm.passthrough.utils import CommonUtils
|
||||
import re
|
||||
|
||||
# Create a temporary endpoint with the model_id to check if encoding is needed
|
||||
temp_endpoint = f"/model/{model_id}/converse"
|
||||
encoded_temp_endpoint = CommonUtils.encode_bedrock_runtime_modelid_arn(temp_endpoint)
|
||||
|
||||
# Extract the encoded model_id from the temporary endpoint
|
||||
encoded_model_id_match = re.search(r'/model/([^/]+)/', encoded_temp_endpoint)
|
||||
if encoded_model_id_match:
|
||||
return encoded_model_id_match.group(1)
|
||||
else:
|
||||
# Fallback to original model_id if extraction fails
|
||||
return model_id
|
||||
|
||||
def get_complete_url(
|
||||
self,
|
||||
api_base: Optional[str],
|
||||
@@ -53,9 +84,13 @@ class BedrockPassthroughConfig(
|
||||
# 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)
|
||||
|
||||
# Encode the model_id if it's an ARN to properly handle special characters
|
||||
encoded_model_id = self._encode_model_id_for_endpoint(model_id)
|
||||
|
||||
# Replace the model name in the endpoint with the encoded model_id
|
||||
endpoint = re.sub(r'model/[^/]+/', f'model/{encoded_model_id}/', endpoint)
|
||||
return self.format_url(endpoint, endpoint_url, request_query_params or {}), endpoint_url
|
||||
|
||||
def sign_request(
|
||||
|
||||
+132
-7
@@ -181,7 +181,7 @@ def test_bedrock_passthrough_with_application_inference_profile():
|
||||
|
||||
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.
|
||||
replace the translated model name in the endpoint URL and be properly encoded.
|
||||
"""
|
||||
config = BedrockPassthroughConfig()
|
||||
|
||||
@@ -204,19 +204,21 @@ def test_bedrock_passthrough_with_application_inference_profile():
|
||||
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
|
||||
# Verify that the URL contains the encoded 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}"
|
||||
# The ARN slash should be encoded as %2F
|
||||
assert "application-inference-profile%2F" in url_str, f"Expected encoded 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"
|
||||
# Verify the complete URL structure with encoded ARN
|
||||
encoded_model_id = "arn:aws:bedrock:eu-west-1:123456789:application-inference-profile%2Fabcdefgh1234"
|
||||
expected_url = f"https://bedrock-runtime.eu-west-1.amazonaws.com/model/{encoded_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"""
|
||||
"""Test Application Inference Profile with converse endpoint and proper ARN encoding"""
|
||||
config = BedrockPassthroughConfig()
|
||||
|
||||
model = "anthropic.claude-sonnet-4-20250514-v1:0"
|
||||
@@ -239,7 +241,8 @@ def test_bedrock_passthrough_with_inference_profile_converse_endpoint():
|
||||
)
|
||||
|
||||
url_str = str(url)
|
||||
assert model_id in url_str
|
||||
# The ARN should be encoded with %2F
|
||||
assert "application-inference-profile%2F" in url_str
|
||||
assert "/converse" in url_str
|
||||
assert model not in url_str
|
||||
|
||||
@@ -304,3 +307,125 @@ def test_bedrock_passthrough_region_extraction_from_inference_profile_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}"
|
||||
|
||||
|
||||
def test_bedrock_passthrough_model_id_arn_encoding():
|
||||
"""
|
||||
Test that model_id ARNs are properly URL-encoded when used in endpoints.
|
||||
|
||||
This is the critical fix for the issue where ARNs with slashes need to be encoded
|
||||
so they're treated as a single path component rather than multiple path segments.
|
||||
|
||||
For example:
|
||||
arn:aws:bedrock:us-east-1:590183661440:application-inference-profile/b943q2qbl3m7
|
||||
should become:
|
||||
arn:aws:bedrock:us-east-1:590183661440:application-inference-profile%2Fb943q2qbl3m7
|
||||
"""
|
||||
config = BedrockPassthroughConfig()
|
||||
|
||||
model = "bedrock-claude-4-5-sonnet"
|
||||
# ARN with a slash that needs encoding
|
||||
model_id = "arn:aws:bedrock:us-east-1:590183661440:application-inference-profile/b943q2qbl3m7"
|
||||
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)
|
||||
|
||||
# The slash in the ARN after application-inference-profile should be encoded as %2F
|
||||
assert "application-inference-profile%2F" in url_str, \
|
||||
f"Expected encoded ARN with %2F in URL, but got: {url_str}"
|
||||
|
||||
# The unencoded version should NOT be in the URL
|
||||
assert "application-inference-profile/" not in url_str, \
|
||||
f"ARN slash should be encoded, but found unencoded version in: {url_str}"
|
||||
|
||||
# Verify the complete expected URL structure
|
||||
expected_encoded_model_id = "arn:aws:bedrock:us-east-1:590183661440:application-inference-profile%2Fb943q2qbl3m7"
|
||||
expected_url = f"https://bedrock-runtime.us-east-1.amazonaws.com/model/{expected_encoded_model_id}/converse"
|
||||
assert url_str == expected_url, f"Expected {expected_url}, but got: {url_str}"
|
||||
|
||||
|
||||
def test_bedrock_passthrough_model_id_arn_encoding_invoke_endpoint():
|
||||
"""
|
||||
Test ARN encoding with /invoke endpoint (not just /converse).
|
||||
"""
|
||||
config = BedrockPassthroughConfig()
|
||||
|
||||
model = "anthropic.claude-sonnet-4-5-20250929-v1:0"
|
||||
model_id = "arn:aws:bedrock:us-east-1:123456789:application-inference-profile/xyz789"
|
||||
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={"model_id": model_id}
|
||||
)
|
||||
|
||||
url_str = str(url)
|
||||
|
||||
# Verify encoding
|
||||
assert "application-inference-profile%2F" in url_str
|
||||
assert "/invoke" in url_str
|
||||
|
||||
expected_encoded_model_id = "arn:aws:bedrock:us-east-1:123456789:application-inference-profile%2Fxyz789"
|
||||
expected_url = f"https://bedrock-runtime.us-east-1.amazonaws.com/model/{expected_encoded_model_id}/invoke"
|
||||
assert url_str == expected_url
|
||||
|
||||
|
||||
def test_bedrock_passthrough_model_id_without_arn():
|
||||
"""
|
||||
Test that non-ARN model_ids (regular model IDs) are not affected by encoding logic.
|
||||
"""
|
||||
config = BedrockPassthroughConfig()
|
||||
|
||||
model = "my-model"
|
||||
# Regular model ID (not an ARN)
|
||||
model_id = "us.anthropic.claude-3-5-sonnet-20240620-v1:0"
|
||||
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)
|
||||
|
||||
# Regular model ID should be used as-is (no encoding needed)
|
||||
assert model_id in url_str
|
||||
assert "%2F" not in url_str, "Non-ARN model IDs should not be encoded"
|
||||
|
||||
expected_url = f"https://bedrock-runtime.us-east-1.amazonaws.com/model/{model_id}/converse"
|
||||
assert url_str == expected_url
|
||||
|
||||
|
||||
Reference in New Issue
Block a user