fix: 🐛 minor fix on api usage

config/chatgpt_config_sample.py

@@ -3,7 +3,10 @@ import dataclasses
 
 @dataclasses.dataclass
 class ChatGPTConfig:
+    # if you're using chatGPT (not API), please use "text-davinci-002-render-sha"
+    # if you're using API, you may configure based on your needs
     model: str = "text-davinci-002-render-sha"
+
     # set up the openai key
     openai_key = ""
     # set the user-agent below

main.py

@@ -9,7 +9,7 @@ if __name__ == "__main__":
     pentestGPTHandler = pentestGPT(reasoning_model="gpt-4", useAPI=False)
 
     # you may use this one if you want to use OpenAI API (without GPT-4)
-    # pentestGPTHandler = pentestGPT(reasoning_model="gpt-4", useAPI=True)
+    # pentestGPTHandler = pentestGPT(reasoning_model="gpt-3.5-turbo", useAPI=True)
 
     # you may use this one if you want to use OpenAI API with GPT-4
     # pentestGPTHandler = pentestGPT(reasoning_model="gpt-4", useAPI=True)

utils/chatgpt_api.py

@@ -41,14 +41,6 @@ class Conversation:
         return self.conversation_id == other.conversation_id
 
 
-def chatgpt_completion(history: List) -> str:
-    response = openai.ChatCompletion.create(
-        model="gpt-3.5-turbo",
-        messages=history,
-    )
-    return response["choices"][0]["message"]["content"]
-
-
 class ChatGPTAPI:
     def __init__(self, config: ChatGPTConfig):
         self.config = config
@@ -56,6 +48,16 @@ class ChatGPTAPI:
         self.history_length = 3  # maintain 3 messages in the history. (3 chat memory)
         self.conversation_dict: Dict[str, Conversation] = {}
 
+    def chatgpt_completion(self, history: List, model="gpt-3.5-turbo") -> str:
+        if self.config.model == "gpt-4":
+            model = "gpt-4"
+            # otherwise, just use the default model (because it is cheaper lol)
+        response = openai.ChatCompletion.create(
+            model=model,
+            messages=history,
+        )
+        return response["choices"][0]["message"]["content"]
+
     def send_new_message(self, message):
         # create a message
         start_time = time.time()
@@ -65,7 +67,7 @@ class ChatGPTAPI:
         message.ask_id = str(uuid1())
         message.ask = data
         message.request_start_timestamp = start_time
-        response = chatgpt_completion(history)
+        response = self.chatgpt_completion(history)
         message.answer = response
         message.request_end_timestamp = time.time()
         message.time_escaped = (
@@ -100,7 +102,7 @@ class ChatGPTAPI:
         message.ask_id = str(uuid1())
         message.ask = data
         message.request_start_timestamp = time.time()
-        response = chatgpt_completion(chat_message)
+        response = self.chatgpt_completion(chat_message)
 
         # update the conversation
         message.answer = response