From e0ebbf0bc69dc78122a268a1c6b98f62e38fd490 Mon Sep 17 00:00:00 2001 From: Serene Aryal Date: Wed, 7 Oct 2026 02:47:03 -0500 Subject: [PATCH] feat(models): accept a pre-configured client in OpenAIResponsesModel Mirrors OpenAIModel's client parameter: an injected client is reused for every request and never closed by the model. With an injected client, stream() reads the rest of the body after the terminal event so the connection returns to the pool, and always closes the response. Resolves #4544 --- .../sdk/model-providers/openai-responses.mdx | 48 ++++ .../sdk/model-providers/openai-responses.ts | 17 ++ .../openai-responses_imports.ts | 6 + .../src/strands/models/openai_responses.py | 118 ++++++++-- .../strands/models/test_openai_responses.py | 205 +++++++++++++++++- 5 files changed, 377 insertions(+), 17 deletions(-) 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