Add support for customizing API Base URL using environment variables (#207)

* Update chatgpt_config.py

Add base api environment variables

* Update README.md

Add guidance for API base changes

* Fix the problem that test-connection cannot run

* Update test_connection.py

Class variable read modification

Author: wyl2003

README.md

@@ -61,7 +61,7 @@
 ## Quick Start
 1. Create a virtual environment if necessary. (`virtualenv -p python3 venv`, `source venv/bin/activate`)
 2. Install the project with `pip3 install git+https://github.com/GreyDGL/PentestGPT`
-3. **Ensure that you have link a payment method to your OpenAI account.** Export your API key with `export OPENAI_KEY='<your key here>'`
+3. **Ensure that you have link a payment method to your OpenAI account.** Export your API key with `export OPENAI_KEY='<your key here>'`,export API base with `export OPENAI_BASEURL='https://api.xxxx.xxx/v1'`if you need.
 4. Test the connection with `pentestgpt-connection`
 5. For Kali Users: use `tmux` as terminal environment. You can do so by simply run `tmux` in the native terminal.
 6. To start: `pentestgpt --logging`
@@ -101,6 +101,7 @@ PentestGPT is tested under `Python 3.10`. Other Python3 versions should work but
 2. To use OpenAI API
    - **Ensure that you have link a payment method to your OpenAI account.**
    - export your API key with `export OPENAI_KEY='<your key here>'`
+   - export API base with `export OPENAI_BASEURL='https://api.xxxx.xxx/v1'`if you need.
    - Test the connection with `pentestgpt-connection`
 3. To verify that the connection is configured properly, you may run `pentestgpt-connection`. After a while, you should see some sample conversation with ChatGPT.
    - A sample output is below

pentestgpt/config/chatgpt_config.py

@@ -12,7 +12,9 @@ class ChatGPTConfig:
     # model: str = "text-davinci-002-render-sha"
     model: str = "gpt-4-browsing"
 
-    api_base: str = "https://api.openai.com/v1"
+    #api_base: str = "https://api.openai.com/v1"
+    #set up the openai api base, default:"https://api.openai.com/v1"
+    api_base: str = os.getenv("OPENAI_BASEURL", "https://api.openai.com/v1")
 
     log_dir: str = "logs"
 

pentestgpt/test_connection.py

@@ -30,7 +30,7 @@ def main():
     parser.add_argument(
         "--baseUrl",
         type=str,
-        default="https://api.openai.com/v1",
+        default=ChatGPTConfig().api_base,
         help="Base URL for OpenAI API,default: https://api.openai.com/v1",
     )