From cdf77556720f524d14dfd540adf4b2e6725549d9 Mon Sep 17 00:00:00 2001 From: Wei Zang Date: Tue, 18 Jun 2024 17:51:28 +0800 Subject: [PATCH 1/2] feat(bedrock): support aws role arn auth --- src/anthropic/lib/bedrock/_auth.py | 11 +++++++++++ src/anthropic/lib/bedrock/_client.py | 8 ++++++++ 2 files changed, 19 insertions(+) diff --git a/src/anthropic/lib/bedrock/_auth.py b/src/anthropic/lib/bedrock/_auth.py index 95c820f6..5cc5b83f 100644 --- a/src/anthropic/lib/bedrock/_auth.py +++ b/src/anthropic/lib/bedrock/_auth.py @@ -64,3 +64,14 @@ def get_auth_headers( prepped = request.prepare() return {key: value for key, value in dict(prepped.headers).items() if value is not None} + +def get_credentials( + *, + aws_role_arn: str | None, +) -> dict: + response = boto3.client('sts').assume_role(RoleArn=aws_role_arn, RoleSessionName='assume-role') + return { + 'aws_access_key': response['Credentials']['AccessKeyId'], + 'aws_secret_key': response['Credentials']['SecretAccessKey'], + 'aws_session_token': response['Credentials']['SessionToken'], + } \ No newline at end of file diff --git a/src/anthropic/lib/bedrock/_client.py b/src/anthropic/lib/bedrock/_client.py index b3f388e5..241b7ca1 100644 --- a/src/anthropic/lib/bedrock/_client.py +++ b/src/anthropic/lib/bedrock/_client.py @@ -95,6 +95,7 @@ def __init__( aws_access_key: str | None = None, aws_region: str | None = None, aws_session_token: str | None = None, + aws_role_arn: str | None = None, base_url: str | httpx.URL | None = None, timeout: float | httpx.Timeout | None | NotGiven = NOT_GIVEN, max_retries: int = DEFAULT_MAX_RETRIES, @@ -122,6 +123,13 @@ def __init__( self.aws_session_token = aws_session_token + if aws_role_arn is not None: + from ._auth import get_credentials + cred = get_credentials(aws_role_arn=aws_role_arn) + self.aws_access_key = cred["aws_access_key"] + self.aws_secret_key = cred["aws_secret_key"] + self.aws_session_token = cred["aws_session_token"] + if base_url is None: base_url = os.environ.get("ANTHROPIC_BEDROCK_BASE_URL") if base_url is None: From 1d1a2dad7266188b5c5fd314e5a6a1c3519f0ad0 Mon Sep 17 00:00:00 2001 From: Wei Zang Date: Tue, 9 Jul 2024 21:02:01 +0800 Subject: [PATCH 2/2] fix(bedrock): fix lint error --- src/anthropic/lib/bedrock/_auth.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/anthropic/lib/bedrock/_auth.py b/src/anthropic/lib/bedrock/_auth.py index 5cc5b83f..dcada500 100644 --- a/src/anthropic/lib/bedrock/_auth.py +++ b/src/anthropic/lib/bedrock/_auth.py @@ -69,6 +69,8 @@ def get_credentials( *, aws_role_arn: str | None, ) -> dict: + import boto3 + response = boto3.client('sts').assume_role(RoleArn=aws_role_arn, RoleSessionName='assume-role') return { 'aws_access_key': response['Credentials']['AccessKeyId'],