Daniel Marques commited on
Commit
225e078
1 Parent(s): 3bf68c9

fix: add websocket in handlerToken

Browse files
Files changed (1) hide show
  1. main.py +1 -3
main.py CHANGED
@@ -70,8 +70,6 @@ class MyCustomSyncHandler(BaseCallbackHandler):
70
  self.end = True
71
 
72
  def on_llm_new_token(self, token: str, **kwargs) -> Any:
73
- print(token)
74
-
75
  if self.callback :
76
  self.callback(token)
77
 
@@ -262,7 +260,7 @@ async def websocket_endpoint(websocket: WebSocket, client_id: int):
262
  if (oldReceiveText != prompt):
263
  handlerToken.callback = websocket.send_text
264
  oldReceiveText = prompt
265
- asyncio.run(QA(prompt))
266
 
267
  except WebSocketDisconnect:
268
  print('disconnect')
 
70
  self.end = True
71
 
72
  def on_llm_new_token(self, token: str, **kwargs) -> Any:
 
 
73
  if self.callback :
74
  self.callback(token)
75
 
 
260
  if (oldReceiveText != prompt):
261
  handlerToken.callback = websocket.send_text
262
  oldReceiveText = prompt
263
+ await QA(prompt)
264
 
265
  except WebSocketDisconnect:
266
  print('disconnect')