feat: enable openai provider use aws profile

Author: JackYPCOnline

pyproject.toml

@@ -50,7 +50,7 @@ litellm = ["litellm>=1.75.9,<=1.83.13", "openai>=1.68.0,<3.0.0"]
 llamaapi = ["llama-api-client>=0.1.0,<1.0.0"]
 mistral = ["mistralai>=1.8.2,<2.0.0"]
 ollama = ["ollama>=0.4.8,<1.0.0"]
-openai = ["openai>=1.68.0,<3.0.0"]
+openai = ["openai>=1.68.0,<3.0.0", "aws-bedrock-token-generator>=1.1.0,<2.0.0"]
 writer = ["writer-sdk>=2.2.0,<3.0.0"]
 sagemaker = [
     "boto3-stubs[sagemaker-runtime]>=1.26.0,<2.0.0",

src/strands/models/openai.py

@@ -21,6 +21,7 @@
 from ..types.exceptions import ContextWindowOverflowException, ModelThrottledException
 from ..types.streaming import StreamEvent
 from ..types.tools import ToolChoice, ToolResult, ToolSpec, ToolUse
+from ._openai_bedrock import AwsConfig, resolve_bedrock_client_args
 from ._validation import _has_location_source, validate_config_keys
 from .model import BaseModelConfig, Model
 
@@ -71,6 +72,7 @@ def __init__(
         self,
         client: Client | None = None,
         client_args: dict[str, Any] | None = None,
+        aws_config: AwsConfig | None = None,
         **model_config: Unpack[OpenAIConfig],
     ) -> None:
         """Initialize provider instance.
@@ -87,23 +89,53 @@ def __init__(
                 Note: The client should not be shared across different asyncio event loops.
             client_args: Arguments for the OpenAI client (legacy approach).
                 For a complete list of supported arguments, see https://pypi.org/project/openai/.
+                May be combined with ``aws_config``; transport-level options like ``http_client``,
+                ``timeout``, or ``default_headers`` are preserved, while ``base_url`` and
+                ``api_key`` are always overridden by ``aws_config`` when both are set.
+            aws_config: Route requests through Amazon Bedrock's Mantle (OpenAI-compatible)
+                endpoint. Provide ``{"region": "us-east-1"}`` at minimum. Accepts optional
+                ``credentials_provider`` (a botocore ``CredentialProvider``) and ``expiry``
+                (a ``datetime.timedelta`` up to 12h). When set, a fresh bearer token is minted
+                on every request via ``aws-bedrock-token-generator`` and the OpenAI client is
+                pointed at ``https://bedrock-mantle.<region>.api.aws/v1``. Cannot be combined
+                with a pre-built ``client``.
             **model_config: Configuration options for the OpenAI model.
 
         Raises:
-            ValueError: If both `client` and `client_args` are provided.
+            ValueError: If ``client`` is combined with ``client_args`` or ``aws_config``,
+                or if ``aws_config`` is missing a region.
         """
         validate_config_keys(model_config, self.OpenAIConfig)
         self.config = dict(model_config)
 
-        # Validate that only one client configuration method is provided
-        if client is not None and client_args is not None and len(client_args) > 0:
+        # Validate that client configuration methods are mutually exclusive where they conflict.
+        # client_args + aws_config is allowed — aws_config will override base_url / api_key only.
+        client_args_provided = client_args is not None and len(client_args) > 0
+        if client is not None and client_args_provided:
             raise ValueError("Only one of 'client' or 'client_args' should be provided, not both.")
+        if aws_config is not None:
+            if client is not None:
+                raise ValueError("'aws_config' cannot be combined with a pre-built 'client'.")
+            if not aws_config.get("region"):
+                raise ValueError("aws_config must include a non-empty 'region'.")
 
         self._custom_client = client
         self.client_args = client_args or {}
+        self._aws_config = aws_config
 
         logger.debug("config=<%s> | initializing", self.config)
 
+    def _resolve_client_args(self) -> dict[str, Any]:
+        """Return the kwargs to pass to ``openai.AsyncOpenAI`` for the current request.
+
+        When ``aws_config`` is set, a fresh Bedrock Mantle bearer token is minted on every
+        call and ``base_url`` / ``api_key`` are overridden. Any other entries from
+        ``client_args`` (e.g. ``http_client``, ``timeout``) are preserved.
+        """
+        if self._aws_config is not None:
+            return resolve_bedrock_client_args(self._aws_config, self.client_args)
+        return self.client_args
+
     @override
     def update_config(self, **model_config: Unpack[OpenAIConfig]) -> None:  # type: ignore[override]
         """Update the OpenAI model configuration with the provided arguments.
@@ -590,11 +622,11 @@ async def _get_client(self) -> AsyncIterator[Any]:
             # Use the injected client (caller manages lifecycle)
             yield self._custom_client
         else:
-            # Create a new client from client_args
+            # Create a new client from resolved args (static client_args or freshly-minted Bedrock creds).
             # We initialize an OpenAI context on every request so as to avoid connection sharing in the underlying
             # httpx client. The asyncio event loop does not allow connections to be shared. For more details, please
             # refer to https://github.com/encode/httpx/discussions/2959.
-            async with openai.AsyncOpenAI(**self.client_args) as client:
+            async with openai.AsyncOpenAI(**self._resolve_client_args()) as client:
                 yield client
 
     @override

src/strands/models/openai_responses.py

@@ -58,6 +58,7 @@
 from ..types.exceptions import ContextWindowOverflowException, ModelThrottledException  # noqa: E402
 from ..types.streaming import StreamEvent  # noqa: E402
 from ..types.tools import ToolChoice, ToolResult, ToolSpec, ToolUse  # noqa: E402
+from ._openai_bedrock import AwsConfig, resolve_bedrock_client_args  # noqa: E402
 from ._validation import validate_config_keys  # noqa: E402
 from .model import BaseModelConfig, Model  # noqa: E402
 
@@ -141,21 +142,52 @@ class OpenAIResponsesConfig(BaseModelConfig, total=False):
         stateful: bool
 
     def __init__(
-        self, client_args: dict[str, Any] | None = None, **model_config: Unpack[OpenAIResponsesConfig]
+        self,
+        client_args: dict[str, Any] | None = None,
+        aws_config: AwsConfig | None = None,
+        **model_config: Unpack[OpenAIResponsesConfig],
     ) -> None:
         """Initialize provider instance.
 
         Args:
             client_args: Arguments for the OpenAI client.
                 For a complete list of supported arguments, see https://pypi.org/project/openai/.
+                May be combined with ``aws_config``; transport-level options like ``http_client``,
+                ``timeout``, or ``default_headers`` are preserved, while ``base_url`` and
+                ``api_key`` are always overridden by ``aws_config`` when both are set.
+            aws_config: Route requests through Amazon Bedrock's Mantle (OpenAI-compatible)
+                endpoint. Provide ``{"region": "us-east-1"}`` at minimum. Accepts optional
+                ``credentials_provider`` (a botocore ``CredentialProvider``) and ``expiry``
+                (a ``datetime.timedelta`` up to 12h). When set, a fresh bearer token is minted
+                on every request via ``aws-bedrock-token-generator`` and the OpenAI client is
+                pointed at ``https://bedrock-mantle.<region>.api.aws/v1``.
             **model_config: Configuration options for the OpenAI Responses API model.
+
+        Raises:
+            ValueError: If ``aws_config`` is missing a region.
         """
         validate_config_keys(model_config, self.OpenAIResponsesConfig)
         self.config = dict(model_config)
+
+        if aws_config is not None and not aws_config.get("region"):
+            raise ValueError("aws_config must include a non-empty 'region'.")
+
         self.client_args = client_args or {}
+        self._aws_config = aws_config
 
         logger.debug("config=<%s> | initializing", self.config)
 
+    def _resolve_client_args(self) -> dict[str, Any]:
+        """Return the kwargs to pass to ``openai.AsyncOpenAI`` for the current request.
+
+        When ``aws_config`` is set, a fresh Bedrock Mantle bearer token is minted on every
+        call and ``base_url`` / ``api_key`` are overridden. Any other entries from
+        ``client_args`` (e.g. ``http_client``, ``timeout``) are preserved.
+        """
+        if self._aws_config is not None:
+            return resolve_bedrock_client_args(self._aws_config, self.client_args)
+        return self.client_args
+
     @property
     @override
     def stateful(self) -> bool:
@@ -215,7 +247,7 @@ async def count_tokens(
             count_tokens_fields = {"model", "input", "instructions", "tools"}
             request = {k: request[k] for k in request.keys() & count_tokens_fields}
 
-            async with openai.AsyncOpenAI(**self.client_args) as client:
+            async with openai.AsyncOpenAI(**self._resolve_client_args()) as client:
                 response = await client.responses.input_tokens.count(**request)
                 total_tokens: int = response.input_tokens
 
@@ -267,7 +299,7 @@ async def stream(
 
         logger.debug("invoking OpenAI Responses API model")
 
-        async with openai.AsyncOpenAI(**self.client_args) as client:
+        async with openai.AsyncOpenAI(**self._resolve_client_args()) as client:
             try:
                 response = await client.responses.create(**request)
 
@@ -447,7 +479,7 @@ async def structured_output(
             ContextWindowOverflowException: If the input exceeds the model's context window.
             ModelThrottledException: If the request is throttled by OpenAI (rate limits).
         """
-        async with openai.AsyncOpenAI(**self.client_args) as client:
+        async with openai.AsyncOpenAI(**self._resolve_client_args()) as client:
             try:
                 response = await client.responses.parse(
                     model=self.get_config()["model_id"],

tests/strands/models/test_openai.py

@@ -1710,3 +1710,102 @@ def test_format_request_messages_multiple_tool_calls_with_images():
         },
     ]
     assert tru_result == exp_result
