diff --git a/config/chatgpt_config_sample.py b/config/chatgpt_config_sample.py index a25c07a..fd4ca55 100644 --- a/config/chatgpt_config_sample.py +++ b/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 diff --git a/main.py b/main.py index a909999..76f8b76 100644 --- a/main.py +++ b/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) diff --git a/utils/chatgpt_api.py b/utils/chatgpt_api.py index 057b7c5..b8774cb 100644 --- a/utils/chatgpt_api.py +++ b/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