From bfa297e82f0198435816397c9a8182c5fbabc7df Mon Sep 17 00:00:00 2001
From: canada4663 <3118399+canada4663@users.noreply.github.com>
Date: Sat, 28 Oct 2023 22:25:45 -0700
Subject: [PATCH] add ability to use custom AWS Bedrock Runtime endpoint via
 environment or .env variable AWS_BEDROCK_RUNTIME_ENDPOINT or as a parameter
 to init_bedrock_client

---
 litellm/llms/bedrock.py | 36 +++++++++++++++++++++++++++---------
 1 file changed, 27 insertions(+), 9 deletions(-)

diff --git a/litellm/llms/bedrock.py b/litellm/llms/bedrock.py
index 62760a48c89..0d2a75faf0f 100644
--- a/litellm/llms/bedrock.py
+++ b/litellm/llms/bedrock.py
@@ -175,8 +175,32 @@ def init_bedrock_client(
         aws_access_key_id = None,
         aws_secret_access_key = None,
         aws_region_name=None,
+        aws_bedrock_runtime_endpoint=None,
     ):
 
+    # check for custom AWS_REGION_NAME and use it if not passed to init_bedrock_client
+    litellm_aws_region_name = get_secret("AWS_REGION_NAME")
+    standard_aws_region_name = get_secret("AWS_REGION")
+    if region_name:
+        pass
+    elif aws_region_name:
+        region_name = aws_region_name
+    elif litellm_aws_region_name:
+        region_name = litellm_aws_region_name
+    elif standard_aws_region_name:
+        region_name = standard_aws_region_name
+    else:
+        raise BedrockError(message="AWS region not set: set AWS_REGION_NAME or AWS_REGION env variable or in .env file")
+
+    # check for custom AWS_BEDROCK_RUNTIME_ENDPOINT and use it if not passed to init_bedrock_client
+    env_aws_bedrock_runtime_endpoint = get_secret("AWS_BEDROCK_RUNTIME_ENDPOINT")
+    if aws_bedrock_runtime_endpoint:
+        endpoint_url = aws_bedrock_runtime_endpoint
+    elif env_aws_bedrock_runtime_endpoint:
+        endpoint_url = env_aws_bedrock_runtime_endpoint
+    else:
+        endpoint_url = f'https://bedrock-runtime.{region_name}.amazonaws.com'
+
     import boto3
     if aws_access_key_id != None:
         # uses auth params passed to completion
@@ -186,23 +210,17 @@ def init_bedrock_client(
             service_name="bedrock-runtime",
             aws_access_key_id=aws_access_key_id,
             aws_secret_access_key=aws_secret_access_key,
-            region_name=aws_region_name,
-            endpoint_url=f'https://bedrock-runtime.{aws_region_name}.amazonaws.com'
+            region_name=region_name,
+            endpoint_url=endpoint_url,
         )
     else:
         # aws_access_key_id is None, assume user is trying to auth using env variables 
         # boto3 automatically reads env variables
 
-        # we need to read region name from env
-        # I assume majority of users use .env for auth
-        region_name = get_secret("AWS_REGION_NAME") # reads env for AWS_REGION_NAME, defaults to None
-        # if litellm config `AWS_REGION_NAME` is not set, attempt to read AWS default env variable name `AWS_REGION`
-        if not region_name:
-            region_name = get_secret("AWS_REGION")
         client = boto3.client(
             service_name="bedrock-runtime",
             region_name=region_name,
-            endpoint_url=f'https://bedrock-runtime.{region_name}.amazonaws.com'
+            endpoint_url=endpoint_url,
         )
 
     return client