diff --git a/debugger_no_modal.py b/debugger_no_modal.py index 7f1da7c86..6880648bc 100644 --- a/debugger_no_modal.py +++ b/debugger_no_modal.py @@ -1,7 +1,8 @@ import sys import os +from time import sleep from constants import DEFAULT_DIR, DEFAULT_MODEL, DEFAULT_MAX_TOKENS, EXTENSION_TO_SKIP - +import argparse def read_file(filename): with open(filename, "r") as file: return file.read() @@ -37,7 +38,10 @@ def walk_directory(directory): return code_contents -def main(prompt, directory=DEFAULT_DIR, model=DEFAULT_MODEL): +def main(args): + prompt=args.prompt + directory= args.directory + model=args.model code_contents = walk_directory(directory) # Now, `code_contents` is a dictionary that contains the content of all your non-image files @@ -86,7 +90,16 @@ def generate_response(system_prompt, user_prompt, model=DEFAULT_MODEL, *args): } # Send the API request - response = openai.ChatCompletion.create(**params) + keep_trying = True + while keep_trying: + try: + response = openai.ChatCompletion.create(**params) + keep_trying = False + except Exception as e: + # e.g. when the API is too busy, we don't want to fail everything + print("Failed to generate response. Error: ", e) + sleep(30) + print("Retrying...") # Get the reply from the API response reply = response.choices[0]["message"]["content"] @@ -94,9 +107,23 @@ def generate_response(system_prompt, user_prompt, model=DEFAULT_MODEL, *args): if __name__ == "__main__": - if len(sys.argv) < 2: - print("Please provide a prompt") - sys.exit(1) - prompt = sys.argv[1] - model = sys.argv[2] if len(sys.argv) > 2 else DEFAULT_MODEL - main(prompt, model) + parser = argparse.ArgumentParser() + parser.add_argument( + "prompt", + help="The prompt to use for the AI. This should be the error message or issue you are facing.", + + ) + parser.add_argument( + "--directory", + "-d", + help="The directory to use for the AI. This should be the directory containing the files you want to debug.", + default=DEFAULT_DIR, + ) + parser.add_argument( + "--model", + "-m", + help="The model to use for the AI. This should be the model ID of the model you want to use.", + default=DEFAULT_MODEL, + ) + args = parser.parse_args() + main(args)