diff --git a/src/finch/_streaming.py b/src/finch/_streaming.py index b0600fc9..913159fd 100644 --- a/src/finch/_streaming.py +++ b/src/finch/_streaming.py @@ -45,10 +45,15 @@ def __stream__(self) -> Iterator[ResponseT]: cast_to = self._cast_to response = self.response process_data = self._client._process_response_data + iterator = self._iter_events() - for sse in self._iter_events(): + for sse in iterator: yield process_data(data=sse.json(), cast_to=cast_to, response=response) + # Ensure the entire stream is consumed + for sse in iterator: + ... + class AsyncStream(Generic[ResponseT]): """Provides the core interface to iterate over an asynchronous stream response.""" @@ -83,10 +88,15 @@ async def __stream__(self) -> AsyncIterator[ResponseT]: cast_to = self._cast_to response = self.response process_data = self._client._process_response_data + iterator = self._iter_events() - async for sse in self._iter_events(): + async for sse in iterator: yield process_data(data=sse.json(), cast_to=cast_to, response=response) + # Ensure the entire stream is consumed + async for sse in iterator: + ... + class ServerSentEvent: def __init__(