From df8043bac9e739b2b140df4d6d62b53e5c17e2c0 Mon Sep 17 00:00:00 2001 From: uid10804 Date: Fri, 31 May 2024 12:02:34 +0200 Subject: [PATCH 1/2] 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) From 559e825e2fbfb5956875dda3a48291fe647aa4cf Mon Sep 17 00:00:00 2001 From: uid10804 Date: Fri, 31 May 2024 16:57:28 +0200 Subject: [PATCH 2/2] refactor(bedrock_httpx): change botocore reference --- litellm/llms/bedrock_httpx.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/litellm/llms/bedrock_httpx.py b/litellm/llms/bedrock_httpx.py index 55bd6d7390..156aa510c3 100644 --- a/litellm/llms/bedrock_httpx.py +++ b/litellm/llms/bedrock_httpx.py @@ -45,7 +45,6 @@ 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: @@ -264,7 +263,9 @@ class BedrockLLM(BaseLLM): # Extract the credentials from the response and convert to Session Credentials sts_credentials = sts_response["Credentials"] - credentials = botocore.credentials.Credentials( + from botocore.credentials import Credentials + + credentials = Credentials( access_key=sts_credentials["AccessKeyId"], secret_key=sts_credentials["SecretAccessKey"], token=sts_credentials["SessionToken"],