From d202e1a963e94424844de10758d077a6a20a9169 Mon Sep 17 00:00:00 2001 From: Curtis Galione Date: Mon, 20 Jul 2026 23:06:30 -0700 Subject: [PATCH 1/3] fix(openai): preserve async streaming response parse --- .../integrations/openai/test_openai.py | 30 +++++++++++++++++++ .../braintrust/integrations/openai/tracing.py | 18 +++++++++-- 2 files changed, 46 insertions(+), 2 deletions(-) diff --git a/py/src/braintrust/integrations/openai/test_openai.py b/py/src/braintrust/integrations/openai/test_openai.py index b95ce084..66027dea 100644 --- a/py/src/braintrust/integrations/openai/test_openai.py +++ b/py/src/braintrust/integrations/openai/test_openai.py @@ -7,6 +7,7 @@ import time import zlib +import httpx import openai import pytest from braintrust import Attachment, logger, wrap_openai @@ -2782,6 +2783,35 @@ async def test_setup_async_creates_spans(self, memory_logger): assert "gpt-4o-mini" in span["metadata"]["model"] assert span["input"] + @pytest.mark.asyncio + async def test_setup_streaming_response_parse_remains_awaitable(self, memory_logger): + """OpenAI Agents awaits parse() on the streaming response wrapper.""" + assert not memory_logger.pop() + + async def handle_request(request): + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content=b"data: [DONE]\n\n", + request=request, + ) + + OpenAIIntegration.setup() + transport = httpx.MockTransport(handle_request) + async with httpx.AsyncClient(transport=transport) as http_client: + client = openai.AsyncOpenAI(api_key="test-api-key", http_client=http_client) + async with client.responses.with_streaming_response.create( + model=TEST_MODEL, + input=TEST_PROMPT, + stream=True, + ) as raw_response: + stream = await raw_response.parse() + assert hasattr(stream, "__aiter__") + await stream.close() + + spans = memory_logger.pop() + assert len(spans) == 1 + class TestAutoInstrumentOpenAI: """Tests for auto_instrument() with OpenAI.""" diff --git a/py/src/braintrust/integrations/openai/tracing.py b/py/src/braintrust/integrations/openai/tracing.py index 6348f481..aa0e3b0c 100644 --- a/py/src/braintrust/integrations/openai/tracing.py +++ b/py/src/braintrust/integrations/openai/tracing.py @@ -844,11 +844,17 @@ class _RawResponseWithTracedStream(NamedWrapper): """Proxy for LegacyAPIResponse that replaces parse() with a traced stream, so that with_raw_response + stream=True preserves both headers and tracing.""" - def __init__(self, raw_response: Any, traced_stream: Any) -> None: + def __init__(self, raw_response: Any, traced_stream: Any, *, async_parse: bool = False) -> None: self._traced_stream = traced_stream + self._async_parse = async_parse super().__init__(raw_response) def parse(self, *args: Any, **kwargs: Any) -> Any: + if self._async_parse: + return self._aparse() + return self._traced_stream + + async def _aparse(self) -> Any: return self._traced_stream @@ -1066,8 +1072,12 @@ async def acreate(self, *args: Any, **kwargs: Any) -> Any: try: start = time.time() create_response = await self.acreate_fn(*args, **kwargs) + parse_was_awaitable = False if hasattr(create_response, "parse"): raw_response = create_response.parse() + if inspect.isawaitable(raw_response): + parse_was_awaitable = True + raw_response = await raw_response log_headers(create_response, span) else: raw_response = create_response @@ -1097,7 +1107,11 @@ async def gen(): should_end = False streamer = gen() if raw_requested and hasattr(create_response, "parse"): - return _RawResponseWithTracedStream(create_response, _AsyncTracedStream(raw_response, streamer)) + return _RawResponseWithTracedStream( + create_response, + _AsyncTracedStream(raw_response, streamer), + async_parse=parse_was_awaitable, + ) return _AsyncTracedStream(raw_response, streamer) else: log_response = _try_to_dict(raw_response) From dab46ef838ab6cfb0493e2be1b54e431f8d43f0b Mon Sep 17 00:00:00 2001 From: Curtis Galione Date: Tue, 21 Jul 2026 00:08:55 -0700 Subject: [PATCH 2/3] test(openai): exercise parsed async stream --- py/src/braintrust/integrations/openai/test_openai.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/py/src/braintrust/integrations/openai/test_openai.py b/py/src/braintrust/integrations/openai/test_openai.py index 66027dea..7e2092bc 100644 --- a/py/src/braintrust/integrations/openai/test_openai.py +++ b/py/src/braintrust/integrations/openai/test_openai.py @@ -2805,8 +2805,11 @@ async def handle_request(request): input=TEST_PROMPT, stream=True, ) as raw_response: - stream = await raw_response.parse() + parse_result = raw_response.parse() + assert inspect.isawaitable(parse_result) + stream = await parse_result assert hasattr(stream, "__aiter__") + assert [event async for event in stream] == [] await stream.close() spans = memory_logger.pop() From 6434e18cbfee954025a3015bda4e5c5e40750445 Mon Sep 17 00:00:00 2001 From: Abhijeet Prasad Date: Tue, 21 Jul 2026 10:53:39 -0700 Subject: [PATCH 3/3] fix(openai): preserve async chat streaming response parse Await async raw-response parsing before wrapping chat streams, and preserve the awaitable parse contract on the traced proxy. Reuse cassette-backed chat and Responses tests to cover real streaming behavior without mocks. --- .../integrations/openai/test_openai.py | 84 +++++++------------ .../braintrust/integrations/openai/tracing.py | 10 ++- 2 files changed, 39 insertions(+), 55 deletions(-) diff --git a/py/src/braintrust/integrations/openai/test_openai.py b/py/src/braintrust/integrations/openai/test_openai.py index 7e2092bc..7725b3c4 100644 --- a/py/src/braintrust/integrations/openai/test_openai.py +++ b/py/src/braintrust/integrations/openai/test_openai.py @@ -7,7 +7,6 @@ import time import zlib -import httpx import openai import pytest from braintrust import Attachment, logger, wrap_openai @@ -1067,16 +1066,23 @@ async def test_openai_chat_streaming_async(memory_logger): for client, is_wrapped in clients: start = time.time() - stream = await client.chat.completions.create( - model=TEST_MODEL, - messages=[{"role": "user", "content": TEST_PROMPT}], - stream=True, - stream_options={"include_usage": True}, - ) - - chunks = [] - async for chunk in stream: - chunks.append(chunk) + create_kwargs = { + "model": TEST_MODEL, + "messages": [{"role": "user", "content": TEST_PROMPT}], + "stream": True, + "stream_options": {"include_usage": True}, + } + if is_wrapped: + async with client.chat.completions.with_streaming_response.create(**create_kwargs) as raw_response: + assert raw_response.headers + parse_result = raw_response.parse() + assert inspect.isawaitable(parse_result) + stream = await parse_result + assert stream.response + chunks = [chunk async for chunk in stream] + else: + stream = await client.chat.completions.create(**create_kwargs) + chunks = [chunk async for chunk in stream] end = time.time() assert chunks @@ -1979,7 +1985,7 @@ async def test_openai_responses_with_raw_response_async(memory_logger): @pytest.mark.asyncio @pytest.mark.vcr async def test_openai_responses_with_raw_response_create_stream_async(memory_logger): - """Async version of test_openai_responses_with_raw_response_create_stream.""" + """Async raw-response variants preserve headers, parsing, streams, and tracing.""" assert not memory_logger.pop() # Unwrapped client: headers accessible, stream iterable via parse(), no spans. @@ -1997,21 +2003,23 @@ async def test_openai_responses_with_raw_response_create_stream_async(memory_log assert "24" in "".join(chunks) or "twenty-four" in "".join(chunks).lower() assert not memory_logger.pop() - # Wrapped client: headers still accessible, parse() yields traced stream, span generated. + # Wrapped client: the newer streaming-response wrapper keeps async parse() and tracing. client = wrap_openai(AsyncOpenAI()) start = time.time() - raw = await client.responses.with_raw_response.create( + async with client.responses.with_streaming_response.create( model=TEST_MODEL, input=TEST_PROMPT, stream=True, - ) - assert raw.headers - stream = raw.parse() - assert stream.response # SDK-specific attribute preserved - chunks = [] - async for chunk in stream: - if chunk.type == "response.output_text.delta": - chunks.append(chunk.delta) + ) as raw: + assert raw.headers + parse_result = raw.parse() + assert inspect.isawaitable(parse_result) + stream = await parse_result + assert stream.response # SDK-specific attribute preserved + chunks = [] + async for chunk in stream: + if chunk.type == "response.output_text.delta": + chunks.append(chunk.delta) end = time.time() assert "24" in "".join(chunks) or "twenty-four" in "".join(chunks).lower() @@ -2783,38 +2791,6 @@ async def test_setup_async_creates_spans(self, memory_logger): assert "gpt-4o-mini" in span["metadata"]["model"] assert span["input"] - @pytest.mark.asyncio - async def test_setup_streaming_response_parse_remains_awaitable(self, memory_logger): - """OpenAI Agents awaits parse() on the streaming response wrapper.""" - assert not memory_logger.pop() - - async def handle_request(request): - return httpx.Response( - 200, - headers={"content-type": "text/event-stream"}, - content=b"data: [DONE]\n\n", - request=request, - ) - - OpenAIIntegration.setup() - transport = httpx.MockTransport(handle_request) - async with httpx.AsyncClient(transport=transport) as http_client: - client = openai.AsyncOpenAI(api_key="test-api-key", http_client=http_client) - async with client.responses.with_streaming_response.create( - model=TEST_MODEL, - input=TEST_PROMPT, - stream=True, - ) as raw_response: - parse_result = raw_response.parse() - assert inspect.isawaitable(parse_result) - stream = await parse_result - assert hasattr(stream, "__aiter__") - assert [event async for event in stream] == [] - await stream.close() - - spans = memory_logger.pop() - assert len(spans) == 1 - class TestAutoInstrumentOpenAI: """Tests for auto_instrument() with OpenAI.""" diff --git a/py/src/braintrust/integrations/openai/tracing.py b/py/src/braintrust/integrations/openai/tracing.py index aa0e3b0c..7f56ff87 100644 --- a/py/src/braintrust/integrations/openai/tracing.py +++ b/py/src/braintrust/integrations/openai/tracing.py @@ -499,8 +499,12 @@ async def acreate(self, *args: Any, **kwargs: Any) -> Any: start = time.time() create_response = await self.acreate_fn(*args, **kwargs) + parse_was_awaitable = False if hasattr(create_response, "parse"): raw_response = create_response.parse() + if inspect.isawaitable(raw_response): + parse_was_awaitable = True + raw_response = await raw_response log_headers(create_response, span) else: raw_response = create_response @@ -529,7 +533,11 @@ async def gen(): should_end = False streamer = gen() if raw_requested and hasattr(create_response, "parse"): - return _RawResponseWithTracedStream(create_response, _AsyncTracedStream(raw_response, streamer)) + return _RawResponseWithTracedStream( + create_response, + _AsyncTracedStream(raw_response, streamer), + async_parse=parse_was_awaitable, + ) return _AsyncTracedStream(raw_response, streamer) else: log_response = _try_to_dict(raw_response)