Merge pull request GreyDGL#123 from jiayuqi7813/main

Added functionality for api_base, which allows you to change openai's baseurl when using API functionality.
waltermarbel · Jun 21, 2023 · 9e59fec · 9e59fec
2 parents ab31da4 + 534d1c9
commit 9e59fec
Show file tree

Hide file tree

Showing 6 changed files with 15 additions and 31 deletions.
diff --git a/main.py b/main.py
diff --git a/pentestgpt/config/chatgpt_config.py b/pentestgpt/config/chatgpt_config.py
@@ -11,6 +11,8 @@ class ChatGPTConfig:
     # model: str = "text-davinci-002-render-sha"
     model: str = "gpt-4-browsing"
 
+    api_base: str = "https://api.openai.com/v1"
+
     # set up the openai key
     openai_key = os.getenv("OPENAI_KEY", None)
     # set the user-agent below

diff --git a/pentestgpt/main.py b/pentestgpt/main.py
@@ -12,10 +12,11 @@ def main():
     parser = argparse.ArgumentParser(description="PentestGPT")
     parser.add_argument("--reasoning_model", type=str, default="gpt-4")
     parser.add_argument("--useAPI", action="store_true", default=False)
+    parser.add_argument("--baseUrl", type=str, default="https://api.openai.com/v1", help="Base URL for OpenAI API,default: https://api.openai.com/v1")
     args = parser.parse_args()
 
     pentestGPTHandler = pentestGPT(
-        reasoning_model=args.reasoning_model, useAPI=args.useAPI
+        reasoning_model=args.reasoning_model, useAPI=args.useAPI, baseUrl=args.baseUrl
     )
 
     # you may use this one if you want to use OpenAI API (without GPT-4)

diff --git a/pentestgpt/test_connection.py b/pentestgpt/test_connection.py
@@ -5,6 +5,7 @@
 from pentestgpt.utils.chatgpt_api import ChatGPTAPI
 from pentestgpt.config.chatgpt_config import ChatGPTConfig
 import openai
+import argparse
 
 from rich.console import Console
 
@@ -13,7 +14,12 @@
 
 
 def main():
-    chatgpt_config = ChatGPTConfig()
+    parser = argparse.ArgumentParser(description="PentestGPTTestConnection")
+    parser.add_argument("--baseUrl", type=str, default="https://api.openai.com/v1", help="Base URL for OpenAI API,default: https://api.openai.com/v1")
+
+    args = parser.parse_args()
+
+    chatgpt_config = ChatGPTConfig(api_base=args.baseUrl)
     console = Console()
 
     # 1. test the connection for chatgpt cookie

diff --git a/pentestgpt/utils/chatgpt_api.py b/pentestgpt/utils/chatgpt_api.py
@@ -45,6 +45,7 @@ def __init__(self, config: ChatGPTConfig):
         self.config = config
         openai.api_key = config.openai_key
         openai.proxy = config.proxies
+        openai.api_base = config.api_base
         self.history_length = 5  # maintain 5 messages in the history. (5 chat memory)
         self.conversation_dict: Dict[str, Conversation] = {}
 

diff --git a/pentestgpt/utils/pentest_gpt.py b/pentestgpt/utils/pentest_gpt.py
@@ -48,19 +48,20 @@ class pentestGPT:
         "default": "The user did not specify the input source. You need to summarize based on the contents.\n",
     }
 
-    def __init__(self, reasoning_model="gpt-4", useAPI=True):
+    def __init__(self, reasoning_model="gpt-4", useAPI=True, baseUrl="https://api.openai.com/v1"):
         self.log_dir = "logs"
         self.save_dir = "test_history"
         self.task_log = (
             {}
         )  # the information that can be saved to continue in the next session
         self.useAPI = useAPI
+        self.baseUrl = baseUrl
         if useAPI is False:
             self.chatGPTAgent = ChatGPT(ChatGPTConfig())
             self.chatGPT4Agent = ChatGPT(ChatGPTConfig(model=reasoning_model))
         else:
             self.chatGPTAgent = ChatGPTAPI(ChatGPTConfig())
-            self.chatGPT4Agent = ChatGPTAPI(ChatGPTConfig(model=reasoning_model))
+            self.chatGPT4Agent = ChatGPTAPI(ChatGPTConfig(model=reasoning_model, api_base=baseUrl))
         self.prompts = PentestGPTPrompt
         self.console = Console()
         self.spinner = Spinner("line", "Processing")