Merge pull request #140 from GreyDGL/local-llm

Local llm
GreyDGL · Jul 25, 2023 · 776aac1 · 776aac1
2 parents 90b4c4f + 85726d1
commit 776aac1
Show file tree

Hide file tree

Showing 7 changed files with 45 additions and 34 deletions.
diff --git a/pentestgpt/test_connection.py b/pentestgpt/test_connection.py
@@ -33,64 +33,62 @@ def main():
     chatgpt_config = ChatGPTConfig(api_base=args.baseUrl)
     console = Console()
 
-    # 1. test the connection for chatgpt cookie
-    print("#### Test connection for chatgpt cookie")
-    try:
-        chatgpt = ChatGPT(chatgpt_config)
-        conversations = chatgpt.get_conversation_history()
-        if conversations is not None:
-            console.print(
-                "1. You're connected with ChatGPT Plus cookie. \nTo start PentestGPT, please use <pentestgpt --reasoning_model=gpt-4>",
-                style="bold green",
-            )
-        else:
-            console.print(
-                "The cookie is not properly configured with ChatGPT Cookie. Please follow README to update cookie through `export CHATGPT_COOKIE=<your cookie>`",
-                style="bold red",
-            )
-    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"
-        )
-
-    # 2. test the connection for chatgpt api with GPT-4
+    # 1. test the connection for chatgpt api with GPT-4
     print("#### Test connection for OpenAI api (GPT-4)")
     try:
         chatgpt_config.model = "gpt-4"
         chatgpt = ChatGPTAPI(chatgpt_config)
         openai.api_key = chatgpt_config.openai_key
         result, conversation_id = chatgpt.send_new_message("Hi how are you?")
         console.print(
-            "2. You're connected with OpenAI API. You have GPT-4 access. To start PentestGPT, please use <pentestgpt --reasoning_model=gpt-4 --useAPI>",
+            "1. You're connected with OpenAI API. You have GPT-4 access. To start PentestGPT, please use <pentestgpt --reasoning_model=gpt-4>",
             style="bold green",
         )
     except Exception as e:  # use a general exception first. Update later for debug
         console.print(
-            "The OpenAI API key is not properly configured. Please follow README to update OpenAI API key through `export OPENAI_KEY=<>`.",
+            "1. The OpenAI API key is not properly configured. Please follow README to update OpenAI API key through `export OPENAI_KEY=<>`.",
             style="bold red",
         )
         print("The error is below:", e)
 
-    # 3. test the connection for chatgpt api with GPT-3.5
+    # 2. test the connection for chatgpt api with GPT-3.5
     print("#### Test connection for OpenAI api (GPT-3.5)")
     try:
-        chatgpt_config.model = "gpt-3.5-turbo"
+        chatgpt_config.model = "gpt-3.5-turbo-16k"
         chatgpt = ChatGPTAPI(chatgpt_config)
         openai.api_key = chatgpt_config.openai_key
         result, conversation_id = chatgpt.send_new_message("Hi how are you?")
         console.print(
-            "3. You're connected with OpenAI API. You have GPT-3.5 access. To start PentestGPT, please use <pentestgpt --reasoning_model=gpt-3.5-turbo --useAPI>",
+            "2. You're connected with OpenAI API. You have GPT-3.5 access. To start PentestGPT, please use <pentestgpt --reasoning_model=gpt-3.5-turbo-16k>",
             style="bold green",
         )
     except Exception as e:  # use a general exception first. Update later for debug
         logger.error(e)
         console.print(
-            "The OpenAI API key is not properly configured. Please follow README to update OpenAI API key through `export OPENAI_KEY=<>`",
+            "2. The OpenAI API key is not properly configured. The likely reason is that you do not link a payment method to OpenAI so your key is not active. \nPlease follow README to update OpenAI API key through `export OPENAI_KEY=<>`",
             style="bold red",
         )
         print("The error is below:", e)
 
