From 240a0da2718948a94b2c610c82af5b1caecf6605 Mon Sep 17 00:00:00 2001 From: David Manouchehri Date: Fri, 17 May 2024 10:19:25 +0000 Subject: [PATCH 1/9] feat(bedrock_httpx.py): Add AWS IAM cred caching for OIDC flow. --- litellm/llms/bedrock_httpx.py | 64 ++++++++++++++++++++++------------- 1 file changed, 41 insertions(+), 23 deletions(-) diff --git a/litellm/llms/bedrock_httpx.py b/litellm/llms/bedrock_httpx.py index dd41a8cb75..c895679f16 100644 --- a/litellm/llms/bedrock_httpx.py +++ b/litellm/llms/bedrock_httpx.py @@ -45,7 +45,9 @@ 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: """ @@ -285,35 +287,51 @@ 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: - oidc_token = get_secret(aws_web_identity_token) + 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, + }) - if oidc_token is None: - raise BedrockError( - message="OIDC token could not be retrieved from secret manager.", - status_code=401, + 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, ) - sts_client = boto3.client( - "sts" - ) + # 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, + ) - # 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_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, + } - 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, - ) + iam_cache.set_cache(key=iam_creds_cache_key, value=json.dumps(iam_creds_dict), ttl=3600 - 60) - return session.get_credentials() + session = boto3.Session(**iam_creds_dict) + + iam_creds = session.get_credentials() + + return iam_creds elif aws_role_name is not None and aws_session_name is not None: sts_client = boto3.client( "sts", From fd16937ad80bdf36298b811002f7e0526c8e7414 Mon Sep 17 00:00:00 2001 From: David Manouchehri Date: Sat, 1 Jun 2024 01:52:06 +0000 Subject: [PATCH 2/9] fix(bedrock_httpx.py): Fix STS region endpoint. --- litellm/llms/bedrock_httpx.py | 1 + 1 file changed, 1 insertion(+) diff --git a/litellm/llms/bedrock_httpx.py b/litellm/llms/bedrock_httpx.py index c895679f16..99e2a8ffe4 100644 --- a/litellm/llms/bedrock_httpx.py +++ b/litellm/llms/bedrock_httpx.py @@ -307,6 +307,7 @@ class BedrockLLM(BaseLLM): sts_client = boto3.client( "sts", region_name=aws_region_name, + endpoint_url=f"https://sts.{aws_region_name}.amazonaws.com" ) # https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRoleWithWebIdentity.html From 3410367610abc62f4899158cd2dacc467e6423da Mon Sep 17 00:00:00 2001 From: David Manouchehri Date: Sat, 1 Jun 2024 15:18:32 +0000 Subject: [PATCH 3/9] test(test_bedrock_completion.py): Add tests to ensure caching isn't breaking anything. --- litellm/tests/test_bedrock_completion.py | 42 +++++++++++++++++++++--- 1 file changed, 37 insertions(+), 5 deletions(-) diff --git a/litellm/tests/test_bedrock_completion.py b/litellm/tests/test_bedrock_completion.py index 047f0cb2e2..20596d265c 100644 --- a/litellm/tests/test_bedrock_completion.py +++ b/litellm/tests/test_bedrock_completion.py @@ -220,13 +220,13 @@ def test_completion_bedrock_claude_sts_oidc_auth(): aws_web_identity_token = "oidc/circleci_v2/" aws_region_name = os.environ["AWS_REGION_NAME"] # aws_role_name = os.environ["AWS_TEMP_ROLE_NAME"] - # TODO: This is using David's IAM role, we should use Litellm's IAM role eventually + # TODO: This is using ai.moda's IAM role, we should use LiteLLM's IAM role eventually aws_role_name = "arn:aws:iam::335785316107:role/litellm-github-unit-tests-circleci" try: litellm.set_verbose = True - response = completion( + response_1 = completion( model="bedrock/anthropic.claude-3-haiku-20240307-v1:0", messages=messages, max_tokens=10, @@ -236,8 +236,40 @@ def test_completion_bedrock_claude_sts_oidc_auth(): aws_role_name=aws_role_name, aws_session_name="my-test-session", ) - # Add any assertions here to check the response - print(response) + print(response_1) + assert len(response_1.choices) > 0 + assert len(response_1.choices[0].message.content) > 0 + + # This second call is to verify that the cache isn't breaking anything + response_2 = completion( + model="bedrock/anthropic.claude-3-haiku-20240307-v1:0", + messages=messages, + max_tokens=5, + temperature=0.2, + aws_region_name=aws_region_name, + aws_web_identity_token=aws_web_identity_token, + aws_role_name=aws_role_name, + aws_session_name="my-test-session", + ) + print(response_2) + assert len(response_2.choices) > 0 + assert len(response_2.choices[0].message.content) > 0 + + # This third call is to verify that the cache isn't used for a different region + response_3 = completion( + model="bedrock/anthropic.claude-3-haiku-20240307-v1:0", + messages=messages, + max_tokens=6, + temperature=0.3, + aws_region_name="us-east-1", + aws_web_identity_token=aws_web_identity_token, + aws_role_name=aws_role_name, + aws_session_name="my-test-session", + ) + print(response_3) + assert len(response_3.choices) > 0 + assert len(response_3.choices[0].message.content) > 0 + except RateLimitError: pass except Exception as e: @@ -254,7 +286,7 @@ def test_completion_bedrock_httpx_command_r_sts_oidc_auth(): aws_web_identity_token = "oidc/circleci_v2/" aws_region_name = os.environ["AWS_REGION_NAME"] # aws_role_name = os.environ["AWS_TEMP_ROLE_NAME"] - # TODO: This is using David's IAM role, we should use Litellm's IAM role eventually + # TODO: This is using ai.moda's IAM role, we should use LiteLLM's IAM role eventually aws_role_name = "arn:aws:iam::335785316107:role/litellm-github-unit-tests-circleci" try: From 1a77c1bf169189c86057844c3204b33959e5cf2d Mon Sep 17 00:00:00 2001 From: David Manouchehri Date: Tue, 11 Jun 2024 15:42:09 +0000 Subject: [PATCH 4/9] Revert "fix(bedrock_httpx.py): Fix STS region endpoint." This reverts commit fd16937ad80bdf36298b811002f7e0526c8e7414. --- litellm/llms/bedrock_httpx.py | 1 - 1 file changed, 1 deletion(-) diff --git a/litellm/llms/bedrock_httpx.py b/litellm/llms/bedrock_httpx.py index 99e2a8ffe4..c895679f16 100644 --- a/litellm/llms/bedrock_httpx.py +++ b/litellm/llms/bedrock_httpx.py @@ -307,7 +307,6 @@ class BedrockLLM(BaseLLM): sts_client = boto3.client( "sts", region_name=aws_region_name, - endpoint_url=f"https://sts.{aws_region_name}.amazonaws.com" ) # https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRoleWithWebIdentity.html From 4de9e46d4229964c4852051955013d733bc7912a Mon Sep 17 00:00:00 2001 From: David Manouchehri Date: Tue, 11 Jun 2024 15:42:18 +0000 Subject: [PATCH 5/9] Revert "feat(bedrock_httpx.py): Add AWS IAM cred caching for OIDC flow." This reverts commit 240a0da2718948a94b2c610c82af5b1caecf6605. --- litellm/llms/bedrock_httpx.py | 64 +++++++++++++---------------------- 1 file changed, 23 insertions(+), 41 deletions(-) diff --git a/litellm/llms/bedrock_httpx.py b/litellm/llms/bedrock_httpx.py index c895679f16..dd41a8cb75 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", From cd21c80a15cbdd31b1a2f377c981fa096d8091d5 Mon Sep 17 00:00:00 2001 From: David Manouchehri Date: Fri, 17 May 2024 10:19:25 +0000 Subject: [PATCH 6/9] feat(bedrock_httpx.py): Add AWS IAM cred caching for OIDC flow. --- litellm/llms/bedrock_httpx.py | 68 +++++++++++++++++++++-------------- 1 file changed, 42 insertions(+), 26 deletions(-) diff --git a/litellm/llms/bedrock_httpx.py b/litellm/llms/bedrock_httpx.py index b011d95129..0c1e7c06fa 100644 --- a/litellm/llms/bedrock_httpx.py +++ b/litellm/llms/bedrock_httpx.py @@ -53,7 +53,9 @@ from litellm.types.llms.openai import ( ChatCompletionToolCallFunctionChunk, ChatCompletionDeltaChunk, ) +from litellm.caching import DualCache +iam_cache = DualCache() class AmazonCohereChatConfig: """ @@ -325,38 +327,52 @@ class BedrockLLM(BaseLLM): ) = params_to_check ### CHECK STS ### - if ( - aws_web_identity_token is not None - and aws_role_name is not None - and aws_session_name is not None - ): - oidc_token = get_secret(aws_web_identity_token) + 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, + }) - if oidc_token is None: - raise BedrockError( - message="OIDC token could not be retrieved from secret manager.", - status_code=401, + 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, ) - sts_client = boto3.client("sts") + # 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, + ) - # 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_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, + } - 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, - ) + iam_cache.set_cache(key=iam_creds_cache_key, value=json.dumps(iam_creds_dict), ttl=3600 - 60) - return session.get_credentials() + session = boto3.Session(**iam_creds_dict) + + iam_creds = session.get_credentials() + + return iam_creds elif aws_role_name is not None and aws_session_name is not None: sts_client = boto3.client( "sts", From 7faf0b95861710b81ea11ce9582155ba16016f68 Mon Sep 17 00:00:00 2001 From: David Manouchehri Date: Sat, 1 Jun 2024 01:52:06 +0000 Subject: [PATCH 7/9] fix(bedrock_httpx.py): Fix STS region endpoint. --- litellm/llms/bedrock_httpx.py | 1 + 1 file changed, 1 insertion(+) diff --git a/litellm/llms/bedrock_httpx.py b/litellm/llms/bedrock_httpx.py index 0c1e7c06fa..6d4c04fb1c 100644 --- a/litellm/llms/bedrock_httpx.py +++ b/litellm/llms/bedrock_httpx.py @@ -348,6 +348,7 @@ class BedrockLLM(BaseLLM): sts_client = boto3.client( "sts", region_name=aws_region_name, + endpoint_url=f"https://sts.{aws_region_name}.amazonaws.com" ) # https://docs.aws.amazon.com/STS/latest/APIReference/API_AssumeRoleWithWebIdentity.html From 9aee0b4a0f2d24ae9edfcc5e5825025ac18ec63a Mon Sep 17 00:00:00 2001 From: David Manouchehri Date: Tue, 11 Jun 2024 16:34:32 +0000 Subject: [PATCH 8/9] fix(bedrock_httpx.py): Fix STS region endpoint for converse flow. --- litellm/llms/bedrock_httpx.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/litellm/llms/bedrock_httpx.py b/litellm/llms/bedrock_httpx.py index 6d4c04fb1c..ee0ae57b42 100644 --- a/litellm/llms/bedrock_httpx.py +++ b/litellm/llms/bedrock_httpx.py @@ -1446,7 +1446,11 @@ class BedrockConverseLLM(BaseLLM): status_code=401, ) - sts_client = boto3.client("sts") + sts_client = boto3.client( + "sts", + region_name=aws_region_name, + endpoint_url=f"https://sts.{aws_region_name}.amazonaws.com" + ) # 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 From 844f15f8b056f2c331b1daefdbdaacfe93a5e95b Mon Sep 17 00:00:00 2001 From: David Manouchehri Date: Tue, 11 Jun 2024 16:59:38 +0000 Subject: [PATCH 9/9] fix(bedrock_httpx.py): Add IAM caching for converse flow. --- litellm/llms/bedrock_httpx.py | 71 ++++++++++++++++++++--------------- 1 file changed, 41 insertions(+), 30 deletions(-) diff --git a/litellm/llms/bedrock_httpx.py b/litellm/llms/bedrock_httpx.py index ee0ae57b42..84b61d4cbd 100644 --- a/litellm/llms/bedrock_httpx.py +++ b/litellm/llms/bedrock_httpx.py @@ -1433,42 +1433,53 @@ class BedrockConverseLLM(BaseLLM): ) = params_to_check ### CHECK STS ### - if ( - aws_web_identity_token is not None - and aws_role_name is not None - and aws_session_name is not None - ): - oidc_token = get_secret(aws_web_identity_token) + 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, + }) - if oidc_token is None: - raise BedrockError( - message="OIDC token could not be retrieved from secret manager.", - status_code=401, + 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, + endpoint_url=f"https://sts.{aws_region_name}.amazonaws.com" ) - sts_client = boto3.client( - "sts", - region_name=aws_region_name, - endpoint_url=f"https://sts.{aws_region_name}.amazonaws.com" - ) + # 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, + ) - # 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_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, + } - 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, - ) + iam_cache.set_cache(key=iam_creds_cache_key, value=json.dumps(iam_creds_dict), ttl=3600 - 60) - return session.get_credentials() + session = boto3.Session(**iam_creds_dict) + + iam_creds = session.get_credentials() + + return iam_creds elif aws_role_name is not None and aws_session_name is not None: sts_client = boto3.client( "sts",