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
def load_openai_key():
with open("openai_key.txt") as f: with open("openai_key.txt") as f:
openai.api_key = f.read().strip() 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,20 +67,19 @@ 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):
try:
with open(file_path, "r") as f: with open(file_path, "r") as f:
original_file_lines = f.readlines() 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()
@ -104,6 +105,12 @@ def apply_changes(file_path, changes_json):
# 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,
) )
try:
apply_changes(script_name, json_response) apply_changes(script_name, json_response)
cprint("Changes applied. Rerunning...", "blue") 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__":
try:
fire.Fire(main) fire.Fire(main)
except Exception as e:
print(str(e))
sys.exit(1)