accelerator.py

from typing import Union
from fastapi.websockets import WebSocket, WebSocketState
from websockets import ConnectionClosedError

class Accelerator:
  ws: Union[WebSocket, None] = None

  def connected(self):
    if self.ws == None: return False
    
    if self.ws.client_state == WebSocketState.CONNECTED:
      return True
    else:
      self.ws = None
      return False

  async def connect(self, ws: WebSocket):
    await ws.accept()
    self.ws = ws

  async def accelerate(self, input):    
    try:
      await self.ws.send_text(input)
      return await self.ws.receive_text()
    except ConnectionClosedError:
      self.ws = None