diff --git a/litellm/llms/bedrock_httpx.py b/litellm/llms/bedrock_httpx.py index c895679f16a..dd41a8cb759 100644 --- a/litellm/llms/bedrock_httpx.py +++ b/litellm/llms/bedrock_httpx.py @@ -45,9 +45,7 @@ import httpx # type: ignore from .bedrock import BedrockError, convert_messages_to_prompt, ModelResponseIterator from litellm.types.llms.bedrock import * import urllib.parse -from litellm.caching import DualCache -iam_cache = DualCache() class AmazonCohereChatConfig: """ @@ -287,51 +285,35 @@ class BedrockLLM(BaseLLM): ### CHECK STS ### if aws_web_identity_token is not None and aws_role_name is not None and aws_session_name is not None: - iam_creds_cache_key = json.dumps({ - "aws_web_identity_token": aws_web_identity_token, - "aws_role_name": aws_role_name, - "aws_session_name": aws_session_name, - "aws_region_name": aws_region_name, - }) + oidc_token = get_secret(aws_web_identity_token) - iam_creds_dict = iam_cache.get_cache(iam_creds_cache_key) - if iam_creds_dict is None: - oidc_token = get_secret(aws_web_identity_token) - - if oidc_token is None: - raise BedrockError( - message="OIDC token could not be retrieved from secret manager.", - status_code=401, - ) - - sts_client = boto3.client( - "sts", - region_name=aws_region_name, + if oidc_token is None: + raise BedrockError( + message="OIDC token could not be retrieved from secret manager.", + status_code=401, ) - # https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRoleWithWebIdentity.html - # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sts/client/assume_role_with_web_identity.html - sts_response = sts_client.assume_role_with_web_identity( - RoleArn=aws_role_name, - RoleSessionName=aws_session_name, - WebIdentityToken=oidc_token, - DurationSeconds=3600, - ) + sts_client = boto3.client( + "sts" + ) - iam_creds_dict = { - "aws_access_key_id": sts_response["Credentials"]["AccessKeyId"], - "aws_secret_access_key": sts_response["Credentials"]["SecretAccessKey"], - "aws_session_token": sts_response["Credentials"]["SessionToken"], - "region_name": aws_region_name, - } + # https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRoleWithWebIdentity.html + # https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/sts/client/assume_role_with_web_identity.html + sts_response = sts_client.assume_role_with_web_identity( + RoleArn=aws_role_name, + RoleSessionName=aws_session_name, + WebIdentityToken=oidc_token, + DurationSeconds=3600, + ) - iam_cache.set_cache(key=iam_creds_cache_key, value=json.dumps(iam_creds_dict), ttl=3600 - 60) + session = boto3.Session( + aws_access_key_id=sts_response["Credentials"]["AccessKeyId"], + aws_secret_access_key=sts_response["Credentials"]["SecretAccessKey"], + aws_session_token=sts_response["Credentials"]["SessionToken"], + region_name=aws_region_name, + ) - session = boto3.Session(**iam_creds_dict) - - iam_creds = session.get_credentials() - - return iam_creds + return session.get_credentials() elif aws_role_name is not None and aws_session_name is not None: sts_client = boto3.client( "sts",