From 8004874bc117d526d343fc0a514ef23a31019eae Mon Sep 17 00:00:00 2001 From: Jim Blomo Date: Wed, 10 Jun 2026 12:42:27 -0700 Subject: [PATCH] [bedrock_sigv4_auth] Add AWS-native Bedrock authentication --- README.md | 17 +- pyproject.toml | 5 + requirements-dev.lock | 7 + requirements.lock | 8 + src/openai/lib/bedrock.py | 478 +++++++++++++++++++++++++++++++++++--- tests/lib/test_bedrock.py | 212 ++++++++++++++++- 6 files changed, 692 insertions(+), 35 deletions(-) diff --git a/README.md b/README.md index 6f87246470..c89273612e 100644 --- a/README.md +++ b/README.md @@ -944,7 +944,7 @@ response = client.responses.create( print(response.output_text) ``` -`BedrockOpenAI` configures AWS bearer auth and the Bedrock Mantle endpoint, then uses the normal SDK resources. AWS controls which endpoints and features are supported; unsupported calls surface the provider's normal HTTP errors through the SDK. +`BedrockOpenAI` configures AWS authentication and the Bedrock Mantle endpoint, then uses the normal SDK resources. AWS controls which endpoints and features are supported; unsupported calls surface the provider's normal HTTP errors through the SDK. Pass `base_url` or set `AWS_BEDROCK_BASE_URL` to override the derived `https://bedrock-mantle..api.aws/openai/v1` endpoint. The legacy module client supports `openai.api_type = "amazon-bedrock"` or `OPENAI_API_TYPE=amazon-bedrock`. @@ -957,6 +957,21 @@ client = BedrockOpenAI( ) ``` +To use the standard AWS credential chain and SigV4 authentication, install the Bedrock extra and omit bearer-token configuration: + +```sh +pip install 'openai[bedrock]' +``` + +```py +client = BedrockOpenAI( + aws_region="us-west-2", + aws_profile="my-profile", # optional; otherwise uses the default AWS credential chain +) +``` + +You can also pass explicit temporary credentials or an `aws_credentials_provider` that returns botocore-compatible credentials. Explicit bearer and AWS credential options are mutually exclusive. Without explicit authentication, `AWS_BEARER_TOKEN_BEDROCK` takes precedence over the default AWS credential chain. + ## Versioning This package generally follows [SemVer](https://semver.org/spec/v2.0.0.html) conventions, though certain backwards-incompatible changes may be released as minor versions: diff --git a/pyproject.toml b/pyproject.toml index 75d0d5e246..d84ebfb1b4 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -47,6 +47,10 @@ aiohttp = ["aiohttp", "httpx_aiohttp>=0.1.9"] realtime = ["websockets >= 13, < 16"] datalib = ["numpy >= 1", "pandas >= 1.2.3", "pandas-stubs >= 1.1.0.11"] voice_helpers = ["sounddevice>=0.5.1", "numpy>=2.0.2"] +bedrock = [ + "botocore[crt]>=1.42.0,<1.43; python_version < '3.10'", + "botocore[crt]>=1.42.0,<2; python_version >= '3.10'", +] [tool.rye] managed = true @@ -65,6 +69,7 @@ dev-dependencies = [ "rich>=13.7.1", "inline-snapshot>=0.28.0", "azure-identity >=1.14.1", + "botocore==1.42.97", "types-tqdm > 4", "types-pyaudio > 0", "trio >=0.22.2", diff --git a/requirements-dev.lock b/requirements-dev.lock index 73312bcee1..7d6764ed7a 100644 --- a/requirements-dev.lock +++ b/requirements-dev.lock @@ -33,11 +33,15 @@ attrs==25.4.0 # via nox # via outcome # via trio +awscrt==0.31.2 + # via botocore azure-core==1.36.0 # via azure-identity azure-identity==1.25.1 backports-asyncio-runner==1.2.0 # via pytest-asyncio +botocore==1.42.97 + # via openai certifi==2026.1.4 # via httpcore # via httpx @@ -100,6 +104,8 @@ iniconfig==2.1.0 inline-snapshot==0.31.1 jiter==0.12.0 # via openai +jmespath==1.1.0 + # via botocore markdown-it-py==3.0.0 # via rich mdurl==0.1.2 @@ -218,6 +224,7 @@ typing-inspection==0.4.2 tzdata==2025.2 # via pandas urllib3==2.5.0 + # via botocore # via requests # via types-requests virtualenv==20.35.4 diff --git a/requirements.lock b/requirements.lock index af6e4f99e8..5f5158bd4f 100644 --- a/requirements.lock +++ b/requirements.lock @@ -26,6 +26,10 @@ async-timeout==5.0.1 # via aiohttp attrs==25.4.0 # via aiohttp +awscrt==0.31.2 + # via botocore +botocore==1.42.97 + # via openai certifi==2026.1.4 # via httpcore # via httpx @@ -53,6 +57,8 @@ idna==3.11 # via yarl jiter==0.12.0 # via openai +jmespath==1.1.0 + # via botocore multidict==6.7.0 # via aiohttp # via yarl @@ -100,6 +106,8 @@ typing-inspection==0.4.2 # via pydantic tzdata==2025.2 # via pandas +urllib3==1.26.20 + # via botocore websockets==15.0.1 # via openai yarl==1.22.0 diff --git a/src/openai/lib/bedrock.py b/src/openai/lib/bedrock.py index 266a2e9358..a7dfae7be7 100644 --- a/src/openai/lib/bedrock.py +++ b/src/openai/lib/bedrock.py @@ -3,13 +3,14 @@ import os import re import inspect +import importlib from typing import Any, Mapping, Callable, Awaitable, cast from typing_extensions import Self, override import httpx from ..auth import WorkloadIdentity -from .._types import NOT_GIVEN, Timeout, NotGiven +from .._types import NOT_GIVEN, Headers, Timeout, NotGiven from .._utils import is_given from .._client import OpenAI, AsyncOpenAI from .._models import SecurityOptions, FinalRequestOptions @@ -18,6 +19,181 @@ BedrockTokenProvider = Callable[[], str] AsyncBedrockTokenProvider = Callable[[], "str | Awaitable[str]"] +AwsCredentialsProvider = Callable[[], object] + + +class _BedrockAwsBearerAuth: + def __init__(self) -> None: + try: + auth_module = importlib.import_module("botocore.auth") + awsrequest_module = importlib.import_module("botocore.awsrequest") + tokens_module = importlib.import_module("botocore.tokens") + except ImportError: + self._bearer_auth_cls = None + self._aws_request_cls = None + self._frozen_auth_token_cls = None + return + + self._bearer_auth_cls = auth_module.BearerAuth + self._aws_request_cls = awsrequest_module.AWSRequest + self._frozen_auth_token_cls = tokens_module.FrozenAuthToken + + def sign(self, request: httpx.Request, token: str) -> None: + if self._bearer_auth_cls is None or self._aws_request_cls is None or self._frozen_auth_token_cls is None: + return + + headers = dict(request.headers) + headers.pop("authorization", None) + aws_request = self._aws_request_cls( + method=request.method, + url=str(request.url), + data=request.read(), + headers=headers, + ) + self._bearer_auth_cls(self._frozen_auth_token_cls(token)).add_auth(aws_request) + request.headers.clear() + request.headers.update(dict(aws_request.headers.items())) + + +class _BedrockAwsAuth: + def __init__( + self, + *, + region: str, + profile: str | None, + access_key_id: str | None, + secret_access_key: str | None, + session_token: str | None, + credentials_provider: AwsCredentialsProvider | None, + ) -> None: + try: + auth_module = importlib.import_module("botocore.auth") + session_module = importlib.import_module("botocore.session") + awsrequest_module = importlib.import_module("botocore.awsrequest") + credentials_module = importlib.import_module("botocore.credentials") + except ImportError as exc: + raise OpenAIError( + "AWS credential authentication requires botocore. Install it with `pip install openai[bedrock]`." + ) from exc + + session = session_module.Session(profile=profile) + service_model = session.get_service_model("bedrock-runtime") + auth_options = cast("list[str]", service_model.metadata.get("auth", [])) + if auth_module.resolve_auth_scheme_preference(["sigv4"], auth_options) != "v4": + raise OpenAIError("The installed botocore version does not support Bedrock SigV4 authentication.") + + self._region = region + self._session = session + self._credentials_provider = credentials_provider + self._explicit_credentials = ( + credentials_module.Credentials(access_key_id, secret_access_key, session_token) + if access_key_id is not None and secret_access_key is not None + else None + ) + self._aws_request_cls = awsrequest_module.AWSRequest + self._sigv4_auth_cls = auth_module.SigV4Auth + + def sign(self, request: httpx.Request) -> None: + credentials = ( + self._credentials_provider() + if self._credentials_provider is not None + else self._explicit_credentials or self._session.get_credentials() + ) + if credentials is None: + raise OpenAIError( + "Could not resolve AWS credentials. Configure the standard AWS credential chain or pass explicit " + "AWS credentials to the Bedrock client." + ) + + get_frozen_credentials = getattr(credentials, "get_frozen_credentials", None) + if callable(get_frozen_credentials): + credentials = get_frozen_credentials() + + headers = dict(request.headers) + headers.pop("authorization", None) + aws_request = self._aws_request_cls( + method=request.method, + url=str(request.url), + data=request.read(), + headers=headers, + ) + self._sigv4_auth_cls(credentials, "bedrock-mantle", self._region).add_auth(aws_request) + request.headers.clear() + request.headers.update(dict(aws_request.headers.items())) + + +def _has_explicit_aws_auth( + *, + aws_profile: str | None, + aws_access_key_id: str | None, + aws_secret_access_key: str | None, + aws_session_token: str | None, + aws_credentials_provider: AwsCredentialsProvider | None, +) -> bool: + return any( + value is not None + for value in ( + aws_profile, + aws_access_key_id, + aws_secret_access_key, + aws_session_token, + aws_credentials_provider, + ) + ) + + +def _validate_explicit_aws_auth( + *, + aws_profile: str | None, + aws_access_key_id: str | None, + aws_secret_access_key: str | None, + aws_session_token: str | None, + aws_credentials_provider: AwsCredentialsProvider | None, +) -> None: + if (aws_access_key_id is None) != (aws_secret_access_key is None): + raise OpenAIError("The `aws_access_key_id` and `aws_secret_access_key` arguments must be provided together.") + + credential_sources = sum( + ( + aws_profile is not None, + aws_access_key_id is not None, + aws_credentials_provider is not None, + ) + ) + if credential_sources > 1: + raise OpenAIError( + "The `aws_profile`, explicit AWS credentials, and `aws_credentials_provider` arguments are mutually exclusive." + ) + + if aws_session_token is not None and aws_access_key_id is None: + raise OpenAIError("The `aws_session_token` argument requires explicit AWS access key credentials.") + + +def _resolve_aws_region(aws_region: str | None) -> str: + region = aws_region or os.environ.get("AWS_REGION") or os.environ.get("AWS_DEFAULT_REGION") + if region is None or not region.strip(): + raise OpenAIError("AWS credential authentication requires `aws_region`, `AWS_REGION`, or `AWS_DEFAULT_REGION`.") + return region.strip() + + +def _resolve_bedrock_env_token() -> str | None: + if "AWS_BEARER_TOKEN_BEDROCK" not in os.environ: + return None + + try: + session_module = importlib.import_module("botocore.session") + except ImportError: + return os.environ.get("AWS_BEARER_TOKEN_BEDROCK") or None + + auth_token = session_module.Session().get_auth_token(signing_name="bedrock") + if auth_token is None: + return None + + get_frozen_token = getattr(auth_token, "get_frozen_token", None) + if callable(get_frozen_token): + auth_token = get_frozen_token() + token = cast(str, auth_token.token) + return token or None def _normalize_bedrock_base_url(base_url: str | httpx.URL) -> httpx.URL: @@ -98,7 +274,14 @@ class BedrockOpenAI(OpenAI): """API client for Amazon Bedrock's OpenAI-compatible endpoint.""" _bedrock_token_provider: BedrockTokenProvider | None + _bedrock_aws_bearer_auth: _BedrockAwsBearerAuth | None + _bedrock_aws_auth: _BedrockAwsAuth | None _uses_region_derived_base_url: bool + _aws_profile: str | None + _aws_access_key_id: str | None + _aws_secret_access_key: str | None + _aws_session_token: str | None + _aws_credentials_provider: AwsCredentialsProvider | None aws_region: str | None def __init__( @@ -107,6 +290,11 @@ def __init__( api_key: str | None = None, bedrock_token_provider: BedrockTokenProvider | None = None, aws_region: str | None = None, + aws_profile: str | None = None, + aws_access_key_id: str | None = None, + aws_secret_access_key: str | None = None, + aws_session_token: str | None = None, + aws_credentials_provider: AwsCredentialsProvider | None = None, organization: str | None = None, project: str | None = None, webhook_secret: str | None = None, @@ -123,30 +311,68 @@ def __init__( """Construct a new synchronous Amazon Bedrock client instance. This automatically infers the following arguments from their corresponding environment variables if they are not provided: - - `api_key` from `AWS_BEARER_TOKEN_BEDROCK` + - bearer authentication from `AWS_BEARER_TOKEN_BEDROCK` - `aws_region` from `AWS_REGION` or `AWS_DEFAULT_REGION` when `base_url` and `AWS_BEDROCK_BASE_URL` are not set - `base_url` from `AWS_BEDROCK_BASE_URL` - `bedrock_token_provider` is invoked before each request when provided. + `bedrock_token_provider` is invoked before each request when provided. When no bearer token is configured, + the client uses the standard AWS credential chain and SigV4 authentication. """ - if api_key is None and bedrock_token_provider is None: - api_key = os.environ.get("AWS_BEARER_TOKEN_BEDROCK") - if callable(cast(object, api_key)): raise OpenAIError("Pass refreshable Bedrock credentials via `bedrock_token_provider`, not `api_key`.") + if api_key == "": + raise OpenAIError("The `api_key` argument must not be empty.") + if api_key is not None and bedrock_token_provider is not None: raise OpenAIError("The `api_key` and `bedrock_token_provider` arguments are mutually exclusive.") - if _enforce_credentials and not api_key and bedrock_token_provider is None: - raise OpenAIError( - "Missing credentials. Please pass an `api_key` or `bedrock_token_provider`, or set the " - "`AWS_BEARER_TOKEN_BEDROCK` environment variable." - ) + explicit_bearer_auth = api_key is not None or bedrock_token_provider is not None + explicit_aws_auth = _has_explicit_aws_auth( + aws_profile=aws_profile, + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + aws_session_token=aws_session_token, + aws_credentials_provider=aws_credentials_provider, + ) + if explicit_bearer_auth and explicit_aws_auth: + raise OpenAIError("Bearer token and AWS credential authentication arguments are mutually exclusive.") + + _validate_explicit_aws_auth( + aws_profile=aws_profile, + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + aws_session_token=aws_session_token, + aws_credentials_provider=aws_credentials_provider, + ) + + if not explicit_bearer_auth and not explicit_aws_auth: + api_key = _resolve_bedrock_env_token() + + use_aws_auth = api_key is None and bedrock_token_provider is None + resolved_region = _resolve_aws_region(aws_region) if use_aws_auth else aws_region self._bedrock_token_provider = bedrock_token_provider + self._bedrock_aws_bearer_auth = _BedrockAwsBearerAuth() if not use_aws_auth else None + self._bedrock_aws_auth = ( + _BedrockAwsAuth( + region=cast(str, resolved_region), + profile=aws_profile, + access_key_id=aws_access_key_id, + secret_access_key=aws_secret_access_key, + session_token=aws_session_token, + credentials_provider=aws_credentials_provider, + ) + if use_aws_auth and _enforce_credentials + else None + ) self._uses_region_derived_base_url = _uses_region_derived_bedrock_base_url(base_url) - self.aws_region = aws_region + self._aws_profile = aws_profile + self._aws_access_key_id = aws_access_key_id + self._aws_secret_access_key = aws_secret_access_key + self._aws_session_token = aws_session_token + self._aws_credentials_provider = aws_credentials_provider + self.aws_region = resolved_region super().__init__( api_key=_bedrock_token_provider(bedrock_token_provider) @@ -156,7 +382,7 @@ def __init__( organization=organization, project=project, webhook_secret=webhook_secret, - base_url=_resolve_bedrock_base_url(base_url, aws_region), + base_url=_resolve_bedrock_base_url(base_url, resolved_region), websocket_base_url=websocket_base_url, timeout=timeout, max_retries=max_retries, @@ -169,11 +395,21 @@ def __init__( @override def _auth_headers(self, security: SecurityOptions) -> dict[str, str]: + if self._bedrock_aws_auth is not None: + return {} + if security.get("bearer_auth", False) or security.get("admin_api_key_auth", False): return self._bearer_auth return {} + @override + def _validate_headers(self, headers: Headers, custom_headers: Headers) -> None: + if self._bedrock_aws_auth is not None: + return + + super()._validate_headers(headers, custom_headers) + @override def _prepare_options(self, options: FinalRequestOptions) -> FinalRequestOptions: if ( @@ -185,6 +421,13 @@ def _prepare_options(self, options: FinalRequestOptions) -> FinalRequestOptions: return super()._prepare_options(options) + @override + def _prepare_request(self, request: httpx.Request) -> None: + if self._bedrock_aws_auth is not None: + self._bedrock_aws_auth.sign(request) + elif self._bedrock_aws_bearer_auth is not None: + self._bedrock_aws_bearer_auth.sign(request, self.api_key) + @override def copy( self, @@ -194,6 +437,11 @@ def copy( workload_identity: WorkloadIdentity | None = None, bedrock_token_provider: BedrockTokenProvider | None = None, aws_region: str | None = None, + aws_profile: str | None = None, + aws_access_key_id: str | None = None, + aws_secret_access_key: str | None = None, + aws_session_token: str | None = None, + aws_credentials_provider: AwsCredentialsProvider | None = None, organization: str | None = None, project: str | None = None, webhook_secret: str | None = None, @@ -219,7 +467,7 @@ def copy( raise OpenAIError("Pass refreshable Bedrock credentials via `bedrock_token_provider`, not `api_key`.") if admin_api_key is not None or workload_identity is not None: - raise OpenAIError("BedrockOpenAI only supports Bedrock bearer token authentication.") + raise OpenAIError("BedrockOpenAI only supports Bedrock bearer token or AWS credential authentication.") if api_key is not None and bedrock_token_provider is not None: raise OpenAIError("The `api_key` and `bedrock_token_provider` arguments are mutually exclusive.") @@ -236,14 +484,33 @@ def copy( elif set_default_query is not None: params = set_default_query - if api_key is not None: + aws_auth_override = _has_explicit_aws_auth( + aws_profile=aws_profile, + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + aws_session_token=aws_session_token, + aws_credentials_provider=aws_credentials_provider, + ) + if api_key is not None or aws_auth_override: next_token_provider = None elif bedrock_token_provider is not None: next_token_provider = bedrock_token_provider else: next_token_provider = self._bedrock_token_provider - next_api_key = api_key if api_key is not None else (None if next_token_provider is not None else self.api_key) + preserve_aws_auth = ( + self._bedrock_aws_auth is not None + and not aws_auth_override + and api_key is None + and next_token_provider is None + ) + next_api_key = ( + api_key + if api_key is not None + else None + if next_token_provider is not None or preserve_aws_auth or aws_auth_override + else self.api_key + ) next_base_url = base_url if next_base_url is None and not (aws_region is not None and self._uses_region_derived_base_url): next_base_url = self.base_url @@ -252,6 +519,35 @@ def copy( api_key=next_api_key, bedrock_token_provider=next_token_provider, aws_region=aws_region if aws_region is not None else self.aws_region, + aws_profile=aws_profile if aws_profile is not None else self._aws_profile if preserve_aws_auth else None, + aws_access_key_id=( + aws_access_key_id + if aws_access_key_id is not None + else self._aws_access_key_id + if preserve_aws_auth + else None + ), + aws_secret_access_key=( + aws_secret_access_key + if aws_secret_access_key is not None + else self._aws_secret_access_key + if preserve_aws_auth + else None + ), + aws_session_token=( + aws_session_token + if aws_session_token is not None + else self._aws_session_token + if preserve_aws_auth + else None + ), + aws_credentials_provider=( + aws_credentials_provider + if aws_credentials_provider is not None + else self._aws_credentials_provider + if preserve_aws_auth + else None + ), organization=organization if organization is not None else self.organization, project=project if project is not None else self.project, webhook_secret=webhook_secret if webhook_secret is not None else self.webhook_secret, @@ -273,7 +569,14 @@ class AsyncBedrockOpenAI(AsyncOpenAI): """Async API client for Amazon Bedrock's OpenAI-compatible endpoint.""" _bedrock_token_provider: AsyncBedrockTokenProvider | None + _bedrock_aws_bearer_auth: _BedrockAwsBearerAuth | None + _bedrock_aws_auth: _BedrockAwsAuth | None _uses_region_derived_base_url: bool + _aws_profile: str | None + _aws_access_key_id: str | None + _aws_secret_access_key: str | None + _aws_session_token: str | None + _aws_credentials_provider: AwsCredentialsProvider | None aws_region: str | None def __init__( @@ -282,6 +585,11 @@ def __init__( api_key: str | None = None, bedrock_token_provider: AsyncBedrockTokenProvider | None = None, aws_region: str | None = None, + aws_profile: str | None = None, + aws_access_key_id: str | None = None, + aws_secret_access_key: str | None = None, + aws_session_token: str | None = None, + aws_credentials_provider: AwsCredentialsProvider | None = None, organization: str | None = None, project: str | None = None, webhook_secret: str | None = None, @@ -298,30 +606,68 @@ def __init__( """Construct a new asynchronous Amazon Bedrock client instance. This automatically infers the following arguments from their corresponding environment variables if they are not provided: - - `api_key` from `AWS_BEARER_TOKEN_BEDROCK` + - bearer authentication from `AWS_BEARER_TOKEN_BEDROCK` - `aws_region` from `AWS_REGION` or `AWS_DEFAULT_REGION` when `base_url` and `AWS_BEDROCK_BASE_URL` are not set - `base_url` from `AWS_BEDROCK_BASE_URL` - `bedrock_token_provider` is invoked before each request when provided. + `bedrock_token_provider` is invoked before each request when provided. When no bearer token is configured, + the client uses the standard AWS credential chain and SigV4 authentication. """ - if api_key is None and bedrock_token_provider is None: - api_key = os.environ.get("AWS_BEARER_TOKEN_BEDROCK") - if callable(cast(object, api_key)): raise OpenAIError("Pass refreshable Bedrock credentials via `bedrock_token_provider`, not `api_key`.") + if api_key == "": + raise OpenAIError("The `api_key` argument must not be empty.") + if api_key is not None and bedrock_token_provider is not None: raise OpenAIError("The `api_key` and `bedrock_token_provider` arguments are mutually exclusive.") - if _enforce_credentials and not api_key and bedrock_token_provider is None: - raise OpenAIError( - "Missing credentials. Please pass an `api_key` or `bedrock_token_provider`, or set the " - "`AWS_BEARER_TOKEN_BEDROCK` environment variable." - ) + explicit_bearer_auth = api_key is not None or bedrock_token_provider is not None + explicit_aws_auth = _has_explicit_aws_auth( + aws_profile=aws_profile, + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + aws_session_token=aws_session_token, + aws_credentials_provider=aws_credentials_provider, + ) + if explicit_bearer_auth and explicit_aws_auth: + raise OpenAIError("Bearer token and AWS credential authentication arguments are mutually exclusive.") + + _validate_explicit_aws_auth( + aws_profile=aws_profile, + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + aws_session_token=aws_session_token, + aws_credentials_provider=aws_credentials_provider, + ) + + if not explicit_bearer_auth and not explicit_aws_auth: + api_key = _resolve_bedrock_env_token() + + use_aws_auth = api_key is None and bedrock_token_provider is None + resolved_region = _resolve_aws_region(aws_region) if use_aws_auth else aws_region self._bedrock_token_provider = bedrock_token_provider + self._bedrock_aws_bearer_auth = _BedrockAwsBearerAuth() if not use_aws_auth else None + self._bedrock_aws_auth = ( + _BedrockAwsAuth( + region=cast(str, resolved_region), + profile=aws_profile, + access_key_id=aws_access_key_id, + secret_access_key=aws_secret_access_key, + session_token=aws_session_token, + credentials_provider=aws_credentials_provider, + ) + if use_aws_auth and _enforce_credentials + else None + ) self._uses_region_derived_base_url = _uses_region_derived_bedrock_base_url(base_url) - self.aws_region = aws_region + self._aws_profile = aws_profile + self._aws_access_key_id = aws_access_key_id + self._aws_secret_access_key = aws_secret_access_key + self._aws_session_token = aws_session_token + self._aws_credentials_provider = aws_credentials_provider + self.aws_region = resolved_region super().__init__( api_key=( @@ -333,7 +679,7 @@ def __init__( organization=organization, project=project, webhook_secret=webhook_secret, - base_url=_resolve_bedrock_base_url(base_url, aws_region), + base_url=_resolve_bedrock_base_url(base_url, resolved_region), websocket_base_url=websocket_base_url, timeout=timeout, max_retries=max_retries, @@ -346,11 +692,21 @@ def __init__( @override def _auth_headers(self, security: SecurityOptions) -> dict[str, str]: + if self._bedrock_aws_auth is not None: + return {} + if security.get("bearer_auth", False) or security.get("admin_api_key_auth", False): return self._bearer_auth return {} + @override + def _validate_headers(self, headers: Headers, custom_headers: Headers) -> None: + if self._bedrock_aws_auth is not None: + return + + super()._validate_headers(headers, custom_headers) + @override async def _prepare_options(self, options: FinalRequestOptions) -> FinalRequestOptions: if ( @@ -362,6 +718,13 @@ async def _prepare_options(self, options: FinalRequestOptions) -> FinalRequestOp return await super()._prepare_options(options) + @override + async def _prepare_request(self, request: httpx.Request) -> None: + if self._bedrock_aws_auth is not None: + self._bedrock_aws_auth.sign(request) + elif self._bedrock_aws_bearer_auth is not None: + self._bedrock_aws_bearer_auth.sign(request, self.api_key) + @override def copy( self, @@ -371,6 +734,11 @@ def copy( workload_identity: WorkloadIdentity | None = None, bedrock_token_provider: AsyncBedrockTokenProvider | None = None, aws_region: str | None = None, + aws_profile: str | None = None, + aws_access_key_id: str | None = None, + aws_secret_access_key: str | None = None, + aws_session_token: str | None = None, + aws_credentials_provider: AwsCredentialsProvider | None = None, organization: str | None = None, project: str | None = None, webhook_secret: str | None = None, @@ -396,7 +764,7 @@ def copy( raise OpenAIError("Pass refreshable Bedrock credentials via `bedrock_token_provider`, not `api_key`.") if admin_api_key is not None or workload_identity is not None: - raise OpenAIError("AsyncBedrockOpenAI only supports Bedrock bearer token authentication.") + raise OpenAIError("AsyncBedrockOpenAI only supports Bedrock bearer token or AWS credential authentication.") if api_key is not None and bedrock_token_provider is not None: raise OpenAIError("The `api_key` and `bedrock_token_provider` arguments are mutually exclusive.") @@ -413,14 +781,33 @@ def copy( elif set_default_query is not None: params = set_default_query - if api_key is not None: + aws_auth_override = _has_explicit_aws_auth( + aws_profile=aws_profile, + aws_access_key_id=aws_access_key_id, + aws_secret_access_key=aws_secret_access_key, + aws_session_token=aws_session_token, + aws_credentials_provider=aws_credentials_provider, + ) + if api_key is not None or aws_auth_override: next_token_provider = None elif bedrock_token_provider is not None: next_token_provider = bedrock_token_provider else: next_token_provider = self._bedrock_token_provider - next_api_key = api_key if api_key is not None else (None if next_token_provider is not None else self.api_key) + preserve_aws_auth = ( + self._bedrock_aws_auth is not None + and not aws_auth_override + and api_key is None + and next_token_provider is None + ) + next_api_key = ( + api_key + if api_key is not None + else None + if next_token_provider is not None or preserve_aws_auth or aws_auth_override + else self.api_key + ) next_base_url = base_url if next_base_url is None and not (aws_region is not None and self._uses_region_derived_base_url): next_base_url = self.base_url @@ -429,6 +816,35 @@ def copy( api_key=next_api_key, bedrock_token_provider=next_token_provider, aws_region=aws_region if aws_region is not None else self.aws_region, + aws_profile=aws_profile if aws_profile is not None else self._aws_profile if preserve_aws_auth else None, + aws_access_key_id=( + aws_access_key_id + if aws_access_key_id is not None + else self._aws_access_key_id + if preserve_aws_auth + else None + ), + aws_secret_access_key=( + aws_secret_access_key + if aws_secret_access_key is not None + else self._aws_secret_access_key + if preserve_aws_auth + else None + ), + aws_session_token=( + aws_session_token + if aws_session_token is not None + else self._aws_session_token + if preserve_aws_auth + else None + ), + aws_credentials_provider=( + aws_credentials_provider + if aws_credentials_provider is not None + else self._aws_credentials_provider + if preserve_aws_auth + else None + ), organization=organization if organization is not None else self.organization, project=project if project is not None else self.project, webhook_secret=webhook_secret if webhook_secret is not None else self.webhook_secret, diff --git a/tests/lib/test_bedrock.py b/tests/lib/test_bedrock.py index dab9abd1cf..656dcb55a0 100644 --- a/tests/lib/test_bedrock.py +++ b/tests/lib/test_bedrock.py @@ -2,12 +2,14 @@ import json from typing import Any, Union, Protocol, cast +from pathlib import Path import httpx import pytest from httpx import URL from respx import MockRouter +import openai.lib.bedrock as bedrock_module from openai import OpenAIError, NotFoundError from tests.utils import update_env from openai._types import Omit @@ -76,6 +78,13 @@ class MockRequestCall(Protocol): request: httpx.Request +class MockAwsCredentials: + def __init__(self, access_key: str, secret_key: str, token: str | None = None) -> None: + self.access_key = access_key + self.secret_key = secret_key + self.token = token + + def make_sync_client(**kwargs: Any) -> BedrockOpenAI: return BedrockOpenAI(http_client=httpx.Client(trust_env=False), **kwargs) @@ -123,6 +132,79 @@ def test_bedrock_config_precedence(client_cls: type[Client]) -> None: assert client.api_key == "explicit token" +@pytest.mark.respx() +def test_env_bearer_does_not_require_botocore(monkeypatch: pytest.MonkeyPatch, respx_mock: MockRouter) -> None: + real_import_module = bedrock_module.importlib.import_module + + def import_module(name: str) -> Any: + if name.startswith("botocore"): + raise ImportError(name) + return real_import_module(name) + + monkeypatch.setattr(bedrock_module.importlib, "import_module", import_module) + respx_mock.post("https://example.com/openai/v1/responses").mock( + return_value=httpx.Response(200, json=RESPONSE_BODY) + ) + with update_env( + AWS_BEDROCK_BASE_URL="https://example.com/openai/v1", + AWS_BEARER_TOKEN_BEDROCK="env token", + ): + client = make_sync_client() + + client.responses.create(model="gpt-4o", input="hello") + + request = cast("list[MockRequestCall]", respx_mock.calls)[0].request + assert request.headers["Authorization"] == "Bearer env token" + + +def test_empty_env_bearer_without_botocore_uses_aws_credentials(monkeypatch: pytest.MonkeyPatch) -> None: + real_import_module = bedrock_module.importlib.import_module + + def import_module(name: str) -> Any: + if name.startswith("botocore"): + raise ImportError(name) + return real_import_module(name) + + monkeypatch.setattr(bedrock_module.importlib, "import_module", import_module) + with update_env(AWS_BEARER_TOKEN_BEDROCK="", AWS_REGION="us-east-1"): + with pytest.raises(OpenAIError, match="requires botocore"): + BedrockOpenAI() + + +@pytest.mark.respx() +def test_env_bearer_uses_botocore_bearer_auth(monkeypatch: pytest.MonkeyPatch, respx_mock: MockRouter) -> None: + auth_module = bedrock_module.importlib.import_module("botocore.auth") + calls = 0 + real_add_auth = auth_module.BearerAuth.add_auth + + def add_auth(auth: object, request: object) -> None: + nonlocal calls + calls += 1 + real_add_auth(auth, request) + + monkeypatch.setattr(auth_module.BearerAuth, "add_auth", add_auth) + respx_mock.post("https://example.com/openai/v1/responses").mock( + return_value=httpx.Response(200, json=RESPONSE_BODY) + ) + with update_env(AWS_BEARER_TOKEN_BEDROCK="env token"): + client = make_sync_client(base_url="https://example.com/openai/v1") + + client.responses.create(model="gpt-4o", input="hello") + + request = cast("list[MockRequestCall]", respx_mock.calls)[0].request + assert request.headers["Authorization"] == "Bearer env token" + assert calls == 1 + + +@pytest.mark.parametrize("client_cls", [BedrockOpenAI, AsyncBedrockOpenAI]) +def test_empty_env_bearer_falls_back_to_aws_credentials(client_cls: type[Client]) -> None: + with update_env(AWS_BEARER_TOKEN_BEDROCK="", AWS_REGION="us-east-1"): + client = make_sync_client() if client_cls is BedrockOpenAI else make_async_client() + + assert client.api_key == "" + assert client._bedrock_aws_auth is not None + + @pytest.mark.parametrize("client_cls", [BedrockOpenAI, AsyncBedrockOpenAI]) def test_bedrock_region_precedence(client_cls: type[Client]) -> None: with update_env(AWS_BEDROCK_BASE_URL=Omit(), AWS_REGION="us-east-1", AWS_DEFAULT_REGION="us-west-2"): @@ -170,8 +252,10 @@ def test_does_not_use_openai_api_key(client_cls: type[Client]) -> None: AWS_BEARER_TOKEN_BEDROCK=Omit(), AWS_BEDROCK_BASE_URL="https://example.com/openai/v1", ): - with pytest.raises(OpenAIError, match="AWS_BEARER_TOKEN_BEDROCK"): - client_cls() + client = client_cls() + + assert client.api_key == "" + assert client._bedrock_aws_auth is not None @pytest.mark.parametrize("client_cls", [BedrockOpenAI, AsyncBedrockOpenAI]) @@ -184,6 +268,33 @@ def test_rejects_static_token_and_provider(client_cls: type[Client]) -> None: ) +@pytest.mark.parametrize("client_cls", [BedrockOpenAI, AsyncBedrockOpenAI]) +def test_rejects_empty_explicit_bearer_token(client_cls: type[Client]) -> None: + with pytest.raises(OpenAIError, match="must not be empty"): + client_cls(base_url="https://example.com/openai/v1", api_key="") + + +@pytest.mark.parametrize("client_cls", [BedrockOpenAI, AsyncBedrockOpenAI]) +def test_rejects_bearer_and_aws_credentials(client_cls: type[Client]) -> None: + with pytest.raises(OpenAIError, match="mutually exclusive"): + client_cls( + base_url="https://example.com/openai/v1", + api_key="token", + aws_access_key_id="access key", + aws_secret_access_key="secret key", + ) + + +@pytest.mark.parametrize("client_cls", [BedrockOpenAI, AsyncBedrockOpenAI]) +def test_rejects_partial_explicit_aws_credentials(client_cls: type[Client]) -> None: + with pytest.raises(OpenAIError, match="must be provided together"): + client_cls( + base_url="https://example.com/openai/v1", + aws_region="us-east-1", + aws_access_key_id="access key", + ) + + @pytest.mark.parametrize("client_cls", [BedrockOpenAI, AsyncBedrockOpenAI]) def test_requires_refreshable_tokens_to_use_provider_option(client_cls: type[Client]) -> None: with pytest.raises(OpenAIError, match="bedrock_token_provider"): @@ -240,6 +351,59 @@ async def test_token_provider_refresh_async(respx_mock: MockRouter) -> None: assert calls[1].request.headers["Authorization"] == "Bearer second" +@pytest.mark.respx() +def test_explicit_aws_credentials_override_ambient_bearer(respx_mock: MockRouter) -> None: + respx_mock.post("https://example.com/openai/v1/responses").mock( + return_value=httpx.Response(200, json=RESPONSE_BODY) + ) + with update_env(AWS_BEARER_TOKEN_BEDROCK="ambient token"): + client = BedrockOpenAI( + base_url="https://example.com/openai/v1", + aws_region="us-east-1", + aws_access_key_id="access key", + aws_secret_access_key="secret key", + aws_session_token="session token", + http_client=httpx.Client(trust_env=False), + ) + + client.responses.create(model="gpt-4o", input="hello") + + request = cast("list[MockRequestCall]", respx_mock.calls)[0].request + assert request.headers["Authorization"].startswith("AWS4-HMAC-SHA256 Credential=access key/") + assert request.headers["X-Amz-Security-Token"] == "session token" + + +@pytest.mark.respx() +def test_aws_credentials_provider_refreshes_before_retries(respx_mock: MockRouter) -> None: + respx_mock.post("https://example.com/openai/v1/responses").mock( + side_effect=[ + httpx.Response(500, json={"error": "server error"}), + httpx.Response(200, json=RESPONSE_BODY), + ] + ) + credentials = iter( + [ + MockAwsCredentials("first access key", "first secret", "first session token"), + MockAwsCredentials("second access key", "second secret", "second session token"), + ] + ) + client = BedrockOpenAI( + base_url="https://example.com/openai/v1", + aws_region="us-east-1", + aws_credentials_provider=lambda: next(credentials), + http_client=httpx.Client(trust_env=False), + max_retries=1, + ) + + client.responses.create(model="gpt-4o", input="hello") + + calls = cast("list[MockRequestCall]", respx_mock.calls) + assert "Credential=first access key/" in calls[0].request.headers["Authorization"] + assert calls[0].request.headers["X-Amz-Security-Token"] == "first session token" + assert "Credential=second access key/" in calls[1].request.headers["Authorization"] + assert calls[1].request.headers["X-Amz-Security-Token"] == "second session token" + + def test_preserves_token_provider_across_with_options() -> None: client = BedrockOpenAI( base_url="https://example.com/openai/v1", @@ -252,6 +416,48 @@ def test_preserves_token_provider_across_with_options() -> None: assert copied_client._refresh_api_key() == "provider token" +def test_preserves_aws_credentials_across_with_options() -> None: + client = BedrockOpenAI( + base_url="https://example.com/openai/v1", + aws_region="us-east-1", + aws_access_key_id="access key", + aws_secret_access_key="secret key", + http_client=httpx.Client(trust_env=False), + ) + + copied_client = client.with_options(timeout=1) + + assert copied_client._bedrock_aws_auth is not None + assert copied_client._aws_access_key_id == "access key" + + +@pytest.mark.parametrize("client_cls", [BedrockOpenAI, AsyncBedrockOpenAI]) +def test_with_options_replaces_the_aws_credential_source(client_cls: type[Client], tmp_path: Path) -> None: + config_path = tmp_path / "config" + config_path.write_text("[profile other-profile]\nregion = us-east-1\n") + explicit_credentials_client = client_cls( + base_url="https://example.com/openai/v1", + aws_region="us-east-1", + aws_access_key_id="access key", + aws_secret_access_key="secret key", + ) + with update_env(AWS_CONFIG_FILE=str(config_path)): + profile_client = explicit_credentials_client.with_options(aws_profile="other-profile") + + assert profile_client._aws_profile == "other-profile" + assert profile_client._aws_access_key_id is None + assert profile_client._aws_secret_access_key is None + + explicit_credentials_client = profile_client.with_options( + aws_access_key_id="replacement access key", + aws_secret_access_key="replacement secret key", + ) + + assert explicit_credentials_client._aws_profile is None + assert explicit_credentials_client._aws_access_key_id == "replacement access key" + assert explicit_credentials_client._aws_secret_access_key == "replacement secret key" + + @pytest.mark.parametrize("client_cls", [BedrockOpenAI, AsyncBedrockOpenAI]) def test_with_options_api_key_replaces_token_provider(client_cls: type[Client]) -> None: client = ( @@ -311,7 +517,7 @@ def test_with_options_aws_region_keeps_explicit_base_url(client_cls: type[Client def test_rejects_non_bedrock_copy_auth(copy_kwargs: dict[str, Any]) -> None: client = make_sync_client(base_url="https://example.com/openai/v1", api_key="token") - with pytest.raises(OpenAIError, match="only supports Bedrock bearer token authentication"): + with pytest.raises(OpenAIError, match="only supports Bedrock bearer token or AWS credential authentication"): client.with_options(**copy_kwargs)