diff --git a/.flake8 b/.flake8 new file mode 100644 index 0000000..6deafc2 --- /dev/null +++ b/.flake8 @@ -0,0 +1,2 @@ +[flake8] +max-line-length = 120 diff --git a/README.md b/README.md index 262a30f..c9e6569 100644 --- a/README.md +++ b/README.md @@ -18,13 +18,19 @@ Add your openAI api key to `openai_key.txt` - _warning!_ by default this uses GP ## Example Usage +To run with gpt-4 (the default, tested option): + python wolverine.py buggy_script.py "subtract" 20 3 +You can also run with other models, but be warned they may not adhere to the edit format as well: + + python wolverine.py --model=gpt-3.5-turbo buggy_script.py "subtract" 20 3 + ## Future Plans This is just a quick prototype I threw together in a few hours. There are many possible extensions and contributions are welcome: -- add flags to customize usage, such as using GPT3.5-turbo instead or asking for user confirmation before running changed code +- add flags to customize usage, such as asking for user confirmation before running changed code - further iterations on the edit format that GPT responds in. Currently it struggles a bit with indentation, but I'm sure that can be improved - a suite of example buggy files that we can test prompts on to ensure reliablity and measure improvement - multiple files / codebases: send GPT everything that appears in the stacktrace diff --git a/requirements.txt b/requirements.txt index 8d1c38c..9d61d58 100644 --- a/requirements.txt +++ b/requirements.txt @@ -5,10 +5,14 @@ attrs==22.2.0 certifi==2022.12.7 charset-normalizer==3.1.0 fire==0.5.0 +flake8==6.0.0 frozenlist==1.3.3 idna==3.4 +mccabe==0.7.0 multidict==6.0.4 openai==0.27.2 +pycodestyle==2.10.0 +pyflakes==3.0.1 requests==2.28.2 six==1.16.0 termcolor==2.2.0 diff --git a/wolverine.py b/wolverine.py index 12345cb..ece00dc 100644 --- a/wolverine.py +++ b/wolverine.py @@ -1,5 +1,6 @@ import argparse import difflib +import fire import json import os import shutil @@ -14,17 +15,18 @@ with open("openai_key.txt") as f: openai.api_key = f.read().strip() -def run_script(script_name, *args): +def run_script(script_name, script_args): + script_args = [str(arg) for arg in script_args] try: result = subprocess.check_output( - [sys.executable, script_name, *args], stderr=subprocess.STDOUT + [sys.executable, script_name, *script_args], stderr=subprocess.STDOUT ) except subprocess.CalledProcessError as e: return e.output.decode("utf-8"), e.returncode return result.decode("utf-8"), 0 -def send_error_to_gpt(file_path, args, error_message, model_name): +def send_error_to_gpt(file_path, args, error_message, model): with open(file_path, "r") as f: file_lines = f.readlines() @@ -52,7 +54,7 @@ def send_error_to_gpt(file_path, args, error_message, model_name): # print(prompt) response = openai.ChatCompletion.create( - model=model_name, + model=model, messages=[ { "role": "user", @@ -113,20 +115,8 @@ def apply_changes(file_path, changes_json): print(line, end="") -def main(): - parser = argparse.ArgumentParser(description="A script to fix Python code using GPT.") - parser.add_argument("script_name", help="The name of the script to fix.") - parser.add_argument("args", nargs="*", help="The arguments for the script.") - parser.add_argument("--model", default="gpt-4", help="The model to use (default: gpt-4).") - parser.add_argument("--revert", action="store_true", help="Revert changes to the script.") - - args = parser.parse_args() - - script_name = args.script_name - script_args = args.args - - # Revert changes if requested - if args.revert: +def main(script_name, *script_args, revert=False, model="gpt-4"): + if revert: backup_file = script_name + ".bak" if os.path.exists(backup_file): shutil.copy(backup_file, script_name) @@ -140,7 +130,7 @@ def main(): shutil.copy(script_name, script_name + ".bak") while True: - output, returncode = run_script(script_name, *script_args) + output, returncode = run_script(script_name, script_args) if returncode == 0: cprint("Script ran successfully.", "blue") @@ -150,10 +140,15 @@ def main(): cprint("Script crashed. Trying to fix...", "blue") print("Output:", output) - json_response = send_error_to_gpt(script_name, script_args, output, args.model) + json_response = send_error_to_gpt( + file_path=script_name, + args=script_args, + error_message=output, + model=model, + ) apply_changes(script_name, json_response) cprint("Changes applied. Rerunning...", "blue") if __name__ == "__main__": - main() + fire.Fire(main)