weyne/PentestGPT
1
0
Fork 0
mirror of https://github.com/weyne85/PentestGPT.git synced 2025-10-29 16:58:59 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: 🐛 minor fix on api usage

Browse Source
This commit is contained in:
Grey_D 2023-04-30 23:08:46 +08:00
parent 3054c3e800
commit 152c226208
3 changed files with 16 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
config/chatgpt_config_sample.py
View File

@ -3,7 +3,10 @@ import dataclasses
@dataclasses.dataclass @dataclasses.dataclass
class ChatGPTConfig: 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" model: str = "text-davinci-002-render-sha"
# set up the openai key # set up the openai key
openai_key = "" openai_key = ""
# set the user-agent below # set the user-agent below

2
main.py
View File

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

22
utils/chatgpt_api.py
View File

@ -41,14 +41,6 @@ class Conversation:
return self.conversation_id == other.conversation_id 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: class ChatGPTAPI:
def __init__(self, config: ChatGPTConfig): def __init__(self, config: ChatGPTConfig):
self.config = config self.config = config
@ -56,6 +48,16 @@ class ChatGPTAPI:
self.history_length = 3 # maintain 3 messages in the history. (3 chat memory) self.history_length = 3 # maintain 3 messages in the history. (3 chat memory)
self.conversation_dict: Dict[str, Conversation] = {} 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): def send_new_message(self, message):
# create a message # create a message
start_time = time.time() start_time = time.time()
@ -65,7 +67,7 @@ class ChatGPTAPI:
message.ask_id = str(uuid1()) message.ask_id = str(uuid1())
message.ask = data message.ask = data
message.request_start_timestamp = start_time message.request_start_timestamp = start_time
response = chatgpt_completion(history) response = self.chatgpt_completion(history)
message.answer = response message.answer = response
message.request_end_timestamp = time.time() message.request_end_timestamp = time.time()
message.time_escaped = ( message.time_escaped = (
@ -100,7 +102,7 @@ class ChatGPTAPI:
message.ask_id = str(uuid1()) message.ask_id = str(uuid1())
message.ask = data message.ask = data
message.request_start_timestamp = time.time() message.request_start_timestamp = time.time()
response = chatgpt_completion(chat_message) response = self.chatgpt_completion(chat_message)
# update the conversation # update the conversation
message.answer = response message.answer = response