fix: 🐛 fix a bug in test_connection.py

GreyDGL · May 9, 2023 · 171065a · 171065a
1 parent 52d7a2b
commit 171065a
Show file tree

Hide file tree

Showing 2 changed files with 19 additions and 9 deletions.
diff --git a/test_connection.py b/test_connection.py
@@ -8,6 +8,8 @@
 import requests
 
 logger = loguru.logger
+logger.add(level="ERROR", sink="logs/chatgpt_connection_test.log")
+
 
 if __name__ == "__main__":
     chatgpt_config = ChatGPTConfig()
@@ -16,8 +18,7 @@
     try:
         chatgpt = ChatGPT(chatgpt_config)
         conversations = chatgpt.get_conversation_history()
-        # print(conversations)
-        if conversations != None:
+        if conversations is not None:
             # print(text, conversation_id)
             print(
                 "1. You're connected with ChatGPT Plus cookie. \nTo start PentestGPT, please use <python3 main.py --reasoning_model=gpt-4>"
@@ -27,6 +28,7 @@
                 "The cookie is not properly configured with ChatGPT Cookie. Please follow README to update cookie in config/chatgpt_config.py"
             )
     except Exception as e:  # use a general exception first. Update later for debug
+        logger.error(e)
         print(
             "The cookie is not properly configured. Please follow README to update cookie in config/chatgpt_config.py"
         )
@@ -45,6 +47,7 @@
             "2. You're connected with OpenAI API. You have GPT-4 access. To start PentestGPT, please use <python3 main.py --reasoning_model=gpt-4 --useAPI>"
         )
     except Exception as e:  # use a general exception first. Update later for debug
+        logger.error(e)
         print(
             "The OpenAI API key is not properly configured. Please follow README to update OpenAI API key in config/chatgpt_config.py"
         )
@@ -63,6 +66,7 @@
             "3. You're connected with OpenAI API. You have GPT-3.5 access. To start PentestGPT, please use <python3 main.py --reasoning_model=gpt-3.5-turbo --useAPI>"
         )
     except Exception as e:  # use a general exception first. Update later for debug
+        logger.error(e)
         print(
             "The OpenAI API key is not properly configured. Please follow README to update OpenAI API key in config/chatgpt_config.py"
         )
diff --git a/utils/chatgpt.py b/utils/chatgpt.py
@@ -16,7 +16,7 @@
 
 logger = loguru.logger
 logger.remove()
-logger.add(level="WARNING", sink="logs/chatgpt.log")
+logger.add(level="ERROR", sink="logs/chatgpt.log")
 
 
 # A sample ChatGPTConfig class has the following structure. All fields can be obtained from the browser's cookie.
@@ -106,17 +106,23 @@ def get_authorization(self):
             return "Bearer " + authorization
         except requests.exceptions.JSONDecodeError as e:
             logger.error(e)
-            print(
-                "Your setting is not correct. Please update it in config/chatgpt_config.py"
+            logger.error(
+                "You encounter an error when communicating with ChatGPT. The most likely reason is that your cookie expired."
             )
-            sys.exit(1)
             return None
 
     def get_latest_message_id(self, conversation_id):
         # Get continuous conversation message id
-        url = f"https://chat.openai.com/backend-api/conversation/{conversation_id}"
-        r = requests.get(url, headers=self.headers, proxies=self.proxies)
-        return r.json()["current_node"]
+        try:
+            url = f"https://chat.openai.com/backend-api/conversation/{conversation_id}"
+            r = requests.get(url, headers=self.headers, proxies=self.proxies)
+            return r.json()["current_node"]
+        except requests.exceptions.JSONDecodeError as e:
+            logger.error(e)
+            logger.error(
+                "You encounter an error when communicating with ChatGPT. The most likely reason is that your cookie expired."
+            )
+            return None
 
     def _parse_message_raw_output(self, response: requests.Response):
         # parse message raw output