fix(openai): add context manager support to traced stream wrappers#120
Open
Abhijeet Prasad (AbhiPrasad) wants to merge 1 commit intomainfrom
Open
fix(openai): add context manager support to traced stream wrappers#120Abhijeet Prasad (AbhiPrasad) wants to merge 1 commit intomainfrom
Abhijeet Prasad (AbhiPrasad) wants to merge 1 commit intomainfrom
Conversation
_TracedStream and _AsyncTracedStream lacked __enter__/__exit__ and __aenter__/__aexit__ methods, causing failures when callers used the OpenAI SDK's HTTP/2 streaming path which goes through LegacyAPIResponse.parse() and expects a context-manager-compatible stream. Also replace AsyncResponseWrapper and bare generator returns with _AsyncTracedStream/_TracedStream on all streaming paths so the wrapper type is consistent regardless of HTTP version. Adds a dedicated nox session (test_openai_http2_streaming) and regression tests covering sync/async context manager usage with h2.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
AI Summary
Fix OpenAI HTTP/2 streaming compatibility in the Braintrust wrapper by preserving the SDK stream interface on wrapped streaming responses.
Root Cause
When the OpenAI SDK goes through the
LegacyAPIResponse.parse()streaming path, Braintrust was returning a bare traced generator instead of a stream-like object. That dropped SDK stream internals such as_iteratorandresponse, which broke downstream consumers on HTTP/2.Primary Fix
The main fix is to return
_TracedStream/_AsyncTracedStreamfrom the wrapped streaming paths instead of returning a bare generator. This preserves the OpenAI stream interface while still routing iteration through Braintrust tracing.Context Manager Update
As part of that change, the traced stream wrappers also now implement context-manager methods:
_TracedStream.__enter__/__exit___AsyncTracedStream.__aenter__/__aexit__This ensures
with stream as s:andasync with stream as s:continue to behave correctly while keeping iteration on the traced wrapper rather than falling back to the raw OpenAI stream.Tests
Added a dedicated HTTP/2 OpenAI regression session and VCR-backed tests covering:
h2is installed only in the dedicated HTTP/2 nox session, not in the general OpenAI test session.