Merge pull request #154 from RiccardoRobb/main

Fast bugfix

Author: GreyDGL

pentestgpt/utils/APIs/chatgpt_api.py

@@ -93,7 +93,7 @@ def _chat_completion(
             logger.error("Token size error; will retry with compressed message ", e)
             # compress the message in two ways.
             ## 1. compress the last message
-            history[-1]["content"] = self.token_compression(history)
+            history[-1]["content"] = self._token_compression(history)
             ## 2. reduce the number of messages in the history. Minimum is 2
             if self.history_length > 2:
                 self.history_length -= 1

pentestgpt/utils/llm_api.py

@@ -96,7 +96,7 @@ def _token_compression(self, complete_messages) -> str:
         -------
             compressed_message: str
         """
-        if self.config.model == "gpt-4":
+        if self.model == "gpt-4":
             token_limit = 8000
         else:
             token_limit = 14000  # leave some budget