Merge pull request #436 from airtai/handle-lifespan-exception

Handle lifespan exception and update script
airtai · Mar 8, 2024 · 8bb0250 · 8bb0250
2 parents fe15bee + e5894fc
commit 8bb0250
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 31 deletions.
diff --git a/captn/captn_agents/application.py b/captn/captn_agents/application.py
@@ -33,39 +33,45 @@ class DailyAnalysisRequest(BaseModel):
 
 def on_connect(iostream: IOWebsockets, num_of_retries: int = 3) -> None:
     try:
-        message = iostream.input()
-    except Exception as e:
-        iostream.print(f"Failed to read the message from the client: {e}")
-        return
-    for i in range(num_of_retries):
         try:
-            request_json = message
-            request = CaptnAgentRequest.model_validate_json(request_json)
-            print("===============================================")
-            print(f"Received request: {request}", flush=True)
-            _, last_message = start_or_continue_conversation(
-                user_id=request.user_id,
-                conv_id=request.conv_id,
-                task=request.message,
-                iostream=iostream,
-                max_round=80,
-                human_input_mode="NEVER",
-                class_name="google_ads_team",
+            message = iostream.input()
+        except Exception as e:
+            iostream.print(
+                "We are sorry, but we are unable to continue the conversation. Please create a new chat in a few minutes to continue."
             )
-            last_message_dict = ast.literal_eval(last_message)
-            iostream.print(json.dumps(last_message_dict))
+            print(f"Failed to read the message from the client: {e}")
             return
-
-        except Exception as e:
-            # TODO: error logging
-            iostream.print(f"Agent conversation failed with an error: {e}")
-            if i < num_of_retries - 1:
-                iostream.print("Retrying the whole conversation...")
-                iostream.print("*" * 100)
-
-    iostream.print(
-        "We are sorry, but we are unable to continue the conversation. Please create a new chat in a few minutes to continue."
-    )
+        for i in range(num_of_retries):
+            try:
+                request_json = message
+                request = CaptnAgentRequest.model_validate_json(request_json)
+                print("===============================================")
+                print(f"Received request: {request}", flush=True)
+                _, last_message = start_or_continue_conversation(
+                    user_id=request.user_id,
+                    conv_id=request.conv_id,
+                    task=request.message,
+                    iostream=iostream,
+                    max_round=80,
+                    human_input_mode="NEVER",
+                    class_name="google_ads_team",
+                )
+                last_message_dict = ast.literal_eval(last_message)
+                iostream.print(json.dumps(last_message_dict))
+                return
+
+            except Exception as e:
+                # TODO: error logging
+                iostream.print(f"Agent conversation failed with an error: {e}")
+                if i < num_of_retries - 1:
+                    iostream.print("Retrying the whole conversation...")
+                    iostream.print("*" * 100)
+
+        iostream.print(
+            "We are sorry, but we are unable to continue the conversation. Please create a new chat in a few minutes to continue."
+        )
+    except Exception as e:
+        print(f"Agent conversation failed with an error: {e}")
 
 
 @router.post("/chat")

diff --git a/scripts/start_webservice.sh b/scripts/start_webservice.sh
@@ -12,4 +12,4 @@ cat <<< "$CLIENT_SECRET" > client_secret.json
 prisma migrate deploy
 prisma generate
 
-uvicorn application:app --port $PORT --host 0.0.0.0 --workers=$NUM_WORKERS --proxy-headers
+uvicorn application:app --port $PORT --host 0.0.0.0 --workers=$NUM_WORKERS --proxy-headers --lifespan on