Update app.py

app.py

@@ -130,36 +130,40 @@ conversational_rag_chain = RunnableWithMessageHistory(
 )
 
 async def echo(websocket):
-    global retriever, conversational_rag_chain
-    async for message in websocket:
-        data = json.loads(message)
-        if data["message"] == "data.":
-            response = store
-            await websocket.send(json.dumps({"response": str(response)}))
-            continue
-        if not "message" in message:
-            return
-        if not "token" in message:
-            return
-        m = data["message"] + "\n\nAssistant: "
-        token = data["token"]
-        docs = retriever.get_relevant_documents(m)
-        rawresponse = conversational_rag_chain.invoke(
-            {"input": m},
-            config={
-                "configurable": {"session_id": token}
-            },
-        )
-        response = rawresponse["answer"]
-        response = response.replace("Assistant: ", "").replace("AI: ", "")
-        response.strip()
-        response = response.split("Human:")[0]
-        while response.startswith("\n"):
-            response = response[1:]
-        await websocket.send(json.dumps({"response": response}))
+    try:
+        async for message in websocket:
+            data = json.loads(message)
+            if data["message"] == "data.":
+                response = store
+                await websocket.send(json.dumps({"response": str(response)}))
+                continue
+            if not "message" in message:
+                return
+            if not "token" in message:
+                return
+            m = data["message"] + "\n\nAssistant: "
+            token = data["token"]
+            docs = retriever.get_relevant_documents(m)
+            rawresponse = conversational_rag_chain.invoke(
+                {"input": m},
+                config={
+                    "configurable": {"session_id": token}
+                },
+            )
+            response = rawresponse["answer"]
+            response = response.replace("Assistant: ", "").replace("AI: ", "")
+            response.strip()
+            response = response.split("Human:")[0]
+            while response.startswith("\n"):
+                response = response[1:]
+            await websocket.send(json.dumps({"response": response}))
+    except ConnectionClosedError:
+        pass
 
 async def main():
-    async with serve(echo, "0.0.0.0", 7860):
-        await asyncio.Future()
-
+    try:
+        async with serve(echo, "0.0.0.0", 7860):
+            await asyncio.Future()
+    except ConnectionClosedError:
+        pass
 asyncio.run(main())