From 3a75e764fca01a12d5c155fde67bff62ae312613 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Fri, 22 Aug 2025 17:33:58 -0700 Subject: [PATCH] test fix location of encode bedrock runtime modelid arn --- litellm/passthrough/main.py | 4 +- litellm/passthrough/utils.py | 53 +++++++++++++++++++ .../pass_through_endpoints/common_utils.py | 51 ------------------ ...test_passthrough_endpoints_common_utils.py | 19 +++---- 4 files changed, 65 insertions(+), 62 deletions(-) diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index 9bde9496d29..f4dc1ef6c84 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -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 diff --git a/litellm/passthrough/utils.py b/litellm/passthrough/utils.py index c52d0e3688d..4bf66d49881 100644 --- a/litellm/passthrough/utils.py +++ b/litellm/passthrough/utils.py @@ -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//invoke + 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 \ No newline at end of file diff --git a/litellm/proxy/pass_through_endpoints/common_utils.py b/litellm/proxy/pass_through_endpoints/common_utils.py index 59d098fdff9..804960cdee9 100644 --- a/litellm/proxy/pass_through_endpoints/common_utils.py +++ b/litellm/proxy/pass_through_endpoints/common_utils.py @@ -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//invoke - 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 \ No newline at end of file diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_endpoints_common_utils.py b/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_endpoints_common_utils.py index 3f0a809bdc1..97ef05100de 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_endpoints_common_utils.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_passthrough_endpoints_common_utils.py @@ -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 \ No newline at end of file