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

test fix location of encode bedrock runtime modelid arn

Browse source
This commit is contained in:
Ishaan Jaff 2025-08-22 17:33:58 -07:00
parent 757274a193
commit 3a75e764fc
4 changed files with 65 additions and 62 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
litellm/passthrough/main.py
View file

@ -24,8 +24,8 @@ import litellm
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.passthrough.utils import CommonUtils
from litellm.utils import client
from litellm.proxy.pass_through_endpoints.common_utils import encode_bedrock_runtime_modelid_arn
base_llm_http_handler = BaseLLMHTTPHandler()
from .utils import BasePassthroughUtils
@ -245,7 +245,7 @@ def llm_passthrough_route(
# 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))
encoded_url_str = CommonUtils.encode_bedrock_runtime_modelid_arn(str(updated_url))
updated_url = httpx.URL(encoded_url_str)
# Add or update query parameters

53
litellm/passthrough/utils.py
View file

@ -37,3 +37,56 @@ class BasePassthroughUtils:
# Combine request headers with custom headers
headers = {**request_headers, **headers}
return headers
class CommonUtils:
@staticmethod
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

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

@ -15,54 +15,3 @@ def get_litellm_virtual_key(request: Request) -> str:
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

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

@ -9,7 +9,8 @@ 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
from litellm.passthrough.utils import CommonUtils
sys.path.insert(
0, os.path.abspath("../../../..")
@ -48,38 +49,38 @@ 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)
result = CommonUtils.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)
result = CommonUtils.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)
result = CommonUtils.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)
result = CommonUtils.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)
result = CommonUtils.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)
result = CommonUtils.encode_bedrock_runtime_modelid_arn(endpoint)
assert result == endpoint
@ -87,11 +88,11 @@ 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)
result = CommonUtils.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)
result = CommonUtils.encode_bedrock_runtime_modelid_arn(endpoint)
assert result == expected