diff --git a/site/src/content/docs/user-guide/sdk/model-providers/openai-responses.mdx b/site/src/content/docs/user-guide/sdk/model-providers/openai-responses.mdx index fe9766a403..c232b6a4b4 100644 --- a/site/src/content/docs/user-guide/sdk/model-providers/openai-responses.mdx +++ b/site/src/content/docs/user-guide/sdk/model-providers/openai-responses.mdx @@ -320,6 +320,54 @@ The OpenAI Responses provider does not currently implement native token counting +### Custom client + +You can pass your own OpenAI client to the OpenAIResponsesModel for Strands Agents to use +directly. The model reuses it for every request, so its connection pool and HTTP settings +carry across agent invocations. You are responsible for handling the lifecycle (e.g., +closing) of the client. + + + + +```python +import asyncio + +from openai import AsyncOpenAI +from strands import Agent +from strands.models.openai_responses import OpenAIResponsesModel + +client = AsyncOpenAI(api_key="") + +agent = Agent(model=OpenAIResponsesModel(model_id="gpt-4o", client=client)) + +async def main(): + await agent.invoke_async("What is 2+2") + await agent.invoke_async("What is 2*2") + await client.close() + +if __name__ == "__main__": + asyncio.run(main()) +``` + +`client` cannot be combined with `client_args` or `bedrock_mantle_config`. + +:::caution +Use an injected client from one event loop: call `invoke_async` or `stream_async`. +Each synchronous `agent(...)` call runs its own event loop, so the pooled connection +belongs to a closed loop and shows up as retries or connection errors. +::: + + + +```typescript +--8<-- "user-guide/sdk/model-providers/openai-responses_imports.ts:custom_client_imports" + +--8<-- "user-guide/sdk/model-providers/openai-responses.ts:custom_client" +``` + + + ## References - [Migrate from OpenAI to Strands](../../migrate/openai.md) — move a hand-written OpenAI loop or an OpenAI Agents SDK app to Strands diff --git a/site/src/content/docs/user-guide/sdk/model-providers/openai-responses.ts b/site/src/content/docs/user-guide/sdk/model-providers/openai-responses.ts index c2e4ce21fa..c71b0d34ce 100644 --- a/site/src/content/docs/user-guide/sdk/model-providers/openai-responses.ts +++ b/site/src/content/docs/user-guide/sdk/model-providers/openai-responses.ts @@ -2,6 +2,7 @@ * TypeScript examples for OpenAI Responses API documentation. */ +import OpenAI from 'openai' import { Agent } from '@strands-agents/sdk' import { OpenAIModel } from '@strands-agents/sdk/models/openai' @@ -142,3 +143,19 @@ import { OpenAIModel } from '@strands-agents/sdk/models/openai' // The model remembers "Alice" via server-side state // --8<-- [end:stateful] } + +// Custom client +{ + // --8<-- [start:custom_client] + const client = new OpenAI({ apiKey: '' }) + + const model = new OpenAIModel({ + client, + modelId: 'gpt-4o', + }) + + const agent = new Agent({ model }) + const response = await agent.invoke('What is 2+2?') + console.log(response) + // --8<-- [end:custom_client] +} diff --git a/site/src/content/docs/user-guide/sdk/model-providers/openai-responses_imports.ts b/site/src/content/docs/user-guide/sdk/model-providers/openai-responses_imports.ts index 7c9328f4a9..ccac6f9d96 100644 --- a/site/src/content/docs/user-guide/sdk/model-providers/openai-responses_imports.ts +++ b/site/src/content/docs/user-guide/sdk/model-providers/openai-responses_imports.ts @@ -5,3 +5,9 @@ import { Agent } from '@strands-agents/sdk' import { OpenAIModel } from '@strands-agents/sdk/models/openai' // --8<-- [end:basic_usage_imports] + +// --8<-- [start:custom_client_imports] +import OpenAI from 'openai' +import { Agent } from '@strands-agents/sdk' +import { OpenAIModel } from '@strands-agents/sdk/models/openai' +// --8<-- [end:custom_client_imports] diff --git a/strands-py/src/strands/models/openai_responses.py b/strands-py/src/strands/models/openai_responses.py index fc387f27e4..701715e559 100644 --- a/strands-py/src/strands/models/openai_responses.py +++ b/strands-py/src/strands/models/openai_responses.py @@ -22,7 +22,8 @@ import json import logging import mimetypes -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator +from contextlib import AsyncExitStack, asynccontextmanager from importlib.metadata import version as get_package_version from types import SimpleNamespace from typing import Any, Protocol, TypedDict, TypeVar, cast @@ -51,6 +52,7 @@ f"OpenAIResponsesModel requires openai>={_MIN_OPENAI_VERSION}. Install with: pip install -U openai" ) from e +import httpx # noqa: E402 import openai # noqa: E402 - must import after version check from ..agent.agent_metadata import AgentMetadata # noqa: E402 @@ -87,6 +89,13 @@ def __init__(self, message: str | None, code: str | None) -> None: self.code = code +async def _close_response_stream(response: Any) -> None: + """Close a streamed response, releasing its connection even when the stream was not read to the end.""" + close = getattr(response, "close", None) + if close is not None: + await close() + + def _encode_media_to_data_url(data: bytes, format_ext: str, media_type: str = "image") -> str: """Encode media bytes to a base64 data URL with size validation. @@ -167,6 +176,8 @@ def __init__( self, client_args: dict[str, Any] | None = None, bedrock_mantle_config: BedrockMantleConfig | None = None, + *, + client: Client | None = None, **model_config: Unpack[OpenAIResponsesConfig], ) -> None: """Initialize provider instance. @@ -180,25 +191,62 @@ def __init__( OpenAI-compatible endpoints, ``bedrock-mantle`` (the default) or ``bedrock-runtime`` via the config's ``endpoint`` key. See :class:`BedrockMantleConfig` for accepted keys. When set, a fresh bearer - token is minted on every request. + token is minted on every request. Cannot be combined with a pre-built + ``client``. + client: Pre-configured OpenAI-compatible client to reuse across requests. + When provided, this client will be reused for all requests and will NOT be closed + by the model. The caller is responsible for managing the client lifecycle. + This is useful for: + - Injecting custom client wrappers (e.g., GuardrailsAsyncOpenAI) + - Reusing connection pools within a single event loop/worker + - Centralizing observability, retries, and networking policy + - Pointing to custom model gateways + Note: Use the client from a single asyncio event loop. Avoid synchronous + ``agent(...)`` calls, which each run on their own event loop; call + ``invoke_async`` or ``stream_async`` from one loop instead. **model_config: Configuration options for the OpenAI Responses API model. + + Raises: + ValueError: If ``client`` is combined with ``client_args`` or ``bedrock_mantle_config``, + or if ``client_args`` sets ``api_key`` or ``base_url`` alongside ``bedrock_mantle_config``. """ validate_config_keys(model_config, self.OpenAIResponsesConfig) self.config = dict(model_config) + self._validate_client_options(client, client_args, bedrock_mantle_config) + + self._custom_client = client self.client_args = client_args or {} self._bedrock_mantle_config = bedrock_mantle_config - if bedrock_mantle_config is not None and client_args: - conflicting = [k for k in ("api_key", "base_url") if k in client_args] - if conflicting: - raise ValueError( - f"client_args must not contain {conflicting} when bedrock_mantle_config is set; " - "these are derived from the Mantle config automatically." - ) - logger.debug("config=<%s> | initializing", self.config) + @staticmethod + def _validate_client_options( + client: Client | None, + client_args: dict[str, Any] | None, + bedrock_mantle_config: BedrockMantleConfig | None, + ) -> None: + """Reject client options that conflict with each other. + + Raises: + ValueError: If ``client`` is combined with ``client_args`` or ``bedrock_mantle_config``, + or if ``client_args`` sets ``api_key`` or ``base_url`` alongside ``bedrock_mantle_config``. + """ + if client is not None and client_args: + raise ValueError("Only one of 'client' or 'client_args' should be provided, not both.") + if client is not None and bedrock_mantle_config is not None: + raise ValueError("'bedrock_mantle_config' cannot be combined with a pre-built 'client'.") + if bedrock_mantle_config is None or not client_args: + return + + conflicting = [key for key in ("api_key", "base_url") if key in client_args] + if conflicting: + raise ValueError( + f"client_args must not contain {conflicting} when bedrock_mantle_config is set; " + "these are derived from the Mantle config automatically." + ) + def _resolve_client_args(self) -> dict[str, Any]: """Return the kwargs to pass to ``openai.AsyncOpenAI`` for the current request. @@ -210,6 +258,45 @@ def _resolve_client_args(self) -> dict[str, Any]: ) return self.client_args + async def _drain_for_connection_reuse(self, response: AsyncIterable[Any]) -> None: + """Read the rest of a completed streamed response so its connection can return to the client's pool. + + Only a pooled connection needs the rest of the body read; a per-request client is closed right after. + Best effort: the response is already complete, so a failed read only costs the connection. + """ + if self._custom_client is None: + return + + try: + async for _event in response: + pass + except (openai.APIError, httpx.HTTPError, httpx.StreamError, ValueError) as error: + logger.debug("error=<%s> | failed to read the rest of the response stream after the terminal event", error) + + @asynccontextmanager + async def _get_client(self) -> AsyncIterator[Any]: + """Get an OpenAI client for making requests. + + This context manager handles client lifecycle management: + - If an injected client was provided during initialization, it yields that client + without closing it (caller manages lifecycle). + - Otherwise, creates a new AsyncOpenAI client from client_args and automatically + closes it when the context exits. + + Note: We create a new client per request to avoid connection sharing in the underlying + httpx client, as the asyncio event loop does not allow connections to be shared. + For more details, see https://github.com/encode/httpx/discussions/2959. + + Yields: + Client: An OpenAI-compatible client instance. + """ + if self._custom_client is not None: + yield self._custom_client + return + + async with openai.AsyncOpenAI(**self._resolve_client_args()) as client: + yield client + @property @override def stateful(self) -> bool: @@ -275,7 +362,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._resolve_client_args()) as client: + async with self._get_client() as client: response = await client.responses.input_tokens.count(**request) total_tokens: int = response.input_tokens @@ -329,9 +416,10 @@ async def stream( logger.debug("invoking OpenAI Responses API model") - async with openai.AsyncOpenAI(**self._resolve_client_args()) as client: + async with self._get_client() as client, AsyncExitStack() as cleanup: try: response = await client.responses.create(**request) + cleanup.push_async_callback(_close_response_stream, response) logger.debug("streaming response from OpenAI Responses API model") @@ -458,6 +546,10 @@ async def stream( if hasattr(event, "response") and hasattr(event.response, "usage"): final_usage = event.response.usage break + + # The loop stops at the terminal event; the rest of the stream must be read before the + # connection can be reused. + await self._drain_for_connection_reuse(response) except (openai.APIError, _OpenAIResponsesStreamError) as error: error_kind = classify_openai_error(error) if error_kind == "throttling": @@ -523,7 +615,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._resolve_client_args()) as client: + async with self._get_client() as client: try: request = self._format_request(prompt, system_prompt=system_prompt) request.pop("stream", None) diff --git a/strands-py/tests/strands/models/test_openai_responses.py b/strands-py/tests/strands/models/test_openai_responses.py index 0b1108e88c..28f4195286 100644 --- a/strands-py/tests/strands/models/test_openai_responses.py +++ b/strands-py/tests/strands/models/test_openai_responses.py @@ -18,11 +18,16 @@ @pytest.fixture -def openai_client(): +def openai_client_cls(): with unittest.mock.patch.object(strands.models.openai_responses.openai, "AsyncOpenAI") as mock_client_cls: - mock_client = unittest.mock.AsyncMock() - mock_client_cls.return_value.__aenter__.return_value = mock_client - yield mock_client + yield mock_client_cls + + +@pytest.fixture +def openai_client(openai_client_cls): + mock_client = unittest.mock.AsyncMock() + openai_client_cls.return_value.__aenter__.return_value = mock_client + return mock_client @pytest.fixture @@ -1235,6 +1240,189 @@ async def test_structured_output_forwards_request_params(openai_client, model_id assert "stream" not in parse_kwargs +@pytest.fixture +def injected_client(): + return unittest.mock.AsyncMock() + + +@pytest.mark.asyncio +async def test_stream_creates_and_closes_client_per_request( + openai_client, openai_client_cls, model_id, messages, agenerator, alist +): + """Guards the default path: without an injected client, each request opens a fresh client and closes it.""" + mock_client = openai_client + mock_client.responses.create = unittest.mock.AsyncMock(side_effect=lambda **_: agenerator([])) + model = OpenAIResponsesModel(model_id=model_id, client_args={"api_key": "k1"}) + + await alist(model.stream(messages)) + await alist(model.stream(messages)) + + assert openai_client_cls.call_args_list == [unittest.mock.call(api_key="k1")] * 2 + assert openai_client_cls.return_value.__aexit__.await_count == 2 + + +@pytest.mark.asyncio +async def test_stream_with_injected_client(openai_client_cls, injected_client, model_id, messages, agenerator, alist): + injected_client.responses.create = unittest.mock.AsyncMock(side_effect=lambda **_: agenerator([])) + model = OpenAIResponsesModel(client=injected_client, model_id=model_id) + + await alist(model.stream(messages)) + await alist(model.stream(messages)) + + assert injected_client.responses.create.await_count == 2 + openai_client_cls.assert_not_called() + injected_client.__aexit__.assert_not_called() + injected_client.close.assert_not_called() + + +@pytest.mark.asyncio +async def test_structured_output_with_injected_client( + openai_client_cls, injected_client, model_id, messages, test_output_model_cls, alist +): + mock_parsed_instance = test_output_model_cls(name="John", age=30) + injected_client.responses.parse = unittest.mock.AsyncMock( + return_value=unittest.mock.Mock(output_parsed=mock_parsed_instance) + ) + model = OpenAIResponsesModel(client=injected_client, model_id=model_id) + + tru_events = await alist(model.structured_output(test_output_model_cls, messages)) + exp_events = [{"output": mock_parsed_instance}] + + assert tru_events == exp_events + injected_client.responses.parse.assert_awaited_once() + openai_client_cls.assert_not_called() + injected_client.__aexit__.assert_not_called() + injected_client.close.assert_not_called() + + +@pytest.mark.asyncio +async def test_count_tokens_with_injected_client(openai_client_cls, injected_client, model_id, messages): + injected_client.responses.input_tokens.count = unittest.mock.AsyncMock( + return_value=unittest.mock.Mock(input_tokens=42) + ) + model = OpenAIResponsesModel(client=injected_client, model_id=model_id, use_native_token_count=True) + + tru_tokens = await model.count_tokens(messages=messages) + + assert tru_tokens == 42 + openai_client_cls.assert_not_called() + injected_client.__aexit__.assert_not_called() + injected_client.close.assert_not_called() + + +def test__init__with_both_client_and_client_args_raises_error(injected_client): + with pytest.raises(ValueError, match="Only one of 'client' or 'client_args' should be provided"): + OpenAIResponsesModel(client=injected_client, client_args={"api_key": "test"}, model_id="test-model") + + +class _ClosableStream: + """Response stream stand-in that records how far it was read and whether it was closed.""" + + def __init__(self, events): + self.remaining = list(events) + self.closed = False + + def __aiter__(self): + return self + + async def __anext__(self): + if not self.remaining: + raise StopAsyncIteration + event = self.remaining.pop(0) + if isinstance(event, Exception): + raise event + return event + + async def close(self): + self.closed = True + + +@pytest.mark.asyncio +@pytest.mark.parametrize("inject_client", [True, False]) +async def test_stream_reads_past_terminal_event_only_with_injected_client( + openai_client, injected_client, model_id, messages, alist, inject_client +): + completed_event = unittest.mock.Mock(type="response.completed", response=unittest.mock.Mock(usage=None)) + trailing_event = unittest.mock.Mock(type="response.output_text.delta", delta="ignored") + response_stream = _ClosableStream([completed_event, trailing_event]) + request_client = injected_client if inject_client else openai_client + request_client.responses.create = unittest.mock.AsyncMock(return_value=response_stream) + model_kwargs = {"client": injected_client} if inject_client else {} + model = OpenAIResponsesModel(model_id=model_id, **model_kwargs) + + tru_events = await alist(model.stream(messages)) + exp_events = [{"messageStart": {"role": "assistant"}}, {"messageStop": {"stopReason": "end_turn"}}] + + assert tru_events == exp_events + assert response_stream.remaining == ([] if inject_client else [trailing_event]) + assert response_stream.closed + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "read_error", + [ + httpx.RemoteProtocolError("peer closed connection without sending complete message body"), + openai.APIError("An error occurred during streaming", request=httpx.Request("POST", "https://x"), body=None), + ], +) +async def test_stream_completes_when_reading_after_terminal_event_fails( + injected_client, model_id, messages, alist, read_error +): + text_event = unittest.mock.Mock(type="response.output_text.delta", delta="Hello") + completed_event = unittest.mock.Mock(type="response.completed", response=unittest.mock.Mock(usage=None)) + response_stream = _ClosableStream([text_event, completed_event, read_error]) + injected_client.responses.create = unittest.mock.AsyncMock(return_value=response_stream) + model = OpenAIResponsesModel(client=injected_client, model_id=model_id) + + tru_events = await alist(model.stream(messages)) + exp_events = [ + {"messageStart": {"role": "assistant"}}, + {"contentBlockStart": {"start": {}}}, + {"contentBlockDelta": {"delta": {"text": "Hello"}}}, + {"contentBlockStop": {}}, + {"messageStop": {"stopReason": "end_turn"}}, + ] + + assert tru_events == exp_events + assert response_stream.closed + + +@pytest.mark.asyncio +async def test_stream_closes_response_without_reading_it_when_consumer_stops_early(injected_client, model_id, messages): + text_event = unittest.mock.Mock(type="response.output_text.delta", delta="Hello") + completed_event = unittest.mock.Mock(type="response.completed", response=unittest.mock.Mock(usage=None)) + response_stream = _ClosableStream([text_event, completed_event]) + injected_client.responses.create = unittest.mock.AsyncMock(return_value=response_stream) + model = OpenAIResponsesModel(client=injected_client, model_id=model_id) + + stream = model.stream(messages) + await anext(stream) + await stream.aclose() + + assert response_stream.remaining == [text_event, completed_event] + assert response_stream.closed + + +@pytest.mark.asyncio +async def test_stream_closes_response_without_reading_it_on_error(injected_client, model_id, messages): + error = ResponseError(message="The model failed while processing the request.", code="server_error") + failed_event = ResponseFailedEvent.model_construct( + type="response.failed", sequence_number=1, response=Response.model_construct(error=error) + ) + trailing_event = unittest.mock.Mock(type="response.output_text.delta", delta="ignored") + response_stream = _ClosableStream([failed_event, trailing_event]) + injected_client.responses.create = unittest.mock.AsyncMock(return_value=response_stream) + model = OpenAIResponsesModel(client=injected_client, model_id=model_id) + + with pytest.raises(RuntimeError, match="The model failed while processing the request"): + async for _ in model.stream(messages): + pass + + assert response_stream.remaining == [trailing_event] + assert response_stream.closed + + @pytest.mark.asyncio async def test_stream_response_failed_raises_provider_error(openai_client, model, messages, agenerator): error = ResponseError(message="The model failed while processing the request.", code="server_error") @@ -2103,6 +2291,15 @@ def test_bedrock_mantle_config_merges_with_client_args(self, openai_client, mock assert resolved["timeout"] == 42 assert resolved["http_client"] is sentinel_http_client + def test_bedrock_mantle_config_conflicts_with_custom_client(self, injected_client): + """Cannot pass both bedrock_mantle_config and a pre-built client.""" + with pytest.raises(ValueError, match="bedrock_mantle_config"): + OpenAIResponsesModel( + model_id="openai.gpt-oss-120b", + client=injected_client, + bedrock_mantle_config={"region": "us-east-1"}, + ) + def test_bedrock_mantle_config_rejects_base_url_in_client_args(self, openai_client): """client_args must not contain base_url or api_key when bedrock_mantle_config is set.""" _ = openai_client