fix(openai): preserve async parsing for streaming responses

py/src/braintrust/integrations/openai/test_openai.py

@@ -1066,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
@@ -1978,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.
@@ -1996,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()
 

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)
@@ -844,11 +852,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 +1080,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 +1115,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)