feat: Add AWS Bedrock client with SigV4 authentication#3170
feat: Add AWS Bedrock client with SigV4 authentication#3170prashanthkvs wants to merge 2 commits intoopenai:mainfrom
Conversation
381be22 to
f3f8097
Compare
|
@3coins Could you please review this PR? |
3coins
left a comment
There was a problem hiding this comment.
@prashanthkvs
Thanks for submitting this change. Left some suggestions and corrections.
There are some parts here that will need help from the maintainers, I have tagged one of them, but feel free to ping on the PR in a few days if you don't see any activity. Will re-review once you have made updates.
|
|
||
| Run: | ||
| export AWS_REGION=us-west-2 | ||
| PYTHONPATH=src python3 examples/bedrock_mantle.py |
There was a problem hiding this comment.
| PYTHONPATH=src python3 examples/bedrock_mantle.py | |
| PYTHONPATH=src python3 examples/aws_client.py |
|
|
||
| Run: | ||
| export AWS_REGION=us-west-2 | ||
| PYTHONPATH=src python3 examples/bedrock_mantle_credential_provider.py |
There was a problem hiding this comment.
| PYTHONPATH=src python3 examples/bedrock_mantle_credential_provider.py | |
| PYTHONPATH=src python3 examples/aws_credential_provider.py |
| async def test_async_wraps_provider_error_in_openai_error(self) -> None: | ||
| def bad_provider() -> None: | ||
| raise RuntimeError("token expired") | ||
|
|
||
| with _patch_ensure_botocore(): | ||
| client = AsyncAwsOpenAI(region="us-west-2", credential_provider=bad_provider) | ||
|
|
||
| request = httpx.Request("POST", "https://example.com/v1/chat/completions", content=b'{"model":"x"}') | ||
| with pytest.raises(OpenAIError, match="Failed to refresh AWS credentials: token expired"): | ||
| await client._prepare_request(request) | ||
|
|
||
| async def test_async_wraps_async_provider_error(self) -> None: | ||
| async def bad_async_provider() -> None: | ||
| raise ValueError("async refresh failed") | ||
|
|
||
| with _patch_ensure_botocore(): | ||
| client = AsyncAwsOpenAI(region="us-west-2", credential_provider=bad_async_provider) | ||
|
|
||
| request = httpx.Request("POST", "https://example.com/v1/chat/completions", content=b'{"model":"x"}') | ||
| with pytest.raises(OpenAIError, match="Failed to refresh AWS credentials: async refresh failed"): | ||
| await client._prepare_request(request) |
There was a problem hiding this comment.
Do any of these need the @pytest.mark.asyncio marker?
There was a problem hiding this comment.
Its not needed as its enabled by default. But added it just to make sure in case that configuration changes in future.
|
|
||
| ## AWS Bedrock Mantle | ||
|
|
||
| To use this library with [AWS Bedrock Mantle](https://docs.aws.amazon.com/bedrock/), use the `AwsOpenAI` |
There was a problem hiding this comment.
We are calling out Bedrock Mantle, and linking to the main Bedrock page, is that the right page to link?
| from .lib.aws import ( | ||
| AwsOpenAI as AwsOpenAI, | ||
| AsyncAwsOpenAI as AsyncAwsOpenAI, | ||
| ) |
There was a problem hiding this comment.
src/openai/__init__.py is Stainless-generated — the AwsOpenAI / AsyncAwsOpenAI exports added there would be dropped on the next generated release unless they're added to the generator config (similar to how the AzureOpenAI exports persist today). The pyproject.toml change for the [aws] extra may have the same concern.
@apcha-oai Could you advise on how you'd like the __init__.py exports and pyproject.toml dependency handled? Should these go through the Stainless generator config, or is there a preferred approach for adding new provider integrations?
| # Resolve base_url — fall back to region-derived endpoint | ||
| if base_url is None: | ||
| if not resolved_region: | ||
| raise ValueError("Must provide base_url, or set region / AWS_REGION / AWS_DEFAULT_REGION") |
There was a problem hiding this comment.
Should we make these OpenAIError, so all credential errors raise the same error?
| import botocore.awsrequest # type: ignore[import-untyped] # pyright: ignore[reportMissingTypeStubs] | ||
|
|
||
| # Exclude httpx transport-level headers that cause SigV4 signature mismatch | ||
| _HEADERS_TO_EXCLUDE = {"connection", "accept-encoding"} |
There was a problem hiding this comment.
This is being allocated on every request. Move this to module level as a frozenset.
|
|
||
| # Sentinel API key used when SigV4 mode is active, so the base OpenAI | ||
| # constructor (which requires a non-None api_key) is satisfied. | ||
| API_KEY_SENTINEL = "<bedrock-mantle-sigv4>" |
There was a problem hiding this comment.
This looks fragile. If a user happens to pass this exact string as api_key, it silently activates SigV4 mode instead of using the key. Consider making it private (_API_KEY_SENTINEL) and using a less guessable value (e.g. a UUID) to reduce the chance of accidental collision.
| completion = client.chat.completions.create( | ||
| model="openai.gpt-oss-120b", | ||
| messages=[ | ||
| { | ||
| "role": "user", | ||
| "content": "How do I output all files in a directory using Python?", | ||
| }, | ||
| ], | ||
| ) |
There was a problem hiding this comment.
Does mantle support responses API? Should we include that first?
There was a problem hiding this comment.
Updated the example to have responses API instead.
| if use_sigv4 and credential_provider is None: | ||
| botocore_credentials = _get_default_credentials() |
There was a problem hiding this comment.
The PR description says “Credentials are resolved at request time,” but credentials are actually resolved before the first request. This means that client construction can fail or block if AWS credentials are temporarily unavailable, metadata credentials are slow or a user only expected credential lookup when making a request.
Should we rather do credential lookup inside _prepare_request() and constructor only verifies botocore is available? The same behavior might apply to _resolve_credentials_async() and tests to verify them.
| credential_provider: Any | None, | ||
| region: str | None, | ||
| base_url: str | None, | ||
| ) -> tuple[bool, Any | None, str, str, Any | None]: |
There was a problem hiding this comment.
We should use NameTuple here, it will increase the readability
| with pytest.raises(OpenAIError, match="botocore must be installed"): | ||
| AsyncAwsOpenAI(region="us-west-2") | ||
|
|
||
| def test_sync_raises_when_no_creds_resolved(self) -> None: |
There was a problem hiding this comment.
might be good to add test_async_raises_when_no_creds_resolved
| self._region, | ||
| base_url, | ||
| self._botocore_credentials, | ||
| ) = _resolve_bedrock_mantle_config( |
There was a problem hiding this comment.
Blocking credential resolution in AsyncAwsOpenAI.__init__
AsyncAwsOpenAI.__init__ is synchronous (it has to be — __init__ can't be async) and calls _get_default_credentials() → botocore.session.get_credentials() → get_frozen_credentials(). On EC2/ECS this does blocking socket I/O against IMDS or the container credentials endpoint. Inside an asyncio event loop, that freezes the loop for the duration of the round trip (10–50ms typical, seconds under IMDS throttling).
In a FastAPI/Starlette app constructing the client at startup, this is a noticeable pause. If clients are constructed per-request (common pattern with dependency injection), every incoming request blocks the loop, and under load IMDS throttling produces a livelock — the loop stops servicing any requests, not just Bedrock ones.
Users won't catch this in dev because local credential chains (env vars, ~/.aws/credentials) don't hit the network. It only manifests on the deployment targets where this library will actually run.
Suggested fix — lazy resolution wrapped with asyncio.to_thread:
class AsyncAwsOpenAI(AsyncOpenAI):
def __init__(self, ...):
# skip credential lookup here; just validate config
self._botocore_credentials = None
self._creds_lock = asyncio.Lock()
async def _prepare_request(self, request):
if not self._use_sigv4:
return
if self._credential_provider is None and self._botocore_credentials is None:
async with self._creds_lock:
if self._botocore_credentials is None:
self._botocore_credentials = await asyncio.to_thread(_get_default_credentials)
credentials = await _resolve_credentials_async(self._credential_provider, self._botocore_credentials)
_sign_httpx_request(request, credentials, self._region)asyncio.to_thread offloads the blocking botocore call to a worker thread so the event loop stays free. Also applies to the per-request get_frozen_credentials() path, which can trigger a refresh call against IMDS.
Async client should not freeze the event loop on construction.
Add AwsOpenAI and AsyncAwsOpenAI clients that support AWS SigV4 request
signing for Bedrock Mantle APIs, alongside standard API key auth.
Key changes:
- Add src/openai/lib/aws.py with sync and async client classes that
sign requests using botocore's SigV4Auth
- Support custom credential providers (sync and async) as well as
automatic credential resolution via the default botocore chain
- Export AwsOpenAI and AsyncAwsOpenAI from the top-level package
- Add examples for basic usage and STS assume-role credential refresh
- Add comprehensive test suite covering SigV4 signing, credential
resolution, API key fallback, and copy/with_options behavior
- Add botocore as a dev dependency in pyproject.toml
- Fix all pyright and mypy lint errors for botocore type stubs
- Make API key sentinel private with non-guessable value and identity check - Use OpenAIError consistently instead of ValueError for config errors - Fix incorrect filenames in example docstrings - Update README with correct Bedrock Mantle docs link and Responses API example - Move _HEADERS_TO_EXCLUDE to module-level frozenset - Defer credential resolution from __init__ to _prepare_request() - Use asyncio.to_thread in AsyncAwsOpenAI to avoid blocking the event loop - Use NamedTuple for _resolve_bedrock_mantle_config return type - Add test_async_raises_when_no_creds_resolved test
f3f8097 to
65a2707
Compare
Purpose
Add
AwsOpenAIandAsyncAwsOpenAIclient classes that enable the OpenAI Python SDK to work with AWS Bedrock Mantle APIs using SigV4 request signing. This allows customers to use their existing AWS credentials (IAM roles, env vars,~/.aws/credentials) to authenticate with Bedrock Mantle endpoints without managing separate API keys.Approach
src/openai/lib/aws.py: ContainsAwsOpenAI(sync) andAsyncAwsOpenAI(async) classes that extend the baseOpenAI/AsyncOpenAIclients. They override_prepare_requestto inject SigV4 signatures on every outgoing HTTP request using botocore'sSigV4Auth.api_keydisables SigV4.RefreshableCredentials).AWS_REGION/AWS_DEFAULT_REGIONenv vars. The base URL is derived ashttps://bedrock-mantle.{region}.api.aws/v1when not explicitly provided.AwsOpenAIandAsyncAwsOpenAIare exported fromopenai.__init__for convenient imports.botocoreis an optional dependency under the[aws]extra (pip install openai[aws]).Files Changed
src/openai/lib/aws.pysrc/openai/__init__.pyAwsOpenAIandAsyncAwsOpenAIpyproject.tomlbotocoreas optional[aws]dependency and dev dependencytests/lib/test_aws.pyexamples/aws_client.pyexamples/aws_credential_provider.pyREADME.mdTests Performed
Unit Tests (40/40 passing)
Authorization,X-Amz-Date,X-Amz-Security-Tokenheaders are injected onchat.completionsandresponses.createcallsapi_key+credential_providerraises error; no region + no base_url raises errorAWS_REGIONandAWS_DEFAULT_REGIONenv var resolution with correct precedenceAwsOpenAIis instance ofOpenAI,AsyncAwsOpenAIis instance ofAsyncOpenAIIntegration Tests (manual)
examples/aws_client.pyagainst live Bedrock Mantle endpoint in us-west-2:examples/aws_credential_provider.py— confirmed SigV4 signing works (401 expected with placeholder credentials)