Update main.py

main.py

@@ -281,9 +281,12 @@ async def __query_ml_predict(qtype: QType, content: str, header: str, token_limi
             print(time() - stopwatch)
             if _r.status_code != 200: 
                 app_logger.info(_r.json())
-                return {"content": "", "style": QType.WH}
-                
-            return {"content": _r.json()['choices'][0]['message']['content'], "style": QType.WH}
+                return {"content": "", "style": None "success": False}
+
+            try:
+                return {"content": _r.json()['choices'][0]['message']['content'], "style": QType.WH, "success": True}
+            except:
+                return {"content": "", "style": None, "success": False}
         
         case QType.STMT:
 
@@ -302,11 +305,20 @@ async def __query_ml_predict(qtype: QType, content: str, header: str, token_limi
                     }),
                     timeout=None
                 )
-            if _r.status_code != 200: 
+            if _r.status_code//100 != 2: 
                 app_logger.info(_r.json())
-                return {"content": "", "style": QType.STMT}
+                return {"content": "", "style": QType.STMT, "success": False}
+
+            try:
+                _r_content = _r.json()['choices'][0]['message']['content']
+            except:
+                return {"content": "", "style":None, "success":False}
+            
+            try:
+                _r_content = _r.json()['choices'][0]['message']['content'].split('\n\n',1)[1]
+            except:
+                _r_content = _r.json()['choices'][0]['message']['content'].split('\n',1)[1]
                 
-            _r_content = _r.json()['choices'][0]['message']['content'].split('\n\n',1)[1]
             async with httpx.AsyncClient() as client:
                 _w = await client.post(
                     url="https://api.awanllm.com/v1/chat/completions",
@@ -322,12 +334,19 @@ async def __query_ml_predict(qtype: QType, content: str, header: str, token_limi
                     timeout=None
             )
 
+            try:
+                _ch = _w.json()['choices'][0]['message']['content']
+            except:
+                return {"content": "", "style": None, "success":False}
+            
             try:
                 _w_content = _w.json()['choices'][0]['message']['content'].split('\n\n',1)[1]
                 #print(time() - stopwatch)
                 return {"content": f"{_r_content}\n{_w_content}", "style": QType.STMT, "success": True}
             except:
-                return {"content": "", "style": None, "success": False}
+                _w_content = _w.json()['choices'][0]['message']['content'].split('\n',1)[1]
+                #print(time() - stopwatch)
+                return {"content": f"{_r_content}\n{_w_content}", "style": QType.STMT, "success": True}
 
 
 async def parse_wh_question(raw_qa_list, pgph_i):