Update main.py

main.py

@@ -29,7 +29,8 @@ HTML = """
     const prompt = document.getElementById("prompt");
     const output = document.getElementById("output");
 
-    const ws = new WebSocket("wss://daniilalpha-answerer-api.hf.space/answer");  
+    // TODO!
+    const ws = new WebSocket("wss://daniilalpha-answerer-api.hf.space/answer_test"); 
     ws.onmessage = (e) => answer(e.data);
 
     function ask(event) {
@@ -64,4 +65,19 @@ async def answer(ws: WebSocket):
   async for el in output:
     await ws.send_text(el) 
 
+  await ws.close()
+
+async def answerer_test():
+  for i in range(10000):
+    yield i
+
+@app.websocket("/answer_test")
+async def answer(ws: WebSocket):
+  await ws.accept()
+
+  input = await ws.receive_text()
+  output = answerer_test()
+  async for el in output:
+    await ws.send_text(el) 
+
   await ws.close()