Merge pull request #1203 from Manouchehri/bedrock-cloudflare-ai-gateway-1

Add aws_bedrock_runtime_endpoint support
This commit is contained in:
Krish Dholakia 2023-12-23 11:44:04 +05:30 committed by GitHub
commit 03fd5da5ae
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23

View file

@ -327,6 +327,7 @@ def completion(
aws_secret_access_key = optional_params.pop("aws_secret_access_key", None)
aws_access_key_id = optional_params.pop("aws_access_key_id", None)
aws_region_name = optional_params.pop("aws_region_name", None)
aws_bedrock_runtime_endpoint = optional_params.pop("aws_bedrock_runtime_endpoint", None)
# use passed in BedrockRuntime.Client if provided, otherwise create a new one
client = optional_params.pop(
@ -336,6 +337,7 @@ def completion(
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
aws_region_name=aws_region_name,
aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint,
),
)
@ -614,12 +616,14 @@ def embedding(
aws_secret_access_key = optional_params.pop("aws_secret_access_key", None)
aws_access_key_id = optional_params.pop("aws_access_key_id", None)
aws_region_name = optional_params.pop("aws_region_name", None)
aws_bedrock_runtime_endpoint = optional_params.pop("aws_bedrock_runtime_endpoint", None)
# use passed in BedrockRuntime.Client if provided, otherwise create a new one
client = init_bedrock_client(
aws_access_key_id=aws_access_key_id,
aws_secret_access_key=aws_secret_access_key,
aws_region_name=aws_region_name,
aws_bedrock_runtime_endpoint=aws_bedrock_runtime_endpoint,
)
## Embedding Call