scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 109

fix application inference profile for pass-through endpoints for bedrock

Browse source
This commit is contained in:
Fang Gong 2025-08-20 00:20:38 -07:00
parent 4df07a5060
commit caad9b3ca7
6 changed files with 274 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
litellm/passthrough/main.py
View file

@ -25,6 +25,8 @@ from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
from litellm.utils import client
from litellm.proxy.pass_through_endpoints.common_utils import encode_bedrock_runtime_modelid_arn
import httpx
base_llm_http_handler = BaseLLMHTTPHandler()
from .utils import BasePassthroughUtils
@ -241,6 +243,12 @@ def llm_passthrough_route(
request_query_params=request_query_params,
litellm_params=litellm_params_dict,
)
# need to encode the id of application-inference-profile for bedrock
if custom_llm_provider == "bedrock" and "application-inference-profile" in endpoint:
encoded_url_str = encode_bedrock_runtime_modelid_arn(str(updated_url))
updated_url = httpx.URL(encoded_url_str)
# Add or update query parameters
provider_api_key = provider_config.get_api_key(api_key)

52
litellm/proxy/pass_through_endpoints/common_utils.py
View file

@ -14,3 +14,55 @@ def get_litellm_virtual_key(request: Request) -> str:
if litellm_api_key:
return f"Bearer {litellm_api_key}"
return request.headers.get("Authorization", "")
def encode_bedrock_runtime_modelid_arn(endpoint: str) -> str:
"""
Encodes any "/" found in the modelId of an AWS Bedrock Runtime Endpoint when arns are passed in.
- modelID value can be an ARN which contains slashes that SHOULD NOT be treated as path separators.
e.g endpoint: /model/<modelId>/invoke
<modelId> containing arns with slashes need to be encoded from
arn:aws:bedrock:ap-southeast-1:123456789012:application-inference-profile/abdefg12334 =>
arn:aws:bedrock:ap-southeast-1:123456789012:application-inference-profile%2Fabdefg12334
so that it is treated as one part of the path.
Otherwise, the encoded endpoint will return 500 error when passed to Bedrock endpoint.
See the apis in https://docs.aws.amazon.com/bedrock/latest/APIReference/API_Operations_Amazon_Bedrock_Runtime.html
for more details on the regex patterns of modelId which we use in the regex logic below.
Args:
endpoint (str): The original endpoint string which may contain ARNs that contain slashes.
Returns:
str: The endpoint with properly encoded ARN slashes
"""
import re
# Early exit: if no ARN detected, return unchanged
if 'arn:aws:' not in endpoint:
return endpoint
# Handle all patterns in one go - more efficient and cleaner
patterns = [
# Custom model with 2 slashes (order matters - do this first)
(r'(custom-model)/([a-z0-9.-]+)/([a-z0-9]+)', r'\1%2F\2%2F\3'),
# All other resource types with 1 slash
(r'(:application-inference-profile)/', r'\1%2F'),
(r'(:inference-profile)/', r'\1%2F'),
(r'(:foundation-model)/', r'\1%2F'),
(r'(:imported-model)/', r'\1%2F'),
(r'(:provisioned-model)/', r'\1%2F'),
(r'(:prompt)/', r'\1%2F'),
(r'(:endpoint)/', r'\1%2F'),
(r'(:prompt-router)/', r'\1%2F'),
(r'(:default-prompt-router)/', r'\1%2F'),
]
for pattern, replacement in patterns:
# Check if pattern exists before applying regex (early exit optimization)
if re.search(pattern, endpoint):
endpoint = re.sub(pattern, replacement, endpoint)
break # Exit after first match since each ARN has only one resource type
return endpoint

14
litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
View file

