From 171065a1dfd7902b2c142bcd10d44f367996f85a Mon Sep 17 00:00:00 2001 From: Grey_D Date: Tue, 9 May 2023 15:13:25 +0800 Subject: [PATCH] =?UTF-8?q?fix:=20=F0=9F=90=9B=20fix=20a=20bug=20in=20test?= =?UTF-8?q?=5Fconnection.py?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- test_connection.py | 8 ++++++-- utils/chatgpt.py | 20 +++++++++++++------- 2 files changed, 19 insertions(+), 9 deletions(-) diff --git a/test_connection.py b/test_connection.py index fcea98f..75fe7ce 100644 --- 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 " @@ -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 " ) 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 " ) 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 index 7b849ab..b30b833 100644 --- 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