+    # 3. test the connection for chatgpt cookie (deprecated)
+    print("#### Test connection for chatgpt cookie")
+    try:
+        chatgpt = ChatGPT(chatgpt_config)
+        conversations = chatgpt.get_conversation_history()
+        if conversations is not None:
+            console.print(
+                "3. You're connected with ChatGPT Plus cookie. \nTo start PentestGPT, please use <pentestgpt --reasoning_model=gpt-4>",
+                style="bold green",
+            )
+        else:
+            console.print(
+                "3. The cookie is not properly configured with ChatGPT Cookie. If you're not using cookie bypass for testing, please neglect this message.",
+                style="bold red",
+            )
+    except Exception as e:  # use a general exception first. Update later for debug
+        logger.error(e)
+        print("The cookie is not properly configured.")
+
 
 if __name__ == "__main__":
     main()
diff --git a/pentestgpt/utils/APIs/chatgpt_api.py b/pentestgpt/utils/APIs/chatgpt_api.py
@@ -43,6 +43,7 @@ def __eq__(self, other):
 
 class ChatGPTAPI(LLMAPI):
     def __init__(self, config_class):
+        self.name = str(config_class.model)
         openai.api_key = os.getenv("OPENAI_KEY", None)
         openai.api_base = config_class.api_base
         self.model = config_class.model

diff --git a/pentestgpt/utils/APIs/gpt4all_api.py b/pentestgpt/utils/APIs/gpt4all_api.py
@@ -45,6 +45,7 @@ def __eq__(self, other):
 
 class GPT4ALLAPI(LLMAPI):
     def __init__(self, config_class):
+        self.name = str(config_class.model)
         self.history_length = (
             2  # maintain 2 messages in the history due to gpt4all limitation.
         )
@@ -71,7 +72,7 @@ def _chat_completion(self, history: List) -> str:
 
 if __name__ == "__main__":
     chatgpt_config = ChatGPTConfig()
-    chatgpt = ChatGPTAPI()
+    chatgpt = GPT4ALLAPI()
 
     # test is below
     # 1. create a new conversation

diff --git a/pentestgpt/utils/APIs/module_import.py b/pentestgpt/utils/APIs/module_import.py
@@ -70,8 +70,14 @@ def dynamic_import(module_name, log_dir) -> object:
         return LLM_class_initialized
 
     else:
-        print("Module not found: " + module_name)
-        return None
+        print(
+            "Module not found: "
+            + module_name
+            + ". Falling back to use the default gpt-3.5-turbo-16k"
+        )
+        # fall back to gpt-3.5-turbo-16k
+        LLM_class_initialized = dynamic_import("gpt-3.5-turbo-16k", log_dir)
+        return LLM_class_initialized
 
 
 if __name__ == "__main__":

diff --git a/pentestgpt/utils/llm_api.py b/pentestgpt/utils/llm_api.py
@@ -45,6 +45,7 @@ def __eq__(self, other):
 
 class LLMAPI:
     def __init__(self, config: ChatGPTConfig):
+        self.name = "LLMAPI_base_class"
         self.config = config
         openai.api_key = config.openai_key
         openai.proxy = config.proxies

diff --git a/pentestgpt/utils/pentest_gpt.py b/pentestgpt/utils/pentest_gpt.py
@@ -95,8 +95,12 @@ def __init__(
             style="bold green",
         )
         self.console.print("The settings are: ")
-        self.console.print(f" - parsing model: {parsing_model}", style="bold green")
-        self.console.print(f" - reasoning model: {reasoning_model}", style="bold green")
+        self.console.print(
+            f" - parsing model: {parsing_model_object.name}", style="bold green"
+        )
+        self.console.print(
+            f" - reasoning model: {reasoning_model_object.name}", style="bold green"
+        )
         self.console.print(f" - use API: {useAPI}", style="bold green")
         self.console.print(f" - log directory: {log_dir}", style="bold green")
 

diff --git a/setup.py b/setup.py
@@ -7,7 +7,7 @@
 
 setup(
     name="pentestgpt",
-    version="0.9.0",
+    version="0.9.1",
     description="PentestGPT, a GPT-empowered penetration testing tool",
     long_description="""
     PentestGPT is a penetration testing tool empowered by ChatGPT.