.

accelerator.py

@@ -1,28 +1,23 @@
-from typing import Union
-from fastapi.websockets import WebSocket, WebSocketState
+from typing import Optional
+from fastapi.websockets import WebSocket
+from websockets import ConnectionClosedError
 
 class Accelerator:
   def __del__(self): self.ws.close()
 
-  ws: Union[WebSocket, None] = None
+  ws: Optional[WebSocket] = None
 
-  def connected(self):
-    if self.ws == None: return False
-    
-    print(f"===== client_state: {self.ws.client_state}")
-
-    if self.ws.client_state == WebSocketState.CONNECTED:
-      return True
-    else:
-      self.ws = None
-      return False
+  def connected(self): return self.ws != None
 
   async def connect(self, ws: WebSocket):
     await ws.accept()
     self.ws = ws
 
-  async def accelerate(self, input):    
-    await self.ws.send_text(input)
-    return await self.ws.receive_text()
-  
+  async def accelerate(self, input):   
+    try:
+      await self.ws.send_text(input)
+      return await self.ws.receive_text()
+    except ConnectionClosedError:
+      self.ws = None
+      return None
 

main.py

@@ -1,5 +1,5 @@
 from asyncio import sleep
-from typing import Union
+from typing import Optional
 from fastapi import FastAPI
 from fastapi.encoders import jsonable_encoder
 from fastapi.websockets import WebSocket, WebSocketDisconnect
@@ -73,23 +73,28 @@ async def answer(ws: WebSocket):
     await sleep(10)
 
 @app.post("/map")
-def map(query: Union[str, None], items: Union[list[str], None]):
+def map(query: Optional[str], items: Optional[list[str]]):
   scores = mapper(query, items)
   return JSONResponse(jsonable_encoder(scores))
 
+async def handle_answerer_local(ws: WebSocket, input: str):
+  output = answerer(input, 32)
+  async for el in output:
+    await ws.send_text(el)
+
+async def handle_answerer_accelerated(ws: WebSocket, input: str):
+  output = await accelerator.accelerate(input)
+  if output: await ws.send_text(output)
+  else: handle_answerer_local(ws, input)
+
 @app.websocket("/answer")
 async def answer(ws: WebSocket):
   await ws.accept()
       
   try: 
     input = await ws.receive_text()
-    if accelerator.connected():
-      output = await accelerator.accelerate(input)
-      await ws.send_text(output)
-    else:
-      output = answerer(input, 32)
-      async for el in output:
-        await ws.send_text(el)
+    if accelerator.connected(): handle_answerer_accelerated(ws, input)
+    else: handle_answerer_local(ws, input)
   except WebSocketDisconnect: return
       
   await ws.close()