OpenAI API: stop streaming on client disconnect (closes #4521)

oobabooga · Nov 9, 2023 · d86f1fd · d86f1fd · bodybreaker · Jun 5, 2024
1 parent effb3ae
commit d86f1fd
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/extensions/openai/script.py b/extensions/openai/script.py
@@ -91,6 +91,10 @@ async def generator():
             async with streaming_semaphore:
                 response = OAIcompletions.stream_completions(to_dict(request_data), is_legacy=is_legacy)
                 for resp in response:
+                    disconnected = await request.is_disconnected()
+                    if disconnected:
+                        break
+
                     yield {"data": json.dumps(resp)}
 
         return EventSourceResponse(generator())  # SSE streaming
@@ -110,6 +114,10 @@ async def generator():
             async with streaming_semaphore:
                 response = OAIcompletions.stream_chat_completions(to_dict(request_data), is_legacy=is_legacy)
                 for resp in response:
+                    disconnected = await request.is_disconnected()
+                    if disconnected:
+                        break
+
                     yield {"data": json.dumps(resp)}
 
         return EventSourceResponse(generator())  # SSE streaming