From d86f1fd2c3692f5fe015dd65b99f1ceb35ac28ca Mon Sep 17 00:00:00 2001
From: oobabooga <112222186+oobabooga@users.noreply.github.com>
Date: Thu, 9 Nov 2023 06:37:32 -0800
Subject: [PATCH] OpenAI API: stop streaming on client disconnect (closes
 #4521)

---
 extensions/openai/script.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/extensions/openai/script.py b/extensions/openai/script.py
index 94e4160e..40574a72 100644
--- a/extensions/openai/script.py
+++ b/extensions/openai/script.py
@@ -91,6 +91,10 @@ async def openai_completions(request: Request, request_data: CompletionRequest):
             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 openai_chat_completions(request: Request, request_data: ChatCompletion
             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