Update main.py

main.py

@@ -1,5 +1,7 @@
 from typing import Union
-from fastapi import FastAPI, WebSocket, WebSocketDisconnect
+from fastapi import FastAPI
+from fastapi.encoders import jsonable_encoder
+from fastapi.websockets import WebSocket, WebSocketDisconnect, WebSocketState
 from fastapi.responses import HTMLResponse, JSONResponse
 
 from answerer import Answerer
@@ -63,30 +65,21 @@ def index():
 
 @app.get("/map")
 def map(query: Union[str, None], items: Union[list[str], None]):
-  indices = mapper(query, items)
-  return JSONResponse(indices)
+  scores = mapper(query, items)
+  return JSONResponse(jsonable_encoder(scores))
 
 @app.websocket("/answer")
 async def answer(ws: WebSocket):
-  print("ws started!")
-
   await ws.accept()
-
-  print("ws accepted!")
       
   try: input = await ws.receive_text()
   except WebSocketDisconnect: return
 
-  print("input received!")
-
-  await ws.send_text("OK!")
-
   output = answerer(input, 32)
-
-  print("output created!")
-
   async for el in output:
-    print(f"sent: '{el}'")
-    await ws.send_text(el)
-
+    if(ws.client_state == WebSocketState.CONNECTED):
+      await ws.send_text(el)
+    elif(ws.client_state == WebSocketState.DISCONNECTED):
+      return
+    
   await ws.close()