DaniilAlpha commited on
Commit
96dc6e1
·
1 Parent(s): ac443a4
Files changed (2) hide show
  1. accelerator.py +12 -17
  2. main.py +14 -9
accelerator.py CHANGED
@@ -1,28 +1,23 @@
1
- from typing import Union
2
- from fastapi.websockets import WebSocket, WebSocketState
 
3
 
4
  class Accelerator:
5
  def __del__(self): self.ws.close()
6
 
7
- ws: Union[WebSocket, None] = None
8
 
9
- def connected(self):
10
- if self.ws == None: return False
11
-
12
- print(f"===== client_state: {self.ws.client_state}")
13
-
14
- if self.ws.client_state == WebSocketState.CONNECTED:
15
- return True
16
- else:
17
- self.ws = None
18
- return False
19
 
20
  async def connect(self, ws: WebSocket):
21
  await ws.accept()
22
  self.ws = ws
23
 
24
- async def accelerate(self, input):
25
- await self.ws.send_text(input)
26
- return await self.ws.receive_text()
27
-
 
 
 
28
 
 
1
+ from typing import Optional
2
+ from fastapi.websockets import WebSocket
3
+ from websockets import ConnectionClosedError
4
 
5
  class Accelerator:
6
  def __del__(self): self.ws.close()
7
 
8
+ ws: Optional[WebSocket] = None
9
 
10
+ def connected(self): return self.ws != None
 
 
 
 
 
 
 
 
 
11
 
12
  async def connect(self, ws: WebSocket):
13
  await ws.accept()
14
  self.ws = ws
15
 
16
+ async def accelerate(self, input):
17
+ try:
18
+ await self.ws.send_text(input)
19
+ return await self.ws.receive_text()
20
+ except ConnectionClosedError:
21
+ self.ws = None
22
+ return None
23
 
main.py CHANGED
@@ -1,5 +1,5 @@
1
  from asyncio import sleep
2
- from typing import Union
3
  from fastapi import FastAPI
4
  from fastapi.encoders import jsonable_encoder
5
  from fastapi.websockets import WebSocket, WebSocketDisconnect
@@ -73,23 +73,28 @@ async def answer(ws: WebSocket):
73
  await sleep(10)
74
 
75
  @app.post("/map")
76
- def map(query: Union[str, None], items: Union[list[str], None]):
77
  scores = mapper(query, items)
78
  return JSONResponse(jsonable_encoder(scores))
79
 
 
 
 
 
 
 
 
 
 
 
80
  @app.websocket("/answer")
81
  async def answer(ws: WebSocket):
82
  await ws.accept()
83
 
84
  try:
85
  input = await ws.receive_text()
86
- if accelerator.connected():
87
- output = await accelerator.accelerate(input)
88
- await ws.send_text(output)
89
- else:
90
- output = answerer(input, 32)
91
- async for el in output:
92
- await ws.send_text(el)
93
  except WebSocketDisconnect: return
94
 
95
  await ws.close()
 
1
  from asyncio import sleep
2
+ from typing import Optional
3
  from fastapi import FastAPI
4
  from fastapi.encoders import jsonable_encoder
5
  from fastapi.websockets import WebSocket, WebSocketDisconnect
 
73
  await sleep(10)
74
 
75
  @app.post("/map")
76
+ def map(query: Optional[str], items: Optional[list[str]]):
77
  scores = mapper(query, items)
78
  return JSONResponse(jsonable_encoder(scores))
79
 
80
+ async def handle_answerer_local(ws: WebSocket, input: str):
81
+ output = answerer(input, 32)
82
+ async for el in output:
83
+ await ws.send_text(el)
84
+
85
+ async def handle_answerer_accelerated(ws: WebSocket, input: str):
86
+ output = await accelerator.accelerate(input)
87
+ if output: await ws.send_text(output)
88
+ else: handle_answerer_local(ws, input)
89
+
90
  @app.websocket("/answer")
91
  async def answer(ws: WebSocket):
92
  await ws.accept()
93
 
94
  try:
95
  input = await ws.receive_text()
96
+ if accelerator.connected(): handle_answerer_accelerated(ws, input)
97
+ else: handle_answerer_local(ws, input)
 
 
 
 
 
98
  except WebSocketDisconnect: return
99
 
100
  await ws.close()