From 5bd9cef9fac21630dc0ee56467a18804b90b5e7e Mon Sep 17 00:00:00 2001
From: Concedo <39025047+LostRuins@users.noreply.github.com>
Date: Fri, 9 Jun 2023 23:22:16 +0800
Subject: [PATCH] merging Proper SSE Token Streaming #220 with end connection
 fix test

---
 koboldcpp.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/koboldcpp.py b/koboldcpp.py
index 981513987..84aaa7699 100644
--- a/koboldcpp.py
+++ b/koboldcpp.py
@@ -289,7 +289,7 @@ class ServerRequestHandler(http.server.SimpleHTTPRequestHandler):
         self.send_header("Connection", "keep-alive")
         self.end_headers()
 
-        current_token = 0;
+        current_token = 0
 
         while not handle.has_finished():
             if current_token < handle.get_stream_count():
@@ -307,7 +307,10 @@ class ServerRequestHandler(http.server.SimpleHTTPRequestHandler):
 
             await asyncio.sleep(0)
 
-        # Implement connection closing here
+        # flush buffers, sleep a bit to make sure all data sent, and then force close the connection
+        self.wfile.flush()
+        await asyncio.sleep(0.1)
+        self.close_connection = True
 
 
     async def handle_request(self, genparams, newprompt, basic_api_flag, stream_flag):
@@ -417,6 +420,7 @@ class ServerRequestHandler(http.server.SimpleHTTPRequestHandler):
             try:
                 genparams = json.loads(body)
             except ValueError as e:
+                utfprint("Body Err: " + str(body))
                 return self.send_response(503)
 
             utfprint("\nInput: " + json.dumps(genparams))