Update accelerator.py

accelerator.py

@@ -1,16 +1,20 @@
 from typing import Union
-from fastapi import WebSocket
+from fastapi.websockets import WebSocket, WebSocketState
 
 class Accelerator:
-  _connected = False
-  def connected(self): return self._connected
-
   ws: Union[WebSocket, None]
 
+  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()
+    while ws.client_state == WebSocketState.CONNECTING: pass
+    
+    if ws.client_state == WebSocketState.CONNECTED:
+      await self.ws.send_text(input)
+      return await self.ws.receive_text()
+    elif ws.client_state == WebSocketState.DISCONNECTED:
+      ws = None