@ -25,7 +25,7 @@ from litellm.proxy.common_utils.http_parsing_utils import (
get_form_data,
get_request_body,
)
from litellm.proxy.pass_through_endpoints.common_utils import get_litellm_virtual_key
from litellm.proxy.pass_through_endpoints.common_utils import get_litellm_virtual_key, encode_bedrock_runtime_modelid_arn
from litellm.proxy.pass_through_endpoints.pass_through_endpoints import (
create_pass_through_route,
)
@ -492,7 +492,12 @@ async def bedrock_llm_proxy_route(
data: Dict[str, Any] = {}
base_llm_response_processor = ProxyBaseLLMRequestProcessing(data=data)
try:
model = endpoint.split("/")[1]
endpoint_parts = endpoint.split("/")
if "application-inference-profile" in endpoint:
# For application-inference-profile, include the profile ID part as well
model = "/".join(endpoint_parts[1:3])
else:
model = endpoint_parts[1]
except Exception:
raise HTTPException(
status_code=400,
@ -500,12 +505,13 @@ async def bedrock_llm_proxy_route(
"error": "Model missing from endpoint. Expected format: /model/<Model>/<endpoint>. Got: "
+ endpoint,
},
)
)
data["method"] = request.method
data["endpoint"] = endpoint
data["data"] = request_body
data["custom_llm_provider"] = "bedrock"
try:
result = await base_llm_response_processor.base_passthrough_process_llm_request(
request=request,

91
tests/test_litellm/passthrough/test_passthrough_main.py
View file

@ -4,7 +4,10 @@ import sys
import pytest
from fastapi.testclient import TestClient
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from unittest.mock import MagicMock, patch
import httpx
sys.path.insert(
0, os.path.abspath("../../..")
) # Adds the parent directory to the system path
@ -45,3 +48,89 @@ def test_llm_passthrough_route():
assert response.status_code == 200
assert response.json == {"message": "Hello, world!"}
def test_bedrock_application_inference_profile_url_encoding():
client = HTTPHandler()
mock_provider_config = MagicMock()
mock_provider_config.get_complete_url.return_value = (
httpx.URL("https://bedrock-runtime.us-east-1.amazonaws.com/model/arn:aws:bedrock:us-east-1:123456789123:application-inference-profile/r742sbn2zckd/converse"),
"https://bedrock-runtime.us-east-1.amazonaws.com"
)
mock_provider_config.get_api_key.return_value = "test-key"
mock_provider_config.validate_environment.return_value = {}
mock_provider_config.sign_request.return_value = ({}, None)
mock_provider_config.is_streaming_request.return_value = False
with patch("litellm.utils.ProviderConfigManager.get_provider_passthrough_config", return_value=mock_provider_config), \
patch("litellm.litellm_core_utils.get_litellm_params.get_litellm_params", return_value={}), \
patch("litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("test-model", "bedrock", "test-key", "test-base")), \
patch.object(client.client, "send", return_value=MagicMock(status_code=200)) as mock_send, \
patch.object(client.client, "build_request") as mock_build_request:
# Mock logging object
mock_logging_obj = MagicMock()
mock_logging_obj.update_environment_variables = MagicMock()
response = llm_passthrough_route(
model="arn:aws:bedrock:us-east-1:123456789123:application-inference-profile/r742sbn2zckd",
endpoint="model/arn:aws:bedrock:us-east-1:123456789123:application-inference-profile/r742sbn2zckd/converse",
method="POST",
custom_llm_provider="bedrock",
client=client,
litellm_logging_obj=mock_logging_obj,
)
# Verify that build_request was called with the encoded URL
mock_build_request.assert_called_once()
call_args = mock_build_request.call_args
# The URL should have the application-inference-profile ID encoded
actual_url = str(call_args.kwargs["url"])
assert "application-inference-profile%2Fr742sbn2zckd" in actual_url
assert response.status_code == 200
def test_bedrock_non_application_inference_profile_no_encoding():
client = HTTPHandler()
# Mock the provider config and its methods
mock_provider_config = MagicMock()
mock_provider_config.get_complete_url.return_value = (
httpx.URL("https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-3-sonnet-20240229-v1:0/converse"),
"https://bedrock-runtime.us-east-1.amazonaws.com"
)
mock_provider_config.get_api_key.return_value = "test-key"
mock_provider_config.validate_environment.return_value = {}
mock_provider_config.sign_request.return_value = ({}, None)
mock_provider_config.is_streaming_request.return_value = False
with patch("litellm.utils.ProviderConfigManager.get_provider_passthrough_config", return_value=mock_provider_config), \
patch("litellm.litellm_core_utils.get_litellm_params.get_litellm_params", return_value={}), \
patch("litellm.litellm_core_utils.get_llm_provider_logic.get_llm_provider", return_value=("test-model", "bedrock", "test-key", "test-base")), \
patch.object(client.client, "send", return_value=MagicMock(status_code=200)) as mock_send, \
patch.object(client.client, "build_request") as mock_build_request:
# Mock logging object
mock_logging_obj = MagicMock()
mock_logging_obj.update_environment_variables = MagicMock()
response = llm_passthrough_route(
model="anthropic.claude-3-sonnet-20240229-v1:0",
endpoint="model/anthropic.claude-3-sonnet-20240229-v1:0/converse",
method="POST",
custom_llm_provider="bedrock",
client=client,
litellm_logging_obj=mock_logging_obj,
)
# Verify that build_request was called with the original URL (no encoding)
mock_build_request.assert_called_once()
call_args = mock_build_request.call_args
# The URL should NOT have application-inference-profile encoding
actual_url = str(call_args.kwargs["url"])
assert "application-inference-profile%2F" not in actual_url
assert "anthropic.claude-3-sonnet-20240229-v1:0" in actual_url
assert response.status_code == 200

61
tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py
View file

@ -21,6 +21,7 @@ from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
create_pass_through_route,
vertex_discovery_proxy_route,
vertex_proxy_route,
bedrock_llm_proxy_route,
)
from litellm.types.passthrough_endpoints.vertex_ai import VertexPassThroughCredentials
@ -853,3 +854,63 @@ async def test_is_streaming_request_fn():
mock_request.headers = {"content-type": "multipart/form-data"}
mock_request.form = AsyncMock(return_value={"stream": "true"})
assert await is_streaming_request_fn(mock_request) is True
class TestBedrockLLMProxyRoute:
@pytest.mark.asyncio
async def test_bedrock_llm_proxy_route_application_inference_profile(self):
mock_request = Mock()
mock_request.method = "POST"
mock_response = Mock()
mock_user_api_key_dict = Mock()
mock_request_body = {"messages": [{"role": "user", "content": "test"}]}
mock_processor = Mock()
mock_processor.base_passthrough_process_llm_request = AsyncMock(return_value="success")
with patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._read_request_body", return_value=mock_request_body), \
patch("litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing", return_value=mock_processor):
# Test application-inference-profile endpoint
endpoint = "model/arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/r742sbn2zckd/converse"
result = await bedrock_llm_proxy_route(
endpoint=endpoint,
request=mock_request,
fastapi_response=mock_response,
user_api_key_dict=mock_user_api_key_dict,
)
mock_processor.base_passthrough_process_llm_request.assert_called_once()
call_kwargs = mock_processor.base_passthrough_process_llm_request.call_args.kwargs
# For application-inference-profile, model should be "arn:aws:bedrock:us-east-1:026090525607:application-inference-profile/r742sbn2zckd"
assert call_kwargs["model"] == "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/r742sbn2zckd"
assert result == "success"
@pytest.mark.asyncio
async def test_bedrock_llm_proxy_route_regular_model(self):
mock_request = Mock()
mock_request.method = "POST"
mock_response = Mock()
mock_user_api_key_dict = Mock()
mock_request_body = {"messages": [{"role": "user", "content": "test"}]}
mock_processor = Mock()
mock_processor.base_passthrough_process_llm_request = AsyncMock(return_value="success")
with patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._read_request_body", return_value=mock_request_body), \
patch("litellm.proxy.common_request_processing.ProxyBaseLLMRequestProcessing", return_value=mock_processor):
# Test regular model endpoint
endpoint = "model/anthropic.claude-3-sonnet-20240229-v1:0/converse"
result = await bedrock_llm_proxy_route(
endpoint=endpoint,
request=mock_request,
fastapi_response=mock_response,
user_api_key_dict=mock_user_api_key_dict,
)
mock_processor.base_passthrough_process_llm_request.assert_called_once()
call_kwargs = mock_processor.base_passthrough_process_llm_request.call_args.kwargs
# For regular models, model should be just the model ID
assert call_kwargs["model"] == "anthropic.claude-3-sonnet-20240229-v1:0"
assert result == "success"

53
tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_endpoints_common_utils.py
View file

@ -9,6 +9,7 @@ import httpx
import pytest
from fastapi import Request, Response
from fastapi.testclient import TestClient
from litellm.proxy.pass_through_endpoints.common_utils import encode_bedrock_runtime_modelid_arn
sys.path.insert(
0, os.path.abspath("../../../..")
@ -42,3 +43,55 @@ async def test_get_litellm_virtual_key():
}
result = get_litellm_virtual_key(mock_request)
assert result == "Bearer test-key-123"
def test_encode_bedrock_runtime_modelid_arn():
# Test application-inference-profile ARN
endpoint = "model/arn:aws:bedrock:us-east-1:123456789123:application-inference-profile/r742sbn2zckd/converse"
expected = "model/arn:aws:bedrock:us-east-1:123456789123:application-inference-profile%2Fr742sbn2zckd/converse"
result = encode_bedrock_runtime_modelid_arn(endpoint)
assert result == expected
# Test inference-profile ARN
endpoint = "model/arn:aws:bedrock:us-east-1:123456789012:inference-profile/test-profile/invoke"
expected = "model/arn:aws:bedrock:us-east-1:123456789012:inference-profile%2Ftest-profile/invoke"
result = encode_bedrock_runtime_modelid_arn(endpoint)
assert result == expected
# Test foundation-model ARN
endpoint = "model/arn:aws:bedrock:us-east-1:123456789012:foundation-model/anthropic.claude-3/converse"
expected = "model/arn:aws:bedrock:us-east-1:123456789012:foundation-model%2Fanthropic.claude-3/converse"
result = encode_bedrock_runtime_modelid_arn(endpoint)
assert result == expected
# Test custom-model ARN (2 slashes)
endpoint = "model/arn:aws:bedrock:us-east-1:123456789012:custom-model/my-model.fine-tuned/abc123/invoke"
expected = "model/arn:aws:bedrock:us-east-1:123456789012:custom-model%2Fmy-model.fine-tuned%2Fabc123/invoke"
result = encode_bedrock_runtime_modelid_arn(endpoint)
assert result == expected
# Test provisioned-model ARN
endpoint = "model/arn:aws:bedrock:us-east-1:123456789012:provisioned-model/test-model/converse"
expected = "model/arn:aws:bedrock:us-east-1:123456789012:provisioned-model%2Ftest-model/converse"
result = encode_bedrock_runtime_modelid_arn(endpoint)
assert result == expected
def test_encode_bedrock_runtime_modelid_arn_no_arn():
# Test regular model ID (no ARN)
endpoint = "model/anthropic.claude-3-sonnet-20240229-v1:0/converse"
result = encode_bedrock_runtime_modelid_arn(endpoint)
assert result == endpoint
def test_encode_bedrock_runtime_modelid_arn_edge_cases():
# Test multiple ARN types (should only encode first match)
endpoint = "model/arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/test1/converse"
expected = "model/arn:aws:bedrock:us-east-1:123456789012:application-inference-profile%2Ftest1/converse"
result = encode_bedrock_runtime_modelid_arn(endpoint)
assert result == expected
# Test ARN with special characters in resource ID
endpoint = "model/arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/test-profile.v1/invoke"
expected = "model/arn:aws:bedrock:us-east-1:123456789012:application-inference-profile%2Ftest-profile.v1/invoke"
result = encode_bedrock_runtime_modelid_arn(endpoint)
assert result == expected