From 2f130b39c01f9df811c6c2e2a63ce2c50df3fac6 Mon Sep 17 00:00:00 2001 From: uid10804 Date: Fri, 31 May 2024 12:02:34 +0200 Subject: [PATCH] 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. --- litellm/llms/bedrock_httpx.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/litellm/llms/bedrock_httpx.py b/litellm/llms/bedrock_httpx.py index 337055dc2e..55bd6d7390 100644 --- a/litellm/llms/bedrock_httpx.py +++ b/litellm/llms/bedrock_httpx.py @@ -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)