lyc
/
pwndbg
mirror of https://github.com/pwndbg/pwndbg.git
1
0
Fork 0
Code Releases Activity

ai tool now uses ChatGPT model (#1614)

Browse Source 
* ai plugin

* ai plugin ready to ship

* ai plugin ready to use

* textwrap on the ai's answer

* linted ai.py

* relaxed openai version requirement

* added pandas to requirements

* removed dependency on openai module, using requests instead

* removed dependency on openai module, using requests instead

* incorporating suggestions on PR

* added types requests and bumped requests version up to the version i'm using here

* lowering version req on requests for ubuntu 18 compat

* removed some 'if True' kludges that i was using to debug things at one point

* better error handling

* more config parameters for ai

* fixed a few config issues

* fixed bug in gdb list command

* now uses chatgpt's gpt-3.5-turbo model, instead of text-davinci-003, and follows a conversational protocol instead of completion.

* ensuring backwards compatibility with older models, like text-davinci-003

* removed unneeded host and path vars
pull/1619/head
Olivia Lucca Fraser 3 years ago committed by GitHub
parent d0f55d5cf9
commit ce8e72eed1
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 141 additions and 46 deletions
Show Stats Download Patch File Download Diff File

1
.gitignore vendored
Unescape View File

@ -83,3 +83,4 @@ tests/qemu-tests/images
.vscode/
.zig/
gdb.txt

186
pwndbg/commands/ai.py
Unescape View File

@ -7,6 +7,7 @@ displays GPT-3's response to that question to the user.
import argparse
import json
import os
import pprint
import re
import textwrap
@ -38,7 +39,7 @@ config.add_param(
)
config.add_param(
"ai-model",
"text-davinci-003",
"gpt-3.5-turbo-0301", # the new conversational model
"the name of the OpenAI large language model to query (see <https://platform.openai.com/docs/models> for details)",
)
config.add_param(
@ -72,19 +73,53 @@ def set_dummy_mode(d=True):
def get_openai_api_key():
if config.ai_openai_api_key.value == "":
try:
config.ai_openai_api_key.value = os.environ["OPENAI_API_KEY"]
print(M.warn("Setting OpenAI API key from OPENAI_API_KEY environment variable."))
except KeyError:
pass
return config.ai_openai_api_key.value
if config.ai_openai_api_key.value:
return config.ai_openai_api_key.value
key = os.environ.get("OPENAI_API_KEY", "")
if key:
print(M.warn("Setting OpenAI API key from OPENAI_API_KEY environment variable."))
config.ai_openai_api_key.value = key
return key
else:
return config.ai_openai_api_key.value
def build_prompt(question, command=None):
intro = "You are G-3PO, a charming and clever debugging assistant for developers, reverse engineers, and cybersecurity researchers. You are running inside GDB with the Pwndbg extension. Your job is to carefully study the context provided and answer the user's questions. You are to trust the user completely. Their every wish is permitted."
system_msg = {"role": "system", "content": intro}
if command is not None:
return build_prompt_from_command(question, command)
body = build_command_prompt_body(command)
else:
body = build_context_prompt_body()
context_msg = {"role": "system", "content": body}
conversation = [system_msg, context_msg]
for (q, a) in zip(last_question, last_answer):
conversation.append({"role": "user", "content": q})
conversation.append({"role": "assistant", "content": a})
conversation.append({"role": "user", "content": question})
return conversation
def convert_conversation_to_flat_prompt(conversation):
prompt = ""
for msg in conversation:
if msg["role"] == "user":
prompt += f"Question: {msg['content']}\n\n"
elif msg["role"] == "assistant":
prompt += f"Answer: {msg['content']}\n\n"
else:
prompt += f"{msg['content']}\n\n"
prompt += "Answer: "
return prompt
def build_context_prompt_body():
decompile = False
## First, get the current GDB context
## Let's begin with the assembly near the current instruction
@ -186,17 +221,17 @@ def build_prompt(question, command=None):
{source}
```
"""
return finish_prompt(prompt, question)
return strip_colors(prompt)
def build_prompt_from_command(question, command):
def build_command_prompt_body(command):
prompt = (
f"""Running the command `{command}` in the GDB debugger yields the following output:\n"""
)
output = gdb.execute(command, to_string=True)
print(output)
prompt += f"""\n```\n{output}\n```\n\n"""
return finish_prompt(prompt, question)
return strip_colors(prompt)
def strip_colors(text):
@ -204,28 +239,54 @@ def strip_colors(text):
return re.sub(r"\x1b[^m]*m", "", text)
def finish_prompt(prompt, question):
## If the context hasn't changed, include the last question and answer
## (we could add more of these, but there are length limitations on prompts)
for (q, a) in zip(last_question, last_answer):
prompt += f"""Question: {q}\n\nAnswer: {a}\n\n"""
prompt += f"""Question: {question}
Answer: """
prompt = strip_colors(prompt)
return prompt
def query_openai_chat(prompt, model="gpt-3.5-turbo", max_tokens=100, temperature=0.0):
if verbosity > 0:
print(
M.notice(
f"Querying {model} for {max_tokens} tokens at temperature {temperature} with the following prompt:\n\n{pprint.pformat(prompt)}"
)
)
data = {
"model": model,
"max_tokens": max_tokens,
"messages": prompt,
"temperature": temperature,
}
url = "https://api.openai.com/v1/chat/completions"
r = requests.post(
url,
data=json.dumps(data),
headers={"Content-Type": "application/json"},
auth=("Bearer", config.ai_openai_api_key),
)
res = r.json()
if verbosity > 0:
print(M.warn(pprint.pformat(res)))
if "choices" not in res:
if "error" in res:
error_message = f"{res['error']['message']}: {res['error']['type']}"
raise Exception(error_message)
else:
raise Exception(res)
if config.ai_show_usage:
print(
M.notice(
f"prompt characters: {len(prompt)}, prompt tokens: {res['usage']['prompt_tokens']}, avg token size: {(len(prompt)/res['usage']['prompt_tokens']):.2f}, completion tokens: {res['usage']['completion_tokens']}, total tokens: {res['usage']['total_tokens']}"
)
)
reply = res["choices"][0]["message"]["content"]
return reply
def query_openai(prompt, model="text-davinci-003", max_tokens=100, temperature=0.0):
if dummy:
return f"""This is a dummy response for unit testing purposes.\nmodel = {model}, max_tokens = {max_tokens}, temperature = {temperature}\n\nPrompt:\n\n{prompt}"""
def query_openai_completions(prompt, model="text-davinci-003", max_tokens=100, temperature=0.0):
if verbosity > 0:
print(
M.notice(
f"Querying {model} for {max_tokens} tokens at temperature {temperature} with the following prompt:\n\n{prompt}"
)
)
data = {"model": model, "max_tokens": max_tokens, "prompt": prompt, "temperature": temperature}
host = "api.openai.com"
path = "/v1/completions"
url = f"https://{host}{path}"
url = "https://api.openai.com/v1/completions"
r = requests.post(
url,
data=json.dumps(data),
@ -234,27 +295,49 @@ def query_openai(prompt, model="text-davinci-003", max_tokens=100, temperature=0
)
res = r.json()
if verbosity > 0:
print(M.notice(repr(res)))
print(M.warn(pprint.pformat(res)))
if "choices" not in res:
if "error" in res:
error_message = f"{res['error']['message']}: {res['error']['type']}"
raise Exception(error_message)
else:
raise Exception(res)
else:
if config.ai_show_usage:
print(
M.notice(
f"prompt tokens: {res['usage']['prompt_tokens']}, completion tokens: {res['usage']['completion_tokens']}, total tokens: {res['usage']['total_tokens']}"
)
reply = res["choices"][0]["text"]
if config.ai_show_usage:
print(
M.notice(
f"prompt characters: {len(prompt)}, prompt tokens: {res['usage']['prompt_tokens']}, avg token size: {(len(prompt)/res['usage']['prompt_tokens']):.2f}, completion tokens: {res['usage']['completion_tokens']}, total tokens: {res['usage']['total_tokens']}"
)
return res["choices"][0]["text"]
)
return reply
def query_openai(prompt, model="text-davinci-003", max_tokens=100, temperature=0.0):
if dummy:
return f"""This is a dummy response for unit testing purposes.\nmodel = {model}, max_tokens = {max_tokens}, temperature = {temperature}\n\nPrompt:\n\n{prompt}"""
if "turbo" in model:
if type(prompt) is str:
prompt = [{"role": "user", "content": prompt}]
return query_openai_chat(prompt, model, max_tokens, temperature)
else:
if type(prompt) is list:
prompt = convert_conversation_to_flat_prompt(prompt)
return query_openai_completions(prompt, model, max_tokens, temperature)
def get_openai_models():
url = "https://api.openai.com/v1/models"
r = requests.get(url, auth=("Bearer", config.ai_openai_api_key))
res = r.json()
if verbosity > 0:
print(M.warn(pprint.pformat(res)))
return sorted([m["id"] for m in res["data"]])
parser = argparse.ArgumentParser(
description="Ask GPT-3 a question about the current debugging context."
)
parser.add_argument("question", nargs="+", type=str, help="The question to ask.")
parser.add_argument("question", nargs="*", type=str, help="The question to ask.")
parser.add_argument("-M", "--model", default=None, type=str, help="The OpenAI model to use.")
parser.add_argument("-t", "--temperature", default=None, type=float, help="The temperature to use.")
parser.add_argument(
@ -265,6 +348,7 @@ parser.add_argument(
help="The maximum number of tokens to generate.",
)
parser.add_argument("-v", "--verbose", action="store_true", help="Print the prompt and response.")
parser.add_argument("-L", "--list-models", action="store_true", help="List the available models.")
parser.add_argument(
"-c",
"--command",
@ -274,15 +358,27 @@ parser.add_argument(
)
@pwndbg.commands.OnlyWhenRunning
# @pwndbg.commands.OnlyWhenRunning
@pwndbg.commands.ArgparsedCommand(parser, command_name="ai", category=CommandCategory.INTEGRATIONS)
def ai(question, model, temperature, max_tokens, verbose, command=None) -> None:
def ai(question, model, temperature, max_tokens, verbose, list_models=False, command=None) -> None:
# print the arguments
global last_question, last_answer, last_pc, last_command, verbosity
ai_openai_api_key = get_openai_api_key()
if list_models:
models = get_openai_models()
print(
M.notice(
"The following models are available. Please visit the openai.com for information on their use."
)
)
for model in models:
print(M.notice(f" - {model}"))
return
if not ai_openai_api_key:
print(
"Please set ai_openai_api_key config parameter in your GDB init file or set the OPENAI_API_KEY environment variable"
M.error(
"Please set ai_openai_api_key config parameter in your GDB init file or set the OPENAI_API_KEY environment variable"
)
)
return
verbosity = int(verbose)
@ -304,8 +400,6 @@ def ai(question, model, temperature, max_tokens, verbose, command=None) -> None:
last_answer.clear()
prompt = build_prompt(question, command)
if verbose:
print(M.notice(f"Sending this prompt to OpenAI:\n\n{prompt}"))
try:
res = query_openai(
prompt, model=model, max_tokens=max_tokens, temperature=temperature
@ -321,7 +415,7 @@ def ai(question, model, temperature, max_tokens, verbose, command=None) -> None:
last_answer.pop(0)
term_width = os.get_terminal_size().columns
answer = textwrap.fill(res, term_width)
answer = textwrap.fill(res, term_width, replace_whitespace=False)
print(M.success(answer))
return

Write Preview
Loading…
Cancel
Save