.

answerer.py

@@ -68,6 +68,9 @@ class Answerer:
       if "\ufffd" not in tmp:
         tokens.clear()
         result += tmp
+        if result.rstrip().endswith("\n\nUser:"):
+          yield result.rstrip().removesuffix("\n\nUser:")
+          break
         yield result
         await sleep(.02)
 

main.py

@@ -4,6 +4,7 @@ from fastapi import FastAPI
 from fastapi.encoders import jsonable_encoder
 from fastapi.websockets import WebSocket, WebSocketDisconnect
 from fastapi.responses import HTMLResponse, JSONResponse
+from websockets import ConnectionClosed
 
 from accelerator import Accelerator
 from answerer import Answerer
@@ -78,7 +79,7 @@ def map(query: Optional[str], items: Optional[list[str]]):
   return JSONResponse(jsonable_encoder(scores))
 
 async def handle_answerer_local(ws: WebSocket, input: str):
-  output = answerer(msg, 128)
+  output = answerer(input, 128)
   el: str
   async for el in output: pass
   ws.send_text(el)