scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

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

Browse source
This commit is contained in:
canada4663 2023-10-28 22:25:45 -07:00
parent 1add0d604e
commit bfa297e82f
1 changed files with 27 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

36
litellm/llms/bedrock.py
View file

@ -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