+
+
+# =============================================================================
+# Bedrock Mantle (aws_config) integration with OpenAIModel
+# =============================================================================
+
+
+class TestOpenAIModelAwsConfig:
+    """Tests for the Bedrock Mantle pathway via the aws_config kwarg."""
+
+    @pytest.fixture
+    def mock_provide_token(self):
+        with unittest.mock.patch("strands.models._openai_bedrock.provide_token") as mock:
+            mock.return_value = "bedrock-api-key-deadbeef&Version=1"
+            yield mock
+
+    def test_aws_config_sets_base_url_and_api_key(self, openai_client, mock_provide_token):
+        """aws_config produces the Mantle base_url and a minted bearer token as api_key."""
+        _ = openai_client
+        model = OpenAIModel(model_id="openai.gpt-oss-120b", aws_config={"region": "us-east-1"})
+
+        # api_key is resolved per-request (lazy), so check via the resolved client_args at call time
+        resolved = model._resolve_client_args()
+        assert resolved["base_url"] == "https://bedrock-mantle.us-east-1.api.aws/v1"
+        assert resolved["api_key"] == "bedrock-api-key-deadbeef&Version=1"
+        # Only region is forwarded when the user did not set optional kwargs,
+        # so provide_token's own defaults (e.g. 12h expiry) apply.
+        mock_provide_token.assert_called_once_with(region="us-east-1")
+
+    def test_aws_config_forwards_credentials_provider_and_expiry(self, openai_client, mock_provide_token):
+        """Optional credentials_provider and expiry are forwarded to provide_token."""
+        _ = openai_client
+        from datetime import timedelta
+
+        provider = unittest.mock.Mock()
+        model = OpenAIModel(
+            model_id="openai.gpt-oss-120b",
+            aws_config={
+                "region": "us-west-2",
+                "credentials_provider": provider,
+                "expiry": timedelta(minutes=15),
+            },
+        )
+        model._resolve_client_args()
+        mock_provide_token.assert_called_once_with(
+            region="us-west-2",
+            aws_credentials_provider=provider,
+            expiry=timedelta(minutes=15),
+        )
+
+    def test_aws_config_mints_token_per_request(self, openai_client, mock_provide_token):
+        """Each call to _resolve_client_args mints a fresh token (long-lived processes)."""
+        _ = openai_client
+        model = OpenAIModel(model_id="openai.gpt-oss-120b", aws_config={"region": "us-east-1"})
+        model._resolve_client_args()
+        model._resolve_client_args()
+        model._resolve_client_args()
+        assert mock_provide_token.call_count == 3
+
+    def test_aws_config_conflicts_with_custom_client(self, openai_client):
+        """Cannot pass both aws_config and a pre-built client."""
+        _ = openai_client
+        custom_client = unittest.mock.Mock()
+        with pytest.raises(ValueError, match="aws_config"):
+            OpenAIModel(
+                model_id="openai.gpt-oss-120b",
+                client=custom_client,
+                aws_config={"region": "us-east-1"},
+            )
+
+    def test_aws_config_merges_with_client_args(self, openai_client, mock_provide_token):
+        """aws_config is allowed alongside client_args; base_url and api_key are overridden,
+        other transport-level options (timeout, http_client, default_headers) are preserved.
+        """
+        _ = openai_client
+        sentinel_http_client = unittest.mock.Mock()
+        model = OpenAIModel(
+            model_id="openai.gpt-oss-120b",
+            client_args={
+                "api_key": "will-be-overridden",
+                "base_url": "https://also-overridden.example.com",
+                "timeout": 42,
+                "http_client": sentinel_http_client,
+                "default_headers": {"X-Trace-Id": "abc"},
+            },
+            aws_config={"region": "us-east-1"},
+        )
+        resolved = model._resolve_client_args()
+        assert resolved["base_url"] == "https://bedrock-mantle.us-east-1.api.aws/v1"
+        assert resolved["api_key"] == "bedrock-api-key-deadbeef&Version=1"
+        assert resolved["timeout"] == 42
+        assert resolved["http_client"] is sentinel_http_client
+        assert resolved["default_headers"] == {"X-Trace-Id": "abc"}
+
+    def test_aws_config_requires_region(self, openai_client):
+        """aws_config must include a region."""
+        _ = openai_client
+        with pytest.raises(ValueError, match="region"):
+            OpenAIModel(model_id="openai.gpt-oss-120b", aws_config={})

tests/strands/models/test_openai_responses.py

@@ -1298,3 +1298,80 @@ async def test_fallback_logs_debug(self, model, openai_client, messages, caplog)
             await model.count_tokens(messages=messages)
 
         assert any("native token counting failed" in record.message for record in caplog.records)
+
+
+# =============================================================================
+# Bedrock Mantle (aws_config) integration with OpenAIResponsesModel
+# =============================================================================
+
+
+class TestOpenAIResponsesModelAwsConfig:
+    """Tests for the Bedrock Mantle pathway via the aws_config kwarg."""
+
+    @pytest.fixture
+    def mock_provide_token(self):
+        with unittest.mock.patch("strands.models._openai_bedrock.provide_token") as mock:
+            mock.return_value = "bedrock-api-key-deadbeef&Version=1"
+            yield mock
+
+    def test_aws_config_sets_base_url_and_api_key(self, openai_client, mock_provide_token):
+        _ = openai_client
+        model = OpenAIResponsesModel(model_id="openai.gpt-oss-120b", aws_config={"region": "us-east-1"})
+        resolved = model._resolve_client_args()
+        assert resolved["base_url"] == "https://bedrock-mantle.us-east-1.api.aws/v1"
+        assert resolved["api_key"] == "bedrock-api-key-deadbeef&Version=1"
+        mock_provide_token.assert_called_once_with(region="us-east-1")
+
+    def test_aws_config_forwards_credentials_provider_and_expiry(self, openai_client, mock_provide_token):
+        _ = openai_client
+        from datetime import timedelta
+
+        provider = unittest.mock.Mock()
+        model = OpenAIResponsesModel(
+            model_id="openai.gpt-oss-120b",
+            aws_config={
+                "region": "us-west-2",
+                "credentials_provider": provider,
+                "expiry": timedelta(minutes=15),
+            },
+        )
+        model._resolve_client_args()
+        mock_provide_token.assert_called_once_with(
+            region="us-west-2",
+            aws_credentials_provider=provider,
+            expiry=timedelta(minutes=15),
+        )
+
+    def test_aws_config_mints_token_per_request(self, openai_client, mock_provide_token):
+        _ = openai_client
+        model = OpenAIResponsesModel(model_id="openai.gpt-oss-120b", aws_config={"region": "us-east-1"})
+        model._resolve_client_args()
+        model._resolve_client_args()
+        assert mock_provide_token.call_count == 2
+
+    def test_aws_config_merges_with_client_args(self, openai_client, mock_provide_token):
+        """aws_config is allowed alongside client_args; base_url and api_key are overridden,
+        other transport-level options are preserved.
+        """
+        _ = openai_client
+        sentinel_http_client = unittest.mock.Mock()
+        model = OpenAIResponsesModel(
+            model_id="openai.gpt-oss-120b",
+            client_args={
+                "api_key": "will-be-overridden",
+                "base_url": "https://also-overridden.example.com",
+                "timeout": 42,
+                "http_client": sentinel_http_client,
+            },
+            aws_config={"region": "us-east-1"},
+        )
+        resolved = model._resolve_client_args()
+        assert resolved["base_url"] == "https://bedrock-mantle.us-east-1.api.aws/v1"
+        assert resolved["api_key"] == "bedrock-api-key-deadbeef&Version=1"
+        assert resolved["timeout"] == 42
+        assert resolved["http_client"] is sentinel_http_client
+
+    def test_aws_config_requires_region(self, openai_client):
+        _ = openai_client
+        with pytest.raises(ValueError, match="region"):
+            OpenAIResponsesModel(model_id="openai.gpt-oss-120b", aws_config={})