Merge pull request #140 from GreyDGL/local-llm

Local llm

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()

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

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

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__":

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

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")
 

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.