added better error handling

moved openai key to a function call for easier calling
pull/25/head
spenceryonce 2023-04-14 09:28:34 -05:00
rodzic 1b9649ed7a
commit d9d43d4b5e
1 zmienionych plików z 60 dodań i 46 usunięć

Wyświetl plik

@ -5,15 +5,15 @@ import os
import shutil import shutil
import subprocess import subprocess
import sys import sys
import openai import openai
from termcolor import cprint from termcolor import cprint
# Set up the OpenAI API # Load the OpenAI API key
with open("openai_key.txt") as f: def load_openai_key():
openai.api_key = f.read().strip() with open("openai_key.txt") as f:
return f.read().strip()
# Run the provided script and return the output and return code
def run_script(script_name, script_args): def run_script(script_name, script_args):
script_args = [str(arg) for arg in script_args] script_args = [str(arg) for arg in script_args]
try: try:
@ -24,8 +24,8 @@ def run_script(script_name, script_args):
return e.output.decode("utf-8"), e.returncode return e.output.decode("utf-8"), e.returncode
return result.decode("utf-8"), 0 return result.decode("utf-8"), 0
# Send the error to GPT and receive suggestions
def send_error_to_gpt(file_path, args, error_message, model): def send_error_to_gpt(file_path, args, error_message, model, prompt_length_limit=4096):
with open(file_path, "r") as f: with open(file_path, "r") as f:
file_lines = f.readlines() file_lines = f.readlines()
@ -50,7 +50,9 @@ def send_error_to_gpt(file_path, args, error_message, model):
"exact format as described above." "exact format as described above."
) )
# print(prompt) # Truncate the prompt if it exceeds the limit
if len(prompt) > prompt_length_limit:
prompt = prompt[:prompt_length_limit]
response = openai.ChatCompletion.create( response = openai.ChatCompletion.create(
model=model, model=model,
@ -65,45 +67,50 @@ def send_error_to_gpt(file_path, args, error_message, model):
return response.choices[0].message.content.strip() return response.choices[0].message.content.strip()
# Apply the changes suggested by GPT
def apply_changes(file_path, changes_json): def apply_changes(file_path, changes_json):
with open(file_path, "r") as f: try:
original_file_lines = f.readlines() with open(file_path, "r") as f:
original_file_lines = f.readlines()
changes = json.loads(changes_json) changes = json.loads(changes_json)
# Filter out explanation elements operation_changes = [change for change in changes if "operation" in change]
operation_changes = [change for change in changes if "operation" in change] explanations = [
explanations = [ change["explanation"] for change in changes if "explanation" in change
change["explanation"] for change in changes if "explanation" in change ]
]
# Sort the changes in reverse line order operation_changes.sort(key=lambda x: x["line"], reverse=True)
operation_changes.sort(key=lambda x: x["line"], reverse=True)
file_lines = original_file_lines.copy() file_lines = original_file_lines.copy()
for change in operation_changes: for change in operation_changes:
operation = change["operation"] operation = change["operation"]
line = change["line"] line = change["line"]
content = change["content"] content = change["content"]
if operation == "Replace": if operation == "Replace":
file_lines[line - 1] = content + "\n" file_lines[line - 1] = content + "\n"
elif operation == "Delete": elif operation == "Delete":
del file_lines[line - 1] del file_lines[line - 1]
elif operation == "InsertAfter": elif operation == "InsertAfter":
file_lines.insert(line, content + "\n") file_lines.insert(line, content + "\n")
with open(file_path, "w") as f: with open(file_path, "w") as f:
f.writelines(file_lines) f.writelines(file_lines)
# Print explanations # Print explanations
cprint("Explanations:", "blue") cprint("Explanations:", "blue")
for explanation in explanations: for explanation in explanations:
cprint(f"- {explanation}", "blue") cprint(f"- {explanation}", "blue")
# Show the diff # Show the diff
print("\nChanges:") print("\nChanges:")
print_diff(original_file_lines, file_lines)
except Exception as e:
raise Exception(f"Failed to apply changes: {str(e)}")
# Print the differences between two file contents
def print_diff(original_file_lines, file_lines):
diff = difflib.unified_diff(original_file_lines, file_lines, lineterm="") diff = difflib.unified_diff(original_file_lines, file_lines, lineterm="")
for line in diff: for line in diff:
if line.startswith("+"): if line.startswith("+"):
@ -113,17 +120,17 @@ def apply_changes(file_path, changes_json):
else: else:
print(line, end="") print(line, end="")
def main(script_name, *script_args, revert=False, model="gpt-4"): def main(script_name, *script_args, revert=False, model="gpt-4"):
openai.api_key = load_openai_key()
if revert: if revert:
backup_file = script_name + ".bak" backup_file = script_name + ".bak"
if os.path.exists(backup_file): if os.path.exists(backup_file):
shutil.copy(backup_file, script_name) shutil.copy(backup_file, script_name)
print(f"Reverted changes to {script_name}") print(f"Reverted changes to {script_name}")
sys.exit(0) return
else: else:
print(f"No backup file found for {script_name}") raise Exception(f"No backup file found for {script_name}")
sys.exit(1)
# Make a backup of the original script # Make a backup of the original script
shutil.copy(script_name, script_name + ".bak") shutil.copy(script_name, script_name + ".bak")
@ -145,9 +152,16 @@ def main(script_name, *script_args, revert=False, model="gpt-4"):
error_message=output, error_message=output,
model=model, model=model,
) )
apply_changes(script_name, json_response) try:
cprint("Changes applied. Rerunning...", "blue") apply_changes(script_name, json_response)
cprint("Changes applied. Rerunning...", "blue")
except Exception as e:
raise Exception(f"Failed to fix the script: {str(e)}")
if __name__ == "__main__": if __name__ == "__main__":
fire.Fire(main) try:
fire.Fire(main)
except Exception as e:
print(str(e))
sys.exit(1)