chore: patch streaming

groq · May 22, 2024 · 9dcba61 · 9dcba61
1 parent 5056621
commit 9dcba61
Show file tree

Hide file tree

Showing 6 changed files with 374 additions and 117 deletions.
diff --git a/src/groq/_streaming.py b/src/groq/_streaming.py
@@ -55,6 +55,8 @@ def __stream__(self) -> Iterator[_T]:
         iterator = self._iter_events()
 
         for sse in iterator:
+            if sse.data.startswith("[DONE]"):
+                break
             yield process_data(data=sse.json(), cast_to=cast_to, response=response)
 
         # Ensure the entire stream is consumed
@@ -119,6 +121,8 @@ async def __stream__(self) -> AsyncIterator[_T]:
         iterator = self._iter_events()
 
         async for sse in iterator:
+            if sse.data.startswith("[DONE]"):
+                break
             yield process_data(data=sse.json(), cast_to=cast_to, response=response)
 
         # Ensure the entire stream is consumed

diff --git a/src/groq/lib/chat_completion_chunk.py b/src/groq/lib/chat_completion_chunk.py