State checking replaced with try-except

main.py

@@ -2,7 +2,7 @@ from asyncio import sleep
 from typing import Union
 from fastapi import FastAPI
 from fastapi.encoders import jsonable_encoder
-from fastapi.websockets import WebSocket, WebSocketDisconnect, WebSocketState
+from fastapi.websockets import WebSocket, WebSocketDisconnect
 from fastapi.responses import HTMLResponse, JSONResponse
 
 from accelerator import Accelerator
@@ -80,21 +80,15 @@ def map(query: Union[str, None], items: Union[list[str], None]):
 async def answer(ws: WebSocket):
   await ws.accept()
       
-  try: input = await ws.receive_text()
-  except WebSocketDisconnect: return
-
-  if accelerator.connected():
-    output = await accelerator.accelerate(input)
-    if(ws.client_state == WebSocketState.CONNECTED):
+  try: 
+    input = await ws.receive_text()
+    if accelerator.connected():
+      output = await accelerator.accelerate(input)
       await ws.send_text(output)
-    elif(ws.client_state == WebSocketState.DISCONNECTED):
-      return
-  else:
-    output = answerer(input, 32)
-    async for el in output:
-      if(ws.client_state == WebSocketState.CONNECTED):
+    else:
+      output = answerer(input, 32)
+      async for el in output:
         await ws.send_text(el)
-      elif(ws.client_state == WebSocketState.DISCONNECTED):
-        return
+  except WebSocketDisconnect: return
       
   await ws.close()