Update main.py

main.py

@@ -5,6 +5,7 @@ from fastapi.encoders import jsonable_encoder
 from fastapi.websockets import WebSocket, WebSocketDisconnect, WebSocketState
 from fastapi.responses import HTMLResponse, JSONResponse
 
+from accelerator import Accelerator
 from answerer import Answerer
 from mapper import Mapper
 
@@ -19,7 +20,7 @@ answerer = Answerer(
   ctx_limit=16*1024,
 )
 
-accelerator: Union[WebSocket, None] = None
+accelerator = Accelerator()
   
 app = FastAPI()
 
@@ -68,9 +69,7 @@ def index():
 
 @app.websocket("/accelerate")
 async def answer(ws: WebSocket):
-  await ws.accept()
-  accelerator = ws
-  while True: await sleep(60)
+  await accelerator.connect(ws)
 
 @app.post("/map")
 def map(query: Union[str, None], items: Union[list[str], None]):
@@ -84,11 +83,10 @@ async def answer(ws: WebSocket):
   try: input = await ws.receive_text()
   except WebSocketDisconnect: return
 
-  if accelerator:
-    await accelerator.send_text(input)
-    output = await accelerator.receive_text()
+  if accelerator.connected():
+    output = await accelerator.accelerate(input)
     if(ws.client_state == WebSocketState.CONNECTED):
-      await ws.send_text(el)
+      await ws.send_text(output)
     elif(ws.client_state == WebSocketState.DISCONNECTED):
       return
   else:
@@ -100,4 +98,3 @@ async def answer(ws: WebSocket):
         return
       
   await ws.close()
-