fix: add botocore credentials extraction and conversion

The code changes in `bedrock_httpx.py` add functionality to extract and convert AWS STS credentials to session credentials using botocore. This fixes later error in add_auth request when token needs to be assigned.
This commit is contained in:
uid10804 2024-05-31 12:02:34 +02:00
parent deb87f71e3
commit 2f130b39c0

View file

@ -45,6 +45,7 @@ import httpx # type: ignore
from .bedrock import BedrockError, convert_messages_to_prompt, ModelResponseIterator
from litellm.types.llms.bedrock import *
import urllib.parse
import botocore
class AmazonCohereChatConfig:
@ -261,7 +262,14 @@ class BedrockLLM(BaseLLM):
RoleArn=aws_role_name, RoleSessionName=aws_session_name
)
return sts_response["Credentials"]
# Extract the credentials from the response and convert to Session Credentials
sts_credentials = sts_response["Credentials"]
credentials = botocore.credentials.Credentials(
access_key=sts_credentials["AccessKeyId"],
secret_key=sts_credentials["SecretAccessKey"],
token=sts_credentials["SessionToken"],
)
return credentials
elif aws_profile_name is not None: ### CHECK SESSION ###
# uses auth values from AWS profile usually stored in ~/.aws/credentials
client = boto3.Session(profile_name=aws_profile_name)