Problem
Two entangled type issues in functions/stream.py:
1. self.stream has no type annotation
BraintrustStream.__init__ assigns self.stream without an annotation (stream.py:99-101). Type checkers cannot reason about the attribute type.
2. _parse_sse_stream return type is too narrow
_parse_sse_stream declares return type Generator[BraintrustStreamChunk, None, None], but the function body yields BraintrustConsoleChunk and BraintrustProgressChunk in addition to the three types in the union:
BraintrustStreamChunk = Union[BraintrustTextChunk, BraintrustJsonChunk, BraintrustErrorChunk]
The generator actually yields 5 chunk types, not 3.
Why these are entangled
Annotating self.stream as Generator[BraintrustStreamChunk, ...] | list[BraintrustStreamChunk] (the obvious fix for issue 1) would carry forward the return type lie from issue 2. Fixing one without the other makes things worse.
Options
- Widen
BraintrustStreamChunk to include all 5 chunk types — simplest, but changes the public type that callers pattern-match against
- Create a separate internal union (e.g.
_InternalStreamChunk) for the generator, keep BraintrustStreamChunk as the public API type — more precise, but adds complexity
- Filter console/progress chunks inside
_parse_sse_stream and don't yield them — but parse_stream explicitly handles them, so they're expected in the stream
Would appreciate guidance on the preferred direction.
Problem
Two entangled type issues in
functions/stream.py:1.
self.streamhas no type annotationBraintrustStream.__init__assignsself.streamwithout an annotation (stream.py:99-101). Type checkers cannot reason about the attribute type.2.
_parse_sse_streamreturn type is too narrow_parse_sse_streamdeclares return typeGenerator[BraintrustStreamChunk, None, None], but the function body yieldsBraintrustConsoleChunkandBraintrustProgressChunkin addition to the three types in the union:The generator actually yields 5 chunk types, not 3.
Why these are entangled
Annotating
self.streamasGenerator[BraintrustStreamChunk, ...] | list[BraintrustStreamChunk](the obvious fix for issue 1) would carry forward the return type lie from issue 2. Fixing one without the other makes things worse.Options
BraintrustStreamChunkto include all 5 chunk types — simplest, but changes the public type that callers pattern-match against_InternalStreamChunk) for the generator, keepBraintrustStreamChunkas the public API type — more precise, but adds complexity_parse_sse_streamand don't yield them — butparse_streamexplicitly handles them, so they're expected in the streamWould appreciate guidance on the preferred direction.