Skip to content

feat(aws): add explicit auth_mode option to Bedrock and AWS clients #1902

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
lx3133584 wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat(aws): add explicit auth_mode option to Bedrock and AWS clients #1902

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions src/anthropic/lib/aws/_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
from ..._utils import asyncify
from ..._client import Anthropic, AsyncAnthropic
from ._credentials import (
AuthMode,
resolve_region,
resolve_api_key,
resolve_base_url,
Expand Down Expand Up @@ -36,6 +37,7 @@ def __init__(
self,
*,
api_key: str | None = None,
auth_mode: AuthMode = "auto",
aws_access_key: str | None = None,
aws_secret_key: str | None = None,
aws_region: str | None = None,
Expand All @@ -56,6 +58,7 @@ def __init__(
webhook_key: str | None = None,
) -> None:
self._skip_auth = skip_auth
self.auth_mode = auth_mode

validate_credentials(aws_access_key=aws_access_key, aws_secret_key=aws_secret_key)

Expand All @@ -64,6 +67,7 @@ def __init__(
resolved_api_key = None
else:
self._use_sigv4 = resolve_auth_mode(
auth_mode=auth_mode,
api_key=api_key,
aws_access_key=aws_access_key,
aws_secret_key=aws_secret_key,
Expand Down Expand Up @@ -170,6 +174,7 @@ def copy( # type: ignore[override] # pyright: ignore[reportIncompatibleMethodO
self,
*,
api_key: str | None = None,
auth_mode: AuthMode = "auto",
aws_access_key: str | None = None,
aws_secret_key: str | None = None,
aws_region: str | None = None,
Expand Down Expand Up @@ -244,6 +249,7 @@ def __init__(
self,
*,
api_key: str | None = None,
auth_mode: AuthMode = "auto",
aws_access_key: str | None = None,
aws_secret_key: str | None = None,
aws_region: str | None = None,
Expand All @@ -264,6 +270,7 @@ def __init__(
webhook_key: str | None = None,
) -> None:
self._skip_auth = skip_auth
self.auth_mode = auth_mode

validate_credentials(aws_access_key=aws_access_key, aws_secret_key=aws_secret_key)

Expand All @@ -272,6 +279,7 @@ def __init__(
resolved_api_key = None
else:
self._use_sigv4 = resolve_auth_mode(
auth_mode=auth_mode,
api_key=api_key,
aws_access_key=aws_access_key,
aws_secret_key=aws_secret_key,
Expand Down Expand Up @@ -378,6 +386,7 @@ def copy( # type: ignore[override] # pyright: ignore[reportIncompatibleMethodO
self,
*,
api_key: str | None = None,
auth_mode: AuthMode = "auto",
aws_access_key: str | None = None,
aws_secret_key: str | None = None,
aws_region: str | None = None,
Expand Down
8 changes: 8 additions & 0 deletions src/anthropic/lib/aws/_credentials.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,9 @@

import os
from typing import Sequence
from typing_extensions import Literal

AuthMode = Literal["auto", "api_key", "sigv4"]


def validate_credentials(
Expand Down Expand Up @@ -30,12 +33,17 @@ def _read_env(*env_vars: str) -> str | None:

def resolve_auth_mode(
*,
auth_mode: AuthMode | None = "auto",
api_key: str | None,
aws_access_key: str | None,
aws_secret_key: str | None,
aws_profile: str | None,
api_key_env_vars: Sequence[str] = ("ANTHROPIC_AWS_API_KEY",),
) -> bool:
if auth_mode == "sigv4":
return True
if auth_mode == "api_key":
return False
"""Determine whether to use SigV4 auth. Returns True for SigV4, False for API key.

Auth precedence:
Expand Down
15 changes: 15 additions & 0 deletions src/anthropic/lib/bedrock/_mantle.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
merge_headers,
)
from ..aws._credentials import (
AuthMode,
resolve_region,
resolve_api_key,
resolve_auth_mode,
Expand Down Expand Up @@ -111,6 +112,7 @@ def _resolve_mantle_config(
aws_region: str | None,
aws_profile: str | None,
skip_auth: bool,
auth_mode: AuthMode = "auto",
base_url: str | httpx2.URL | None,
default_headers: Mapping[str, str] | None,
) -> tuple[str | None, str | httpx2.URL, bool, dict[str, str]]:
Expand All @@ -125,6 +127,7 @@ def _resolve_mantle_config(
validate_credentials(aws_access_key=aws_access_key, aws_secret_key=aws_secret_key)

use_sigv4 = resolve_auth_mode(
auth_mode=auth_mode,
api_key=api_key,
aws_access_key=aws_access_key,
aws_secret_key=aws_secret_key,
Expand Down Expand Up @@ -171,6 +174,7 @@ class AnthropicBedrockMantle(BaseMantleClient[httpx2.Client, Stream[Any]], SyncA
aws_session_token: str | None
aws_profile: str | None
skip_auth: bool
auth_mode: AuthMode

_use_sigv4: bool

Expand All @@ -183,6 +187,7 @@ def __init__(
aws_region: str | None = None,
aws_profile: str | None = None,
api_key: str | None = None,
auth_mode: AuthMode = "auto",
skip_auth: bool = False,
base_url: str | httpx2.URL | None = None,
timeout: float | Timeout | None | NotGiven = not_given,
Expand All @@ -195,6 +200,7 @@ def __init__(
) -> None:
resolved_api_key, resolved_base_url, use_sigv4, merged_headers = _resolve_mantle_config(
api_key=api_key,
auth_mode=auth_mode,
aws_access_key=aws_access_key,
aws_secret_key=aws_secret_key,
aws_region=aws_region,
Expand Down Expand Up @@ -225,6 +231,7 @@ def __init__(
self.aws_session_token = aws_session_token
self.aws_profile = aws_profile
self.skip_auth = skip_auth
self.auth_mode = auth_mode
self._use_sigv4 = use_sigv4

self.messages = Messages(self)
Expand Down Expand Up @@ -289,6 +296,7 @@ def copy(
aws_session_token: str | None = None,
aws_region: str | None = None,
aws_profile: str | None = None,
auth_mode: AuthMode | None = None,
skip_auth: bool | None = None,
base_url: str | httpx2.URL | None = None,
timeout: float | Timeout | None | NotGiven = not_given,
Expand Down Expand Up @@ -329,6 +337,7 @@ def copy(
aws_session_token=aws_session_token or self.aws_session_token,
aws_region=aws_region or self.aws_region,
aws_profile=aws_profile or self.aws_profile,
auth_mode=self.auth_mode if auth_mode is None else auth_mode,
skip_auth=skip_auth if skip_auth is not None else self.skip_auth,
base_url=base_url or self.base_url,
timeout=self.timeout if isinstance(timeout, NotGiven) else timeout,
Expand Down Expand Up @@ -367,6 +376,7 @@ class AsyncAnthropicBedrockMantle(BaseMantleClient[httpx2.AsyncClient, AsyncStre
aws_session_token: str | None
aws_profile: str | None
skip_auth: bool
auth_mode: AuthMode

_use_sigv4: bool

Expand All @@ -379,6 +389,7 @@ def __init__(
aws_region: str | None = None,
aws_profile: str | None = None,
api_key: str | None = None,
auth_mode: AuthMode = "auto",
skip_auth: bool = False,
base_url: str | httpx2.URL | None = None,
timeout: float | Timeout | None | NotGiven = not_given,
Expand All @@ -391,6 +402,7 @@ def __init__(
) -> None:
resolved_api_key, resolved_base_url, use_sigv4, merged_headers = _resolve_mantle_config(
api_key=api_key,
auth_mode=auth_mode,
aws_access_key=aws_access_key,
aws_secret_key=aws_secret_key,
aws_region=aws_region,
Expand Down Expand Up @@ -421,6 +433,7 @@ def __init__(
self.aws_session_token = aws_session_token
self.aws_profile = aws_profile
self.skip_auth = skip_auth
self.auth_mode = auth_mode
self._use_sigv4 = use_sigv4

self.messages = AsyncMessages(self)
Expand Down Expand Up @@ -485,6 +498,7 @@ def copy(
aws_session_token: str | None = None,
aws_region: str | None = None,
aws_profile: str | None = None,
auth_mode: AuthMode | None = None,
skip_auth: bool | None = None,
base_url: str | httpx2.URL | None = None,
timeout: float | Timeout | None | NotGiven = not_given,
Expand Down Expand Up @@ -525,6 +539,7 @@ def copy(
aws_session_token=aws_session_token or self.aws_session_token,
aws_region=aws_region or self.aws_region,
aws_profile=aws_profile or self.aws_profile,
auth_mode=self.auth_mode if auth_mode is None else auth_mode,
skip_auth=skip_auth if skip_auth is not None else self.skip_auth,
base_url=base_url or self.base_url,
timeout=self.timeout if isinstance(timeout, NotGiven) else timeout,
Expand Down
24 changes: 24 additions & 0 deletions tests/lib/test_aws_auth.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,30 @@ def test_secret_key_only_raises(self) -> None:


class TestResolveAuthMode:
def test_explicit_auth_mode_sigv4_overrides_env_api_key(self, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("ANTHROPIC_AWS_API_KEY", "env-key")
assert (
resolve_auth_mode(
auth_mode="sigv4",
api_key=None,
aws_access_key=None,
aws_secret_key=None,
aws_profile=None,
)
is True
)

def test_explicit_auth_mode_api_key_overrides_creds(self) -> None:
assert (
resolve_auth_mode(
auth_mode="api_key",
api_key="my-key",
aws_access_key="AKID",
aws_secret_key="secret",
aws_profile=None,
)
is False
)
def test_api_key_arg_returns_false(self) -> None:
assert (
resolve_auth_mode(
Expand Down
31 changes: 31 additions & 0 deletions tests/lib/test_bedrock_mantle.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -163,6 +163,37 @@ def test_beta_skills_not_available(self) -> None:


class TestAuthPrecedence:

def test_explicit_auth_mode_sigv4_forces_sigv4(self, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token")
client = AnthropicBedrockMantle(
auth_mode="sigv4",
aws_region="us-east-1",
)
assert client._use_sigv4 is True
assert client.auth_headers == {}

def test_explicit_auth_mode_api_key_forces_api_key(self) -> None:
client = AnthropicBedrockMantle(
auth_mode="api_key",
api_key="my-bearer-key",
aws_access_key="AKID",
aws_secret_key="secret",
aws_region="us-east-1",
)
assert client._use_sigv4 is False
assert client.api_key == "my-bearer-key"
assert client.auth_headers == {"Authorization": "Bearer my-bearer-key"}

def test_copy_preserves_auth_mode(self, monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", "env-bearer-token")
client = AnthropicBedrockMantle(
auth_mode="sigv4",
aws_region="us-east-1",
)
copied = client.copy()
assert copied.auth_mode == "sigv4"
assert copied._use_sigv4 is True
def test_api_key_arg_uses_api_key_mode(self) -> None:
client = AnthropicBedrockMantle(
api_key="my-key",
Expand Down