Error handling in websocket

BP-Live · Dec 8, 2023 · b844038 · b844038
1 parent ef0f3ea
commit b844038
Showing 1 changed file with 63 additions and 58 deletions.
diff --git a/routers/ws.py b/routers/ws.py
@@ -1,3 +1,4 @@
+import traceback
 import asyncio
 import random
 from fastapi import (
@@ -14,83 +15,87 @@
 from common import gpt_api
 
 async def websocket_endpoint(websocket: WebSocket):
-    await websocket.accept()
+    try:
+        await websocket.accept()
 
-    await websocket.send_json({"progress": 0})
+        await websocket.send_json({"progress": 0})
 
-    data = await websocket.receive_json()
-    prompt = data["prompt"]
+        data = await websocket.receive_json()
+        prompt = data["prompt"]
 
-    response_json = gpt_api.gpt_response(prompt)
+        response_json = gpt_api.gpt_response(prompt)
 
-    if response_json["location"].lower() == "user":
-        lng, lat = data["longitude"], data["latitude"]
-    else:
-        geocode = bkk_api.geocode_location(response_json["location"])[0]["geometry"]["location"]
-        lng, lat = geocode["lng"], geocode["lat"]
+        if response_json["location"].lower() == "user":
+            lng, lat = data["longitude"], data["latitude"]
+        else:
+            geocode = bkk_api.geocode_location(response_json["location"])[0]["geometry"]["location"]
+            lng, lat = geocode["lng"], geocode["lat"]
 
-    await websocket.send_json(
-        {
-            "metadata": {
-                "type": response_json["business_type"],
-                "name": response_json["business_name"],
-                "location": {"lat": lat, "lng": lng},
+        await websocket.send_json(
+            {
+                "metadata": {
+                    "type": response_json["business_type"],
+                    "name": response_json["business_name"],
+                    "location": {"lat": lat, "lng": lng},
+                }
             }
-        }
-    )
+        )
 
-    await websocket.send_json({"progress": random.randint(5, 15)})
+        await websocket.send_json({"progress": random.randint(5, 15)})
 
-    await asyncio.sleep(1)
+        await asyncio.sleep(1)
 
-    await websocket.send_json(
-        {
-            "pros": [
-                response_json["pros"]["pro1"],
-                response_json["pros"]["pro2"],
-                response_json["pros"]["pro3"],
-            ]
-        }
-    )
+        await websocket.send_json(
+            {
+                "pros": [
+                    response_json["pros"]["pro1"],
+                    response_json["pros"]["pro2"],
+                    response_json["pros"]["pro3"],
+                ]
+            }
+        )
 
-    await websocket.send_json({"progress": random.randint(20, 30)})
+        await websocket.send_json({"progress": random.randint(20, 30)})
 
-    await asyncio.sleep(1)
+        await asyncio.sleep(1)
 
-    await websocket.send_json(
-        {
-            "cons": [
-                response_json["cons"]["con1"],
-                response_json["cons"]["con2"],
-                response_json["cons"]["con3"],
-            ]
-        }
-    )
+        await websocket.send_json(
+            {
+                "cons": [
+                    response_json["cons"]["con1"],
+                    response_json["cons"]["con2"],
+                    response_json["cons"]["con3"],
+                ]
+            }
+        )
 
-    await websocket.send_json({"progress": random.randint(35, 45)})
+        await websocket.send_json({"progress": random.randint(35, 45)})
 
-    await asyncio.sleep(1)
+        await asyncio.sleep(1)
 
-    await websocket.send_json(
-        {
-            "competitors": [
-                {"lat": 37.7749, "lng": -122.42},
-                {"lat": 37.7769, "lng": -122.4194},
-                {"lat": 37.7759, "lng": -122.4184},
-            ]
-        }
-    )
+        await websocket.send_json(
+            {
+                "competitors": [
+                    {"lat": 37.7749, "lng": -122.42},
+                    {"lat": 37.7769, "lng": -122.4194},
+                    {"lat": 37.7759, "lng": -122.4184},
+                ]
+            }
+        )
 
-    await websocket.send_json({"progress": random.randint(50, 60)})
+        await websocket.send_json({"progress": random.randint(50, 60)})
 
-    await asyncio.sleep(1)
+        await asyncio.sleep(1)
 
-    await websocket.send_json({"progress": random.randint(65, 75)})
+        await websocket.send_json({"progress": random.randint(65, 75)})
 
-    await asyncio.sleep(1)
+        await asyncio.sleep(1)
 
-    await websocket.send_json({"progress": random.randint(80, 85)})
+        await websocket.send_json({"progress": random.randint(80, 85)})
 
-    await asyncio.sleep(1)
+        await asyncio.sleep(1)
 
-    await websocket.send_json({"progress": 100})
+        await websocket.send_json({"progress": 100})
+    except Exception as e:
+        await websocket.send_json({"error": str(e), "trace": traceback.format_exc()})
+    await websocket.close()