DaniilAlpha commited on
Commit
71f5570
·
1 Parent(s): c7f3620

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +6 -9
main.py CHANGED
@@ -5,6 +5,7 @@ from fastapi.encoders import jsonable_encoder
5
  from fastapi.websockets import WebSocket, WebSocketDisconnect, WebSocketState
6
  from fastapi.responses import HTMLResponse, JSONResponse
7
 
 
8
  from answerer import Answerer
9
  from mapper import Mapper
10
 
@@ -19,7 +20,7 @@ answerer = Answerer(
19
  ctx_limit=16*1024,
20
  )
21
 
22
- accelerator: Union[WebSocket, None] = None
23
 
24
  app = FastAPI()
25
 
@@ -68,9 +69,7 @@ def index():
68
 
69
  @app.websocket("/accelerate")
70
  async def answer(ws: WebSocket):
71
- await ws.accept()
72
- accelerator = ws
73
- while True: await sleep(60)
74
 
75
  @app.post("/map")
76
  def map(query: Union[str, None], items: Union[list[str], None]):
@@ -84,11 +83,10 @@ async def answer(ws: WebSocket):
84
  try: input = await ws.receive_text()
85
  except WebSocketDisconnect: return
86
 
87
- if accelerator:
88
- await accelerator.send_text(input)
89
- output = await accelerator.receive_text()
90
  if(ws.client_state == WebSocketState.CONNECTED):
91
- await ws.send_text(el)
92
  elif(ws.client_state == WebSocketState.DISCONNECTED):
93
  return
94
  else:
@@ -100,4 +98,3 @@ async def answer(ws: WebSocket):
100
  return
101
 
102
  await ws.close()
103
-
 
5
  from fastapi.websockets import WebSocket, WebSocketDisconnect, WebSocketState
6
  from fastapi.responses import HTMLResponse, JSONResponse
7
 
8
+ from accelerator import Accelerator
9
  from answerer import Answerer
10
  from mapper import Mapper
11
 
 
20
  ctx_limit=16*1024,
21
  )
22
 
23
+ accelerator = Accelerator()
24
 
25
  app = FastAPI()
26
 
 
69
 
70
  @app.websocket("/accelerate")
71
  async def answer(ws: WebSocket):
72
+ await accelerator.connect(ws)
 
 
73
 
74
  @app.post("/map")
75
  def map(query: Union[str, None], items: Union[list[str], None]):
 
83
  try: input = await ws.receive_text()
84
  except WebSocketDisconnect: return
85
 
86
+ if accelerator.connected():
87
+ output = await accelerator.accelerate(input)
 
88
  if(ws.client_state == WebSocketState.CONNECTED):
89
+ await ws.send_text(output)
90
  elif(ws.client_state == WebSocketState.DISCONNECTED):
91
  return
92
  else:
 
98
  return
99
 
100
  await ws.close()