feat: Add AWS Bedrock Mantle client with SigV4 authentication

prashanthkvs · prashanthkvs · commit f3f809760c27 · 2026-04-29T08:08:15.000-07:00
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
diff --git a/README.md b/README.md
@@ -933,6 +933,41 @@ In addition to the options provided in the base `OpenAI` client, the following o
 
 An example of using the client with Microsoft Entra ID (formerly known as Azure Active Directory) can be found [here](https://github.com/openai/openai-python/blob/main/examples/azure_ad.py).
 
+## AWS Bedrock Mantle
+
+To use this library with [AWS Bedrock Mantle](https://docs.aws.amazon.com/bedrock/), use the `AwsOpenAI`
+class instead of the `OpenAI` class.
+
+> [!IMPORTANT]
+> This requires `botocore` to be installed for SigV4 request signing. Install it with: `pip install 'openai[aws]'`
+
+```py
+from openai import AwsOpenAI
+
+# uses the default botocore credential chain (env vars, ~/.aws/credentials, IAM role, etc.)
+client = AwsOpenAI(
+    region="us-west-2",
+)
+
+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?",
+        },
+    ],
+)
+print(completion.choices[0].message.content)
+```
+
+In addition to the options provided in the base `OpenAI` client, the following options are provided:
+
+- `region` (or the `AWS_REGION` / `AWS_DEFAULT_REGION` environment variable)
+- `credential_provider` - a callable that returns credentials with `access_key`, `secret_key`, and optional `token` attributes
+
+An example of using the client with a custom credential provider and STS assume-role refresh can be found [here](https://github.com/openai/openai-python/blob/main/examples/aws_credential_provider.py).
+
 ## 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/examples/aws_client.py b/examples/aws_client.py
@@ -0,0 +1,80 @@
+"""Example: Using AwsOpenAI (sync) and AsyncAwsOpenAI (async) with SigV4 signing.
+
+Requires:
+  - botocore installed (pip install botocore)
+  - AWS credentials configured (env vars, ~/.aws/credentials, IAM role, etc.)
+  - AWS_REGION or AWS_DEFAULT_REGION set (or pass region= explicitly)
+
+Run:
+  export AWS_REGION=us-west-2
+  PYTHONPATH=src python3 examples/bedrock_mantle.py
+"""
+
+import asyncio
+
+from openai.lib.aws import AwsOpenAI, AsyncAwsOpenAI
+
+# --- Synchronous usage ---
+
+client = AwsOpenAI(region="us-west-2")
+
+response = client.chat.completions.create(
+    model="openai.gpt-oss-120b",
+    messages=[{"role": "user", "content": "Hello, how are you?"}],
+)
+
+print("Sync:", response.choices[0].message.content)
+
+
+# --- Asynchronous usage ---
+
+
+async def main() -> None:
+    async_client = AsyncAwsOpenAI(region="us-west-2")
+
+    response = await async_client.chat.completions.create(
+        model="openai.gpt-oss-120b",
+        messages=[{"role": "user", "content": "Hello from async!"}],
+    )
+
+    print("Async:", response.choices[0].message.content)
+
+
+asyncio.run(main())
+
+
+# --- Streaming usage (sync) ---
+
+print("\nStreaming: ", end="")
+stream = client.chat.completions.create(
+    model="openai.gpt-oss-120b",
+    messages=[{"role": "user", "content": "Count from 1 to 5."}],
+    stream=True,
+)
+for chunk in stream:
+    delta = chunk.choices[0].delta.content
+    if delta:
+        print(delta, end="", flush=True)
+print()
+
+
+# --- Streaming usage (async) ---
+
+
+async def stream_async() -> None:
+    async_client = AsyncAwsOpenAI(region="us-west-2")
+
+    print("Async streaming: ", end="")
+    stream = await async_client.chat.completions.create(
+        model="openai.gpt-oss-120b",
+        messages=[{"role": "user", "content": "Count from 1 to 5."}],
+        stream=True,
+    )
+    async for chunk in stream:
+        delta = chunk.choices[0].delta.content
+        if delta:
+            print(delta, end="", flush=True)
+    print()
+
+
+asyncio.run(stream_async())
diff --git a/examples/aws_credential_provider.py b/examples/aws_credential_provider.py
@@ -0,0 +1,144 @@
+"""Example: Using AwsOpenAI with a custom credential provider and auto-refresh.
+
+This shows how to:
+  1. Use a custom credential provider that returns fresh credentials on each call
+  2. Use botocore's RefreshableCredentials for automatic STS assume-role refresh
+  3. Use an async credential provider with AsyncAwsOpenAI
+
+Requires:
+  - botocore installed (pip install botocore)
+  - boto3 installed (pip install boto3) — for the STS assume-role example
+  - AWS credentials configured for the initial session
+  - AWS_REGION or AWS_DEFAULT_REGION set (or pass region= explicitly)
+
+Run:
+  export AWS_REGION=us-west-2
+  PYTHONPATH=src python3 examples/bedrock_mantle_credential_provider.py
+"""
+
+from __future__ import annotations
+
+import asyncio
+from typing import Any, Callable
+from dataclasses import dataclass
+
+from openai.lib.aws import AwsOpenAI, AsyncAwsOpenAI
+
+# ---------------------------------------------------------------------------
+# 1. Simple custom credential provider
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class MyCredentials:
+    """Minimal object satisfying the Credentials protocol."""
+
+    access_key: str
+    secret_key: str
+    token: str | None = None
+
+
+def my_credential_provider() -> MyCredentials:
+    """Return credentials from your own secret store, vault, etc.
+
+    This callable is invoked before every request, so returning fresh
+    credentials here is all you need for auto-refresh.
+    """
+    # Replace with your actual credential fetching logic
+    return MyCredentials(
+        access_key="AKIA...",
+        secret_key="wJalr...",
+        token="FwoGZX...",  # optional session token
+    )
+
+
+client = AwsOpenAI(
+    region="us-west-2",
+    credential_provider=my_credential_provider,
+)
+
+response = client.chat.completions.create(
+    model="openai.gpt-oss-120b",
+    messages=[{"role": "user", "content": "Hello from custom credentials!"}],
+)
+print("Custom provider:", response.choices[0].message.content)
+
+
+# ---------------------------------------------------------------------------
+# 2. Auto-refreshing STS assume-role credentials via botocore
+# ---------------------------------------------------------------------------
+
+
+def make_sts_credential_provider(role_arn: str, session_name: str = "bedrock-mantle") -> Callable[[], Any]:
+    """Create a credential provider that assumes an IAM role and auto-refreshes.
+
+    botocore's RefreshableCredentials handles expiry checks and refresh
+    transparently — accessing .access_key / .secret_key / .token on the
+    returned object triggers a refresh if the credentials are expired.
+    """
+    import botocore.session  # type: ignore[import-untyped, import-not-found]
+    import botocore.credentials  # type: ignore[import-untyped, import-not-found]
+
+    session: Any = botocore.session.get_session()  # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType]
+    sts: Any = session.create_client("sts")  # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType]
+
+    def fetch_credentials() -> dict[str, Any]:
+        resp: Any = sts.assume_role(RoleArn=role_arn, RoleSessionName=session_name)["Credentials"]  # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType]
+        return {
+            "access_key": resp["AccessKeyId"],
+            "secret_key": resp["SecretAccessKey"],
+            "token": resp["SessionToken"],
+            "expiry_time": resp["Expiration"].isoformat(),  # pyright: ignore[reportUnknownMemberType]
+        }
+
+    refreshable: Any = botocore.credentials.RefreshableCredentials.create_from_metadata(  # pyright: ignore[reportUnknownVariableType, reportUnknownMemberType]
+        metadata=fetch_credentials(),
+        refresh_using=fetch_credentials,
+        method="sts-assume-role",
+    )
+
+    # Return a provider that gives back the refreshable object.
+    # Accessing its attributes auto-refreshes when expired.
+    def provider() -> Any:
+        return refreshable  # pyright: ignore[reportUnknownVariableType]
+
+    return provider
+
+
+# Uncomment to use:
+# sts_client = AwsOpenAI(
+#     region="us-west-2",
+#     credential_provider=make_sts_credential_provider("arn:aws:iam::123456789012:role/MyRole"),
+# )
+
+
+# ---------------------------------------------------------------------------
+# 3. Async credential provider
+# ---------------------------------------------------------------------------
+
+
+async def async_credential_provider() -> MyCredentials:
+    """An async provider — useful when credentials come from an async API."""
+    # Simulate async credential fetch (e.g., from an async HTTP vault client)
+    await asyncio.sleep(0)
+    return MyCredentials(
+        access_key="AKIA...",
+        secret_key="wJalr...",
+        token="FwoGZX...",
+    )
+
+
+async def main() -> None:
+    async_client = AsyncAwsOpenAI(
+        region="us-west-2",
+        credential_provider=async_credential_provider,
+    )
+
+    response = await async_client.chat.completions.create(
+        model="openai.gpt-oss-120b",
+        messages=[{"role": "user", "content": "Hello from async credentials!"}],
+    )
+    print("Async provider:", response.choices[0].message.content)
+
+
+asyncio.run(main())
diff --git a/pyproject.toml b/pyproject.toml
@@ -50,6 +50,7 @@ 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"]
+aws = ["botocore >= 1.29.0"]
 
 [tool.rye]
 managed = true
@@ -68,6 +69,7 @@ dev-dependencies = [
     "rich>=13.7.1",
     "inline-snapshot>=0.28.0",
     "azure-identity >=1.14.1",
+    "botocore >=1.29.0",
     "types-tqdm > 4",
     "types-pyaudio > 0",
     "trio >=0.22.2",
diff --git a/src/openai/__init__.py b/src/openai/__init__.py
@@ -97,6 +97,10 @@
     from ._utils._resources_proxy import resources as resources
 
 from .lib import azure as _azure, pydantic_function_tool as pydantic_function_tool
+from .lib.aws import (
+    AwsOpenAI as AwsOpenAI,
+    AsyncAwsOpenAI as AsyncAwsOpenAI,
+)
 from .version import VERSION as VERSION
 from .lib.azure import AzureOpenAI as AzureOpenAI, AsyncAzureOpenAI as AsyncAzureOpenAI
 from .lib._old_api import *
diff --git a/src/openai/lib/aws.py b/src/openai/lib/aws.py
diff --git a/tests/lib/test_aws.py b/tests/lib/test_aws.py
