diff --git a/README.md b/README.md index 7af6143c88..6f87246470 100644 --- a/README.md +++ b/README.md @@ -926,6 +926,37 @@ 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). +## Amazon Bedrock + +To use this library with [Amazon Bedrock's OpenAI-compatible API](https://docs.aws.amazon.com/bedrock/latest/userguide/models-api-compatibility.html), use the `BedrockOpenAI` class instead of the `OpenAI` class. + +```py +from openai import BedrockOpenAI + +# gets the bearer token from AWS_BEARER_TOKEN_BEDROCK and the region from AWS_REGION/AWS_DEFAULT_REGION +client = BedrockOpenAI() + +response = client.responses.create( + model="openai.gpt-5.4", + input="Say hello!", +) + +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. + +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`. + +Set `AWS_BEARER_TOKEN_BEDROCK` to an [Amazon Bedrock API key](https://docs.aws.amazon.com/bedrock/latest/userguide/api-keys.html). To refresh tokens yourself, pass a provider instead of `api_key`: + +```py +client = BedrockOpenAI( + aws_region="us-west-2", + bedrock_token_provider=lambda: refresh_bedrock_token(), +) +``` + ## 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/bedrock.py b/examples/bedrock.py new file mode 100644 index 0000000000..24dafb5b80 --- /dev/null +++ b/examples/bedrock.py @@ -0,0 +1,13 @@ +from openai import BedrockOpenAI + +client = BedrockOpenAI() + +# For refreshed Bedrock bearer tokens: +# client = BedrockOpenAI(aws_region="us-west-2", bedrock_token_provider=get_bedrock_token) + +response = client.responses.create( + model="openai.gpt-5.4", + input="Say hello!", +) + +print(response.output_text) diff --git a/src/openai/__init__.py b/src/openai/__init__.py index cbaef0615f..8d4265970e 100644 --- a/src/openai/__init__.py +++ b/src/openai/__init__.py @@ -79,6 +79,8 @@ "AsyncStream", "OpenAI", "AsyncOpenAI", + "BedrockOpenAI", + "AsyncBedrockOpenAI", "file_from_path", "BaseModel", "DEFAULT_TIMEOUT", @@ -96,9 +98,10 @@ if not _t.TYPE_CHECKING: from ._utils._resources_proxy import resources as resources -from .lib import azure as _azure, pydantic_function_tool as pydantic_function_tool +from .lib import azure as _azure, bedrock as _bedrock, pydantic_function_tool as pydantic_function_tool from .version import VERSION as VERSION from .lib.azure import AzureOpenAI as AzureOpenAI, AsyncAzureOpenAI as AsyncAzureOpenAI +from .lib.bedrock import BedrockOpenAI as BedrockOpenAI, AsyncBedrockOpenAI as AsyncBedrockOpenAI from .lib._old_api import * from .lib.streaming import ( AssistantEventHandler as AssistantEventHandler, @@ -150,7 +153,7 @@ http_client: _httpx.Client | None = None -_ApiType = _te.Literal["openai", "azure"] +_ApiType = _te.Literal["openai", "azure", "amazon-bedrock"] api_type: _ApiType | None = _t.cast(_ApiType, _os.environ.get("OPENAI_API_TYPE")) @@ -162,6 +165,10 @@ azure_ad_token_provider: _azure.AzureADTokenProvider | None = None +_bedrock_api_key: str | None = None + +bedrock_token_provider: _bedrock.BedrockTokenProvider | None = None + class _ModuleClient(OpenAI): # Note: we have to use type: ignores here as overriding class members @@ -294,10 +301,23 @@ class _AzureModuleClient(_ModuleClient, AzureOpenAI): # type: ignore ... +class _BedrockModuleClient(_ModuleClient, BedrockOpenAI): # type: ignore + @property # type: ignore + @override + def api_key(self) -> str | None: + return _bedrock_api_key if _bedrock_api_key is not None else api_key + + @api_key.setter # type: ignore + def api_key(self, value: str | None) -> None: # type: ignore + global _bedrock_api_key + + _bedrock_api_key = value + + class _AmbiguousModuleClientUsageError(OpenAIError): def __init__(self) -> None: super().__init__( - "Ambiguous use of module client; please set `openai.api_type` or the `OPENAI_API_TYPE` environment variable to `openai` or `azure`" + "Ambiguous use of module client; please set `openai.api_type` or the `OPENAI_API_TYPE` environment variable to `openai`, `azure`, or `amazon-bedrock`" ) @@ -370,6 +390,22 @@ def _load_client() -> OpenAI: # type: ignore[reportUnusedFunction] ) return _client + if api_type == "amazon-bedrock": + _client = _BedrockModuleClient( # type: ignore + api_key=api_key, + bedrock_token_provider=bedrock_token_provider, + organization=organization, + project=project, + webhook_secret=webhook_secret, + base_url=base_url, + timeout=timeout, + max_retries=max_retries, + default_headers=default_headers, + default_query=default_query, + http_client=http_client, + ) + return _client + _client = _ModuleClient( api_key=api_key, admin_api_key=admin_api_key, diff --git a/src/openai/lib/bedrock.py b/src/openai/lib/bedrock.py new file mode 100644 index 0000000000..0a1aec36ee --- /dev/null +++ b/src/openai/lib/bedrock.py @@ -0,0 +1,417 @@ +from __future__ import annotations + +import os +import re +import inspect +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 .._utils import is_given +from .._client import OpenAI, AsyncOpenAI +from .._models import SecurityOptions, FinalRequestOptions +from .._exceptions import OpenAIError +from .._base_client import DEFAULT_MAX_RETRIES + +BedrockTokenProvider = Callable[[], str] +AsyncBedrockTokenProvider = Callable[[], "str | Awaitable[str]"] + + +def _normalize_bedrock_base_url(base_url: str | httpx.URL) -> httpx.URL: + """Normalize a Bedrock Responses URL variant back to the provider API root.""" + url = httpx.URL(base_url) + path = url.path.rstrip("/") + responses_match = re.search(r"/responses(?:/.*)?$", path) + if responses_match is not None: + path = path[: responses_match.start()] + + return url.copy_with(path=path or "/") + + +def _resolve_bedrock_base_url(base_url: str | httpx.URL | None, aws_region: str | None) -> httpx.URL: + """Resolve Bedrock base URL precedence from explicit, env, then region config.""" + if isinstance(base_url, str) and not base_url.strip(): + base_url = None + + if base_url is None: + env_base_url = os.environ.get("AWS_BEDROCK_BASE_URL") + if env_base_url is not None and env_base_url.strip(): + base_url = env_base_url + + if base_url is None: + 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( + "Must provide one of the `base_url` or `aws_region` arguments, or set the " + "`AWS_BEDROCK_BASE_URL`, `AWS_REGION`, or `AWS_DEFAULT_REGION` environment variable." + ) + + base_url = f"https://bedrock-mantle.{region}.api.aws/openai/v1" + + return _normalize_bedrock_base_url(base_url) + + +def _bedrock_token_provider(provider: BedrockTokenProvider) -> BedrockTokenProvider: + """Adapt a sync Bedrock token provider to the base client's api_key callback.""" + + def get_token() -> str: + token = cast(object, provider()) + if not isinstance(token, str) or not token: + raise ValueError(f"Expected `bedrock_token_provider` argument to return a string but it returned {token}") + + return token + + return get_token + + +def _async_bedrock_token_provider(provider: AsyncBedrockTokenProvider) -> Callable[[], Awaitable[str]]: + """Adapt a sync or async Bedrock token provider to the async api_key callback.""" + + async def get_token() -> str: + token = cast(object, provider()) + if inspect.isawaitable(token): + token = await token + + if not isinstance(token, str) or not token: + raise ValueError(f"Expected `bedrock_token_provider` argument to return a string but it returned {token}") + + return token + + return get_token + + +class BedrockOpenAI(OpenAI): + """API client for Amazon Bedrock's OpenAI-compatible endpoint.""" + + _bedrock_token_provider: BedrockTokenProvider | None + aws_region: str | None + + def __init__( + self, + *, + api_key: str | None = None, + bedrock_token_provider: BedrockTokenProvider | None = None, + aws_region: str | None = None, + organization: str | None = None, + project: str | None = None, + webhook_secret: str | None = None, + base_url: str | httpx.URL | None = None, + websocket_base_url: str | httpx.URL | None = None, + timeout: float | Timeout | None | NotGiven = NOT_GIVEN, + max_retries: int = DEFAULT_MAX_RETRIES, + default_headers: Mapping[str, str] | None = None, + default_query: Mapping[str, object] | None = None, + http_client: httpx.Client | None = None, + _strict_response_validation: bool = False, + _enforce_credentials: bool = True, + ) -> None: + """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` + - `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. + """ + 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 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." + ) + + self._bedrock_token_provider = bedrock_token_provider + self.aws_region = aws_region + + super().__init__( + api_key=_bedrock_token_provider(bedrock_token_provider) + if bedrock_token_provider is not None + else api_key or "", + admin_api_key="", + organization=organization, + project=project, + webhook_secret=webhook_secret, + base_url=_resolve_bedrock_base_url(base_url, aws_region), + websocket_base_url=websocket_base_url, + timeout=timeout, + max_retries=max_retries, + default_headers=default_headers, + default_query=default_query, + http_client=http_client, + _strict_response_validation=_strict_response_validation, + _enforce_credentials=False, + ) + + @override + def _auth_headers(self, security: SecurityOptions) -> dict[str, str]: + if security.get("bearer_auth", False) or security.get("admin_api_key_auth", False): + return self._bearer_auth + + return {} + + @override + def _prepare_options(self, options: FinalRequestOptions) -> FinalRequestOptions: + if ( + self._api_key_provider is not None + and options.security.get("admin_api_key_auth", False) + and not options.security.get("bearer_auth", False) + ): + self._refresh_api_key() + + return super()._prepare_options(options) + + @override + def copy( + self, + *, + api_key: str | BedrockTokenProvider | None = None, + admin_api_key: str | None = None, + workload_identity: WorkloadIdentity | None = None, + bedrock_token_provider: BedrockTokenProvider | None = None, + aws_region: str | None = None, + organization: str | None = None, + project: str | None = None, + webhook_secret: str | None = None, + websocket_base_url: str | httpx.URL | None = None, + base_url: str | httpx.URL | None = None, + timeout: float | Timeout | None | NotGiven = NOT_GIVEN, + http_client: httpx.Client | None = None, + max_retries: int | NotGiven = NOT_GIVEN, + default_headers: Mapping[str, str] | None = None, + set_default_headers: Mapping[str, str] | None = None, + default_query: Mapping[str, object] | None = None, + set_default_query: Mapping[str, object] | None = None, + _enforce_credentials: bool | None = None, + _extra_kwargs: Mapping[str, Any] = {}, + ) -> Self: + if default_headers is not None and set_default_headers is not None: + raise ValueError("The `default_headers` and `set_default_headers` arguments are mutually exclusive") + + if default_query is not None and set_default_query is not None: + raise ValueError("The `default_query` and `set_default_query` arguments are mutually exclusive") + + if callable(api_key): + 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.") + + 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.") + + headers = self._custom_headers + if default_headers is not None: + headers = {**headers, **default_headers} + elif set_default_headers is not None: + headers = set_default_headers + + params = self._custom_query + if default_query is not None: + params = {**params, **default_query} + elif set_default_query is not None: + params = set_default_query + + next_token_provider = ( + bedrock_token_provider if bedrock_token_provider is not None else 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) + + return self.__class__( + 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, + 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, + websocket_base_url=websocket_base_url if websocket_base_url is not None else self.websocket_base_url, + base_url=base_url if base_url is not None else self.base_url, + timeout=self.timeout if isinstance(timeout, NotGiven) else timeout, + http_client=http_client or self._client, + max_retries=max_retries if is_given(max_retries) else self.max_retries, + default_headers=headers, + default_query=params, + _enforce_credentials=True if _enforce_credentials is None else _enforce_credentials, + **_extra_kwargs, + ) + + with_options = copy + + +class AsyncBedrockOpenAI(AsyncOpenAI): + """Async API client for Amazon Bedrock's OpenAI-compatible endpoint.""" + + _bedrock_token_provider: AsyncBedrockTokenProvider | None + aws_region: str | None + + def __init__( + self, + *, + api_key: str | None = None, + bedrock_token_provider: AsyncBedrockTokenProvider | None = None, + aws_region: str | None = None, + organization: str | None = None, + project: str | None = None, + webhook_secret: str | None = None, + base_url: str | httpx.URL | None = None, + websocket_base_url: str | httpx.URL | None = None, + timeout: float | Timeout | None | NotGiven = NOT_GIVEN, + max_retries: int = DEFAULT_MAX_RETRIES, + default_headers: Mapping[str, str] | None = None, + default_query: Mapping[str, object] | None = None, + http_client: httpx.AsyncClient | None = None, + _strict_response_validation: bool = False, + _enforce_credentials: bool = True, + ) -> None: + """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` + - `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. + """ + 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 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." + ) + + self._bedrock_token_provider = bedrock_token_provider + self.aws_region = aws_region + + super().__init__( + api_key=( + _async_bedrock_token_provider(bedrock_token_provider) + if bedrock_token_provider is not None + else api_key or "" + ), + admin_api_key="", + organization=organization, + project=project, + webhook_secret=webhook_secret, + base_url=_resolve_bedrock_base_url(base_url, aws_region), + websocket_base_url=websocket_base_url, + timeout=timeout, + max_retries=max_retries, + default_headers=default_headers, + default_query=default_query, + http_client=http_client, + _strict_response_validation=_strict_response_validation, + _enforce_credentials=False, + ) + + @override + def _auth_headers(self, security: SecurityOptions) -> dict[str, str]: + if security.get("bearer_auth", False) or security.get("admin_api_key_auth", False): + return self._bearer_auth + + return {} + + @override + async def _prepare_options(self, options: FinalRequestOptions) -> FinalRequestOptions: + if ( + self._api_key_provider is not None + and options.security.get("admin_api_key_auth", False) + and not options.security.get("bearer_auth", False) + ): + await self._refresh_api_key() + + return await super()._prepare_options(options) + + @override + def copy( + self, + *, + api_key: str | AsyncBedrockTokenProvider | None = None, + admin_api_key: str | None = None, + workload_identity: WorkloadIdentity | None = None, + bedrock_token_provider: AsyncBedrockTokenProvider | None = None, + aws_region: str | None = None, + organization: str | None = None, + project: str | None = None, + webhook_secret: str | None = None, + websocket_base_url: str | httpx.URL | None = None, + base_url: str | httpx.URL | None = None, + timeout: float | Timeout | None | NotGiven = NOT_GIVEN, + http_client: httpx.AsyncClient | None = None, + max_retries: int | NotGiven = NOT_GIVEN, + default_headers: Mapping[str, str] | None = None, + set_default_headers: Mapping[str, str] | None = None, + default_query: Mapping[str, object] | None = None, + set_default_query: Mapping[str, object] | None = None, + _enforce_credentials: bool | None = None, + _extra_kwargs: Mapping[str, Any] = {}, + ) -> Self: + if default_headers is not None and set_default_headers is not None: + raise ValueError("The `default_headers` and `set_default_headers` arguments are mutually exclusive") + + if default_query is not None and set_default_query is not None: + raise ValueError("The `default_query` and `set_default_query` arguments are mutually exclusive") + + if callable(api_key): + 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.") + + 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.") + + headers = self._custom_headers + if default_headers is not None: + headers = {**headers, **default_headers} + elif set_default_headers is not None: + headers = set_default_headers + + params = self._custom_query + if default_query is not None: + params = {**params, **default_query} + elif set_default_query is not None: + params = set_default_query + + next_token_provider = ( + bedrock_token_provider if bedrock_token_provider is not None else 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) + + return self.__class__( + 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, + 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, + websocket_base_url=websocket_base_url if websocket_base_url is not None else self.websocket_base_url, + base_url=base_url if base_url is not None else self.base_url, + timeout=self.timeout if isinstance(timeout, NotGiven) else timeout, + http_client=http_client or self._client, + max_retries=max_retries if is_given(max_retries) else self.max_retries, + default_headers=headers, + default_query=params, + _enforce_credentials=True if _enforce_credentials is None else _enforce_credentials, + **_extra_kwargs, + ) + + with_options = copy diff --git a/tests/lib/test_bedrock.py b/tests/lib/test_bedrock.py new file mode 100644 index 0000000000..9995e4c431 --- /dev/null +++ b/tests/lib/test_bedrock.py @@ -0,0 +1,439 @@ +from __future__ import annotations + +import json +from typing import Any, Union, Protocol, cast + +import httpx +import pytest +from httpx import URL +from respx import MockRouter + +from openai import OpenAIError, NotFoundError +from tests.utils import update_env +from openai._types import Omit +from openai.lib.bedrock import BedrockOpenAI, AsyncBedrockOpenAI + +Client = Union[BedrockOpenAI, AsyncBedrockOpenAI] + +RESPONSE_BODY: dict[str, Any] = { + "id": "resp_123", + "object": "response", + "created_at": 0, + "status": "completed", + "background": False, + "error": None, + "incomplete_details": None, + "instructions": None, + "max_output_tokens": None, + "max_tool_calls": None, + "model": "gpt-4o", + "output": [], + "parallel_tool_calls": True, + "previous_response_id": None, + "prompt_cache_key": None, + "reasoning": {"effort": None, "summary": None}, + "safety_identifier": None, + "service_tier": "default", + "store": True, + "temperature": 1.0, + "text": {"format": {"type": "text"}, "verbosity": "medium"}, + "tool_choice": "auto", + "tools": [], + "top_logprobs": 0, + "top_p": 1.0, + "truncation": "disabled", + "usage": { + "input_tokens": 0, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens": 0, + "output_tokens_details": {"reasoning_tokens": 0}, + "total_tokens": 0, + }, + "user": None, + "metadata": {}, +} +COMPACTED_RESPONSE_BODY: dict[str, Any] = { + "id": "resp_123", + "created_at": 0, + "object": "response.compaction", + "output": [], + "usage": RESPONSE_BODY["usage"], +} +INPUT_ITEMS_BODY: dict[str, Any] = { + "object": "list", + "data": [], + "first_id": "item_123", + "last_id": "item_123", + "has_more": False, +} +INPUT_TOKENS_BODY: dict[str, Any] = { + "object": "response.input_tokens", + "input_tokens": 1, +} + + +class MockRequestCall(Protocol): + request: httpx.Request + + +def make_sync_client(**kwargs: Any) -> BedrockOpenAI: + return BedrockOpenAI(http_client=httpx.Client(trust_env=False), **kwargs) + + +def make_async_client(**kwargs: Any) -> AsyncBedrockOpenAI: + return AsyncBedrockOpenAI(http_client=httpx.AsyncClient(trust_env=False), **kwargs) + + +def response_created_sse() -> str: + event: dict[str, Any] = {"type": "response.created", "sequence_number": 0, "response": RESPONSE_BODY} + return f"event: response.created\ndata: {json.dumps(event)}\n\n" + + +@pytest.mark.parametrize("client_cls", [BedrockOpenAI, AsyncBedrockOpenAI]) +def test_region_derived_base_url(client_cls: type[Client]) -> None: + with update_env(AWS_BEDROCK_BASE_URL=Omit(), AWS_REGION="us-east-1", AWS_DEFAULT_REGION=Omit()): + client = ( + make_sync_client(api_key="token") if client_cls is BedrockOpenAI else make_async_client(api_key="token") + ) + + assert client.base_url == URL("https://bedrock-mantle.us-east-1.api.aws/openai/v1/") + + +@pytest.mark.parametrize("client_cls", [BedrockOpenAI, AsyncBedrockOpenAI]) +def test_bedrock_config_precedence(client_cls: type[Client]) -> None: + with update_env( + AWS_BEDROCK_BASE_URL="https://env.example.com/openai/v1", + AWS_BEARER_TOKEN_BEDROCK="env token", + AWS_REGION="us-east-1", + AWS_DEFAULT_REGION="us-west-2", + ): + client = ( + make_sync_client( + base_url="https://explicit.example.com/openai/v1/responses", + api_key="explicit token", + ) + if client_cls is BedrockOpenAI + else make_async_client( + base_url="https://explicit.example.com/openai/v1/responses", + api_key="explicit token", + ) + ) + + assert client.base_url == URL("https://explicit.example.com/openai/v1/") + assert client.api_key == "explicit token" + + +@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"): + explicit_region_client = ( + make_sync_client(aws_region="eu-west-1", api_key="token") + if client_cls is BedrockOpenAI + else make_async_client(aws_region="eu-west-1", api_key="token") + ) + aws_region_client = ( + make_sync_client(api_key="token") if client_cls is BedrockOpenAI else make_async_client(api_key="token") + ) + + with update_env(AWS_BEDROCK_BASE_URL=Omit(), AWS_REGION=Omit(), AWS_DEFAULT_REGION="us-west-2"): + default_region_client = ( + make_sync_client(api_key="token") if client_cls is BedrockOpenAI else make_async_client(api_key="token") + ) + + assert explicit_region_client.base_url == URL("https://bedrock-mantle.eu-west-1.api.aws/openai/v1/") + assert aws_region_client.base_url == URL("https://bedrock-mantle.us-east-1.api.aws/openai/v1/") + assert default_region_client.base_url == URL("https://bedrock-mantle.us-west-2.api.aws/openai/v1/") + + +@pytest.mark.parametrize("client_cls", [BedrockOpenAI, AsyncBedrockOpenAI]) +def test_normalizes_responses_url(client_cls: type[Client]) -> None: + client = ( + make_sync_client(base_url="https://example.com/openai/v1/responses", api_key="token") + if client_cls is BedrockOpenAI + else make_async_client(base_url="https://example.com/openai/v1/responses", api_key="token") + ) + + assert client.base_url == URL("https://example.com/openai/v1/") + + +@pytest.mark.parametrize("client_cls", [BedrockOpenAI, AsyncBedrockOpenAI]) +def test_requires_endpoint_configuration(client_cls: type[Client]) -> None: + with update_env(AWS_BEDROCK_BASE_URL=Omit(), AWS_REGION=Omit(), AWS_DEFAULT_REGION=Omit()): + with pytest.raises(OpenAIError, match="Must provide one of the `base_url` or `aws_region`"): + client_cls(api_key="token") + + +@pytest.mark.parametrize("client_cls", [BedrockOpenAI, AsyncBedrockOpenAI]) +def test_does_not_use_openai_api_key(client_cls: type[Client]) -> None: + with update_env( + OPENAI_API_KEY="openai token", + 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() + + +@pytest.mark.parametrize("client_cls", [BedrockOpenAI, AsyncBedrockOpenAI]) +def test_rejects_static_token_and_provider(client_cls: type[Client]) -> None: + with pytest.raises(OpenAIError, match="mutually exclusive"): + client_cls( + base_url="https://example.com/openai/v1", + api_key="token", + bedrock_token_provider=lambda: "provider token", + ) + + +@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"): + client_cls( + base_url="https://example.com/openai/v1", + api_key=lambda: "provider token", # type: ignore[arg-type] + ) + + +@pytest.mark.respx() +def test_token_provider_refresh_sync(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), + ] + ) + tokens = iter(["first", "second"]) + client = BedrockOpenAI( + base_url="https://example.com/openai/v1", + bedrock_token_provider=lambda: next(tokens), + 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 calls[0].request.headers["Authorization"] == "Bearer first" + assert calls[1].request.headers["Authorization"] == "Bearer second" + + +@pytest.mark.asyncio +@pytest.mark.respx() +async def test_token_provider_refresh_async(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), + ] + ) + tokens = iter(["first", "second"]) + client = AsyncBedrockOpenAI( + base_url="https://example.com/openai/v1", + bedrock_token_provider=lambda: next(tokens), + http_client=httpx.AsyncClient(trust_env=False), + max_retries=1, + ) + + await client.responses.create(model="gpt-4o", input="hello") + + calls = cast("list[MockRequestCall]", respx_mock.calls) + assert calls[0].request.headers["Authorization"] == "Bearer first" + assert calls[1].request.headers["Authorization"] == "Bearer second" + + +def test_preserves_token_provider_across_with_options() -> None: + client = BedrockOpenAI( + base_url="https://example.com/openai/v1", + bedrock_token_provider=lambda: "provider token", + http_client=httpx.Client(trust_env=False), + ) + + copied_client = client.with_options(timeout=1) + + assert copied_client._refresh_api_key() == "provider token" + + +@pytest.mark.parametrize( + "copy_kwargs", + [ + {"admin_api_key": "admin token"}, + {"workload_identity": cast(Any, {})}, + ], +) +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"): + client.with_options(**copy_kwargs) + + +@pytest.mark.respx() +def test_passes_non_responses_resources_through(respx_mock: MockRouter) -> None: + respx_mock.post("https://example.com/openai/v1/chat/completions").mock( + return_value=httpx.Response( + 404, + json={"error": {"message": "AWS does not support chat completions here"}}, + headers={"x-request-id": "req_chat"}, + ) + ) + client = make_sync_client(base_url="https://example.com/openai/v1", api_key="token") + + with pytest.raises(NotFoundError, match="AWS does not support chat completions here") as exc: + client.chat.completions.create(model="gpt-4o", messages=[]) + + assert exc.value.request_id == "req_chat" + calls = cast("list[MockRequestCall]", respx_mock.calls) + assert calls[0].request.url == URL("https://example.com/openai/v1/chat/completions") + + +@pytest.mark.asyncio +@pytest.mark.respx() +async def test_passes_non_responses_resources_through_async(respx_mock: MockRouter) -> None: + respx_mock.post("https://example.com/openai/v1/chat/completions").mock( + return_value=httpx.Response( + 404, + json={"error": {"message": "AWS does not support chat completions here"}}, + headers={"x-request-id": "req_chat"}, + ) + ) + client = make_async_client(base_url="https://example.com/openai/v1", api_key="token") + + with pytest.raises(NotFoundError, match="AWS does not support chat completions here") as exc: + await client.chat.completions.create(model="gpt-4o", messages=[]) + + assert exc.value.request_id == "req_chat" + calls = cast("list[MockRequestCall]", respx_mock.calls) + assert calls[0].request.url == URL("https://example.com/openai/v1/chat/completions") + + +@pytest.mark.respx() +def test_passes_responses_features_through(respx_mock: MockRouter) -> None: + respx_mock.post("https://example.com/openai/v1/responses").mock( + return_value=httpx.Response(200, json=RESPONSE_BODY) + ) + client = make_sync_client(base_url="https://example.com/openai/v1", api_key="token") + + response = client.responses.create( + model="gpt-4o", + input="hello", + tools=[{"type": "web_search_preview"}], # type: ignore[list-item] + ) + + assert response.id == "resp_123" + calls = cast("list[MockRequestCall]", respx_mock.calls) + assert json.loads(calls[0].request.content)["tools"] == [{"type": "web_search_preview"}] + + +@pytest.mark.respx() +def test_passes_admin_security_routes_through(respx_mock: MockRouter) -> None: + respx_mock.get("https://example.com/openai/v1/organization/invites").mock( + return_value=httpx.Response( + 404, + json={"error": {"message": "AWS does not support organization invites here"}}, + headers={"x-request-id": "req_admin"}, + ) + ) + client = make_sync_client(base_url="https://example.com/openai/v1", api_key="token") + + with pytest.raises(NotFoundError, match="AWS does not support organization invites here"): + list(client.admin.organization.invites.list()) + + calls = cast("list[MockRequestCall]", respx_mock.calls) + assert calls[0].request.headers["Authorization"] == "Bearer token" + + +@pytest.mark.respx() +def test_refreshes_token_provider_for_admin_security_routes(respx_mock: MockRouter) -> None: + respx_mock.get("https://example.com/openai/v1/organization/invites").mock( + side_effect=[ + httpx.Response(500, json={"error": "server error"}), + httpx.Response( + 404, + json={"error": {"message": "AWS does not support organization invites here"}}, + headers={"x-request-id": "req_admin"}, + ), + ] + ) + tokens = iter(["first", "second"]) + client = BedrockOpenAI( + base_url="https://example.com/openai/v1", + bedrock_token_provider=lambda: next(tokens), + http_client=httpx.Client(trust_env=False), + max_retries=1, + ) + + with pytest.raises(NotFoundError, match="AWS does not support organization invites here"): + list(client.admin.organization.invites.list()) + + calls = cast("list[MockRequestCall]", respx_mock.calls) + assert calls[0].request.headers["Authorization"] == "Bearer first" + assert calls[1].request.headers["Authorization"] == "Bearer second" + + +@pytest.mark.respx() +def test_allows_responses_http_methods(respx_mock: MockRouter) -> None: + respx_mock.post("https://example.com/openai/v1/responses").mock( + return_value=httpx.Response(200, json=RESPONSE_BODY) + ) + respx_mock.get("https://example.com/openai/v1/responses/resp_123?starting_after=1&stream=true").mock( + return_value=httpx.Response(200, text=response_created_sse(), headers={"Content-Type": "text/event-stream"}) + ) + respx_mock.get("https://example.com/openai/v1/responses/resp_123?stream=true").mock( + return_value=httpx.Response(200, text=response_created_sse(), headers={"Content-Type": "text/event-stream"}) + ) + respx_mock.get("https://example.com/openai/v1/responses/resp_123").mock( + return_value=httpx.Response(200, json=RESPONSE_BODY) + ) + respx_mock.post("https://example.com/openai/v1/responses/resp_123/cancel").mock( + return_value=httpx.Response(200, json=RESPONSE_BODY) + ) + respx_mock.post("https://example.com/openai/v1/responses/compact").mock( + return_value=httpx.Response(200, json=COMPACTED_RESPONSE_BODY) + ) + respx_mock.get("https://example.com/openai/v1/responses/resp_123/input_items").mock( + return_value=httpx.Response(200, json=INPUT_ITEMS_BODY) + ) + respx_mock.post("https://example.com/openai/v1/responses/input_tokens").mock( + return_value=httpx.Response(200, json=INPUT_TOKENS_BODY) + ) + client = make_sync_client(base_url="https://example.com/openai/v1", api_key="token") + + assert client.responses.create(model="gpt-4o", input="hello", background=True).id == "resp_123" + assert client.responses.retrieve("resp_123").id == "resp_123" + assert [event.type for event in client.responses.retrieve("resp_123", starting_after=1, stream=True)] == [ + "response.created" + ] + assert [event.type for event in client.responses.retrieve("resp_123", stream=True)] == ["response.created"] + assert client.responses.cancel("resp_123").id == "resp_123" + assert client.responses.compact(model="gpt-4o").object == "response.compaction" + assert list(client.responses.input_items.list("resp_123")) == [] + assert client.responses.input_tokens.count(model="gpt-4o", input="hello").input_tokens == 1 + + calls = cast("list[MockRequestCall]", respx_mock.calls) + assert {call.request.headers["Authorization"] for call in calls} == {"Bearer token"} + + +@pytest.mark.respx() +def test_allows_sse_and_response_wrappers(respx_mock: MockRouter) -> None: + respx_mock.post("https://example.com/openai/v1/responses").mock( + side_effect=[ + httpx.Response(200, text=response_created_sse(), headers={"Content-Type": "text/event-stream"}), + httpx.Response(200, json=RESPONSE_BODY), + httpx.Response(200, json=RESPONSE_BODY), + ] + ) + client = make_sync_client(base_url="https://example.com/openai/v1", api_key="token") + + events = list(client.responses.create(model="gpt-4o", input="hello", stream=True)) + assert [event.type for event in events] == ["response.created"] + + raw_response = client.responses.with_raw_response.create(model="gpt-4o", input="hello") + assert raw_response.parse().id == "resp_123" + + with client.responses.with_streaming_response.create(model="gpt-4o", input="hello") as response: + assert response.parse().id == "resp_123" + + +def test_does_not_guard_responses_connect() -> None: + client = make_sync_client(base_url="https://example.com/openai/v1", api_key="token") + + assert client.responses.connect() is not None diff --git a/tests/test_module_client.py b/tests/test_module_client.py index 6371ae7057..7889d4b41e 100644 --- a/tests/test_module_client.py +++ b/tests/test_module_client.py @@ -30,6 +30,8 @@ def reset_state() -> None: openai.azure_endpoint = None openai.azure_ad_token = None openai.azure_ad_token_provider = None + openai._bedrock_api_key = None + openai.bedrock_token_provider = None @pytest.fixture(autouse=True) @@ -102,6 +104,7 @@ def test_http_client_option() -> None: from typing import Iterator from openai.lib.azure import AzureOpenAI +from openai.lib.bedrock import BedrockOpenAI @contextlib.contextmanager @@ -184,3 +187,64 @@ def test_azure_azure_ad_token_provider_version_and_endpoint_env() -> None: assert isinstance(client, AzureOpenAI) assert client._azure_ad_token_provider is not None assert client._azure_ad_token_provider() == "token" + + +def test_bedrock_token_and_region_env() -> None: + with fresh_env(): + openai.api_type = "amazon-bedrock" + _os.environ["AWS_BEARER_TOKEN_BEDROCK"] = "example Bedrock token" + _os.environ["AWS_REGION"] = "us-west-2" + + client = openai.responses._client + assert isinstance(client, BedrockOpenAI) + assert client.base_url == URL("https://bedrock-mantle.us-west-2.api.aws/openai/v1/") + + +def test_bedrock_api_type_env() -> None: + with fresh_env(): + _os.environ["OPENAI_API_TYPE"] = "amazon-bedrock" + _os.environ["AWS_BEARER_TOKEN_BEDROCK"] = "example Bedrock token" + _os.environ["AWS_REGION"] = "us-west-2" + reset_state() + + client = openai.responses._client + assert isinstance(client, BedrockOpenAI) + assert openai.api_type == "amazon-bedrock" + + +def test_bedrock_api_type_uses_bedrock_credentials() -> None: + with fresh_env(): + openai.api_type = "amazon-bedrock" + _os.environ["OPENAI_API_KEY"] = "openai api key" + _os.environ["AWS_BEARER_TOKEN_BEDROCK"] = "example Bedrock token" + _os.environ["AWS_REGION"] = "us-west-2" + + client = openai.responses._client + assert isinstance(client, BedrockOpenAI) + assert client.api_key == "example Bedrock token" + assert openai.api_key is None + + +def test_bedrock_api_type_uses_explicit_module_api_key() -> None: + with fresh_env(): + openai.api_type = "amazon-bedrock" + openai.api_key = "explicit Bedrock token" + _os.environ["AWS_BEARER_TOKEN_BEDROCK"] = "env Bedrock token" + _os.environ["AWS_REGION"] = "us-west-2" + + client = openai.responses._client + assert isinstance(client, BedrockOpenAI) + assert client.api_key == "explicit Bedrock token" + assert openai.api_key == "explicit Bedrock token" + + +def test_bedrock_api_type_uses_token_provider_without_mutating_module_api_key() -> None: + with fresh_env(): + openai.api_type = "amazon-bedrock" + openai.bedrock_token_provider = lambda: "provider Bedrock token" + _os.environ["AWS_REGION"] = "us-west-2" + + client = openai.responses._client + assert isinstance(client, BedrockOpenAI) + assert client._refresh_api_key() == "provider Bedrock token" + assert openai.api_key is None