From 4715123f55a04df50e4d2af6257131902d8cff07 Mon Sep 17 00:00:00 2001 From: Yiximail Date: Wed, 31 May 2023 09:03:40 +0800 Subject: [PATCH] Add a `/api/v1/stop-stream` API that allows the user to interrupt the generation (#2392) --- extensions/api/blocking_api.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/extensions/api/blocking_api.py b/extensions/api/blocking_api.py index 8c2326f4..6bcd840c 100644 --- a/extensions/api/blocking_api.py +++ b/extensions/api/blocking_api.py @@ -5,7 +5,7 @@ from threading import Thread from extensions.api.util import build_parameters, try_start_cloudflared from modules import shared from modules.chat import generate_chat_reply -from modules.text_generation import encode, generate_reply +from modules.text_generation import encode, generate_reply, stop_everything_event class Handler(BaseHTTPRequestHandler): @@ -78,6 +78,19 @@ class Handler(BaseHTTPRequestHandler): self.wfile.write(response.encode('utf-8')) + elif self.path == '/api/v1/stop-stream': + self.send_response(200) + self.send_header('Content-Type', 'application/json') + self.end_headers() + + stop_everything_event() + + response = json.dumps({ + 'results': 'success' + }) + + self.wfile.write(response.encode('utf-8')) + elif self.path == '/api/v1/token-count': self.send_response(200) self.send_header('Content-Type', 'application/json')