fix: add safe dictionary access for bedrock credentials (#11860)

This commit is contained in:
Dr.MerdanBay 2024-12-20 12:13:39 +09:00 committed by GitHub
parent 463fbe2680
commit bb2f46d7cc
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 16 additions and 5 deletions

View File

@ -1,11 +1,19 @@
from collections.abc import Mapping
import boto3
from botocore.config import Config
from core.model_runtime.errors.invoke import InvokeBadRequestError
def get_bedrock_client(service_name: str, credentials: Mapping[str, str]):
region_name = credentials.get("aws_region")
if not region_name:
raise InvokeBadRequestError("aws_region is required")
client_config = Config(region_name=region_name)
aws_access_key_id = credentials.get("aws_access_key_id")
aws_secret_access_key = credentials.get("aws_secret_access_key")
def get_bedrock_client(service_name, credentials=None):
client_config = Config(region_name=credentials["aws_region"])
aws_access_key_id = credentials["aws_access_key_id"]
aws_secret_access_key = credentials["aws_secret_access_key"]
if aws_access_key_id and aws_secret_access_key:
# use aksk to call bedrock
client = boto3.client(

View File

@ -62,7 +62,10 @@ class BedrockRerankModel(RerankModel):
}
)
modelId = model
region = credentials["aws_region"]
region = credentials.get("aws_region")
# region is a required field
if not region:
raise InvokeBadRequestError("aws_region is required in credentials")
model_package_arn = f"arn:aws:bedrock:{region}::foundation-model/{modelId}"
rerankingConfiguration = {
"type": "BEDROCK_RERANKING_MODEL",