Mirrors/text-generation-webui
1
0
Fork 0
mirror of https://github.com/oobabooga/text-generation-webui.git synced 2024-11-22 08:07:56 +01:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
oobabooga 2023-11-09 06:37:32 -08:00
parent effb3aef42
commit d86f1fd2c3
1 changed files with 8 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
extensions/openai/script.py
View File

@ -91,6 +91,10 @@ async def openai_completions(request: Request, request_data: CompletionRequest):
async with streaming_semaphore: async with streaming_semaphore:
response = OAIcompletions.stream_completions(to_dict(request_data), is_legacy=is_legacy) response = OAIcompletions.stream_completions(to_dict(request_data), is_legacy=is_legacy)
for resp in response: for resp in response:
disconnected = await request.is_disconnected()
if disconnected:
break
yield {"data": json.dumps(resp)} yield {"data": json.dumps(resp)}
return EventSourceResponse(generator()) # SSE streaming return EventSourceResponse(generator()) # SSE streaming
@ -110,6 +114,10 @@ async def openai_chat_completions(request: Request, request_data: ChatCompletion
async with streaming_semaphore: async with streaming_semaphore:
response = OAIcompletions.stream_chat_completions(to_dict(request_data), is_legacy=is_legacy) response = OAIcompletions.stream_chat_completions(to_dict(request_data), is_legacy=is_legacy)
for resp in response: for resp in response:
disconnected = await request.is_disconnected()
if disconnected:
break
yield {"data": json.dumps(resp)} yield {"data": json.dumps(resp)}
return EventSourceResponse(generator()) # SSE streaming return EventSourceResponse(generator()) # SSE streaming