File tree Expand file tree Collapse file tree 1 file changed +12
-2
lines changed Expand file tree Collapse file tree 1 file changed +12
-2
lines changed Original file line number Diff line number Diff line change @@ -47,8 +47,9 @@ def __stream__(self) -> Iterator[ResponseT]:
4747 cast_to = self ._cast_to
4848 response = self .response
4949 process_data = self ._client ._process_response_data
50+ iterator = self ._iter_events ()
5051
51- for sse in self . _iter_events () :
52+ for sse in iterator :
5253 if sse .data .startswith ("[DONE]" ):
5354 break
5455
@@ -63,6 +64,10 @@ def __stream__(self) -> Iterator[ResponseT]:
6364
6465 yield process_data (data = data , cast_to = cast_to , response = response )
6566
67+ # Ensure the entire stream is consumed
68+ for sse in iterator :
69+ ...
70+
6671
6772class AsyncStream (Generic [ResponseT ]):
6873 """Provides the core interface to iterate over an asynchronous stream response."""
@@ -97,8 +102,9 @@ async def __stream__(self) -> AsyncIterator[ResponseT]:
97102 cast_to = self ._cast_to
98103 response = self .response
99104 process_data = self ._client ._process_response_data
105+ iterator = self ._iter_events ()
100106
101- async for sse in self . _iter_events () :
107+ async for sse in iterator :
102108 if sse .data .startswith ("[DONE]" ):
103109 break
104110
@@ -113,6 +119,10 @@ async def __stream__(self) -> AsyncIterator[ResponseT]:
113119
114120 yield process_data (data = data , cast_to = cast_to , response = response )
115121
122+ # Ensure the entire stream is consumed
123+ async for sse in iterator :
124+ ...
125+
116126
117127class ServerSentEvent :
118128 def __init__ (
You can’t perform that action at this time.
0 commit comments