From c301d0c5caf84c35c09524106657da9677fff117 Mon Sep 17 00:00:00 2001 From: David Manouchehri Date: Thu, 16 May 2024 18:05:55 +0000 Subject: [PATCH 1/5] feat (bedrock_httpx.py): Add OIDC support for Amazon Bedrock with httpx. --- litellm/llms/bedrock_httpx.py | 30 +++++++++++++++++++++++++++++- 1 file changed, 29 insertions(+), 1 deletion(-) diff --git a/litellm/llms/bedrock_httpx.py b/litellm/llms/bedrock_httpx.py index 337055dc2ee..a51f57b3e01 100644 --- a/litellm/llms/bedrock_httpx.py +++ b/litellm/llms/bedrock_httpx.py @@ -217,6 +217,7 @@ class BedrockLLM(BaseLLM): aws_session_name: Optional[str] = None, aws_profile_name: Optional[str] = None, aws_role_name: Optional[str] = None, + aws_web_identity_token: Optional[str] = None, ): """ Return a boto3.Credentials object @@ -231,6 +232,7 @@ class BedrockLLM(BaseLLM): aws_session_name, aws_profile_name, aws_role_name, + aws_web_identity_token, ] # Iterate over parameters and update if needed @@ -247,10 +249,34 @@ class BedrockLLM(BaseLLM): aws_session_name, aws_profile_name, aws_role_name, + aws_web_identity_token, ) = params_to_check ### CHECK STS ### - if aws_role_name is not None and aws_session_name is not None: + 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 oidc_token is None: + raise BedrockError( + message="OIDC token could not be retrieved from secret manager.", + status_code=401, + ) + + 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, + ) + + return sts_response["Credentials"] + elif aws_role_name is not None and aws_session_name is not None: sts_client = boto3.client( "sts", aws_access_key_id=aws_access_key_id, # [OPTIONAL] @@ -582,6 +608,7 @@ class BedrockLLM(BaseLLM): aws_bedrock_runtime_endpoint = optional_params.pop( "aws_bedrock_runtime_endpoint", None ) # https://bedrock-runtime.{region_name}.amazonaws.com + aws_web_identity_token = optional_params.pop("aws_web_identity_token", None) ### SET REGION NAME ### if aws_region_name is None: @@ -609,6 +636,7 @@ class BedrockLLM(BaseLLM): aws_session_name=aws_session_name, aws_profile_name=aws_profile_name, aws_role_name=aws_role_name, + aws_web_identity_token=aws_web_identity_token, ) ### SET RUNTIME ENDPOINT ### From 0ed01f430d0cd44f8041176c987508e76661a196 Mon Sep 17 00:00:00 2001 From: David Manouchehri Date: Thu, 16 May 2024 18:57:13 +0000 Subject: [PATCH 2/5] fix(bedrock_httpx.py): Fix OIDC support. --- litellm/llms/bedrock_httpx.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/litellm/llms/bedrock_httpx.py b/litellm/llms/bedrock_httpx.py index a51f57b3e01..de0a08426c4 100644 --- a/litellm/llms/bedrock_httpx.py +++ b/litellm/llms/bedrock_httpx.py @@ -275,7 +275,14 @@ class BedrockLLM(BaseLLM): DurationSeconds=3600, ) - return sts_response["Credentials"] + 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, + ) + + return session.get_credentials() elif aws_role_name is not None and aws_session_name is not None: sts_client = boto3.client( "sts", From 08ee4519b6b70002803739f9e7a511c6838bd960 Mon Sep 17 00:00:00 2001 From: David Manouchehri Date: Thu, 16 May 2024 19:21:40 +0000 Subject: [PATCH 3/5] Add unit test for bedrock httpx oidc auth. --- litellm/tests/test_bedrock_completion.py | 27 ++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/litellm/tests/test_bedrock_completion.py b/litellm/tests/test_bedrock_completion.py index e7346c3f34c..f84cd9b7004 100644 --- a/litellm/tests/test_bedrock_completion.py +++ b/litellm/tests/test_bedrock_completion.py @@ -238,6 +238,33 @@ def test_completion_bedrock_claude_sts_oidc_auth(): except Exception as e: pytest.fail(f"Error occurred: {e}") +def test_completion_bedrock_httpx_command_r_sts_oidc_auth(): + print("\ncalling bedrock httpx command r with oidc auth") + import os + + aws_web_identity_token = "oidc/circleci_v2/" + aws_region_name = os.environ["AWS_REGION_NAME"] + aws_role_name = os.environ["AWS_TEMP_ROLE_NAME"] + + try: + litellm.set_verbose = True + + response = completion( + model="bedrock/cohere.command-r-v1:0", + messages=messages, + max_tokens=10, + temperature=0.1, + 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", + ) + # Add any assertions here to check the response + print(response) + except RateLimitError: + pass + except Exception as e: + pytest.fail(f"Error occurred: {e}") def test_bedrock_claude_3(): try: From eb57b2c6d779274e0a43d18a47184a11099f5de6 Mon Sep 17 00:00:00 2001 From: David Manouchehri Date: Fri, 31 May 2024 13:48:52 +0000 Subject: [PATCH 4/5] Fix OIDC tests. This should always work, even without IAM set up . --- litellm/tests/test_secret_manager.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/litellm/tests/test_secret_manager.py b/litellm/tests/test_secret_manager.py index f990fed3f19..d18d0ea63d5 100644 --- a/litellm/tests/test_secret_manager.py +++ b/litellm/tests/test_secret_manager.py @@ -54,7 +54,10 @@ def test_oidc_github(): print(f"secret_val: {redact_oidc_signature(secret_val)}") -@pytest.mark.skip(reason="Cannot run without being in a CircleCI Runner") +@pytest.mark.skipif( + os.environ.get("CIRCLE_OIDC_TOKEN") is None, + reason="Cannot run without being in CircleCI Runner", +) def test_oidc_circleci(): secret_val = get_secret( "oidc/circleci/https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-text-express-v1/invoke" @@ -63,7 +66,10 @@ def test_oidc_circleci(): print(f"secret_val: {redact_oidc_signature(secret_val)}") -@pytest.mark.skip(reason="Cannot run without being in a CircleCI Runner") +@pytest.mark.skipif( + os.environ.get("CIRCLE_OIDC_TOKEN_V2") is None, + reason="Cannot run without being in CircleCI Runner", +) def test_oidc_circleci_v2(): secret_val = get_secret( "oidc/circleci_v2/https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-text-express-v1/invoke" From d70d484e10e9dcf1a8a4004838d9c1fd6ea52d51 Mon Sep 17 00:00:00 2001 From: David Manouchehri Date: Fri, 31 May 2024 13:50:49 +0000 Subject: [PATCH 5/5] Fix: Use David's AWS account to pass unit tests. --- litellm/tests/test_bedrock_completion.py | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/litellm/tests/test_bedrock_completion.py b/litellm/tests/test_bedrock_completion.py index f84cd9b7004..047f0cb2e2e 100644 --- a/litellm/tests/test_bedrock_completion.py +++ b/litellm/tests/test_bedrock_completion.py @@ -209,20 +209,25 @@ def test_completion_bedrock_claude_sts_client_auth(): # test_completion_bedrock_claude_sts_client_auth() -@pytest.mark.skip(reason="We don't have Circle CI OIDC credentials as yet") +@pytest.mark.skipif( + os.environ.get("CIRCLE_OIDC_TOKEN_V2") is None, + reason="Cannot run without being in CircleCI Runner", +) def test_completion_bedrock_claude_sts_oidc_auth(): print("\ncalling bedrock claude with oidc auth") import os aws_web_identity_token = "oidc/circleci_v2/" aws_region_name = os.environ["AWS_REGION_NAME"] - aws_role_name = os.environ["AWS_TEMP_ROLE_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 + aws_role_name = "arn:aws:iam::335785316107:role/litellm-github-unit-tests-circleci" try: litellm.set_verbose = True response = completion( - model="bedrock/anthropic.claude-instant-v1", + model="bedrock/anthropic.claude-3-haiku-20240307-v1:0", messages=messages, max_tokens=10, temperature=0.1, @@ -238,13 +243,19 @@ def test_completion_bedrock_claude_sts_oidc_auth(): except Exception as e: pytest.fail(f"Error occurred: {e}") +@pytest.mark.skipif( + os.environ.get("CIRCLE_OIDC_TOKEN_V2") is None, + reason="Cannot run without being in CircleCI Runner", +) def test_completion_bedrock_httpx_command_r_sts_oidc_auth(): print("\ncalling bedrock httpx command r with oidc auth") import os aws_web_identity_token = "oidc/circleci_v2/" aws_region_name = os.environ["AWS_REGION_NAME"] - aws_role_name = os.environ["AWS_TEMP_ROLE_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 + aws_role_name = "arn:aws:iam::335785316107:role/litellm-github-unit-tests-circleci" try: litellm.set_verbose = True