prompt to run command by default

yolo.py

@@ -1,5 +1,9 @@
 #!/usr/bin/env python3
 
+# MIT License
+# Copyright (c) 2023 wunderwuzzi23
+# Greetings from Seattle! 
+
 import os
 import openai
 import sys
@@ -7,21 +11,31 @@ import subprocess
 from termcolor import colored
 from colorama import init
 
-
-ask = False             # safety switch -a
-yolo = ""               # user's answer to safety switch
-command_start_idx = 1   # command starts at which argv index?
+yolo_safety_switch = True  # by default, user is prompted to run each command
+ask_flag = False           # safety switch -a command line argument
+yolo = ""                  # user's answer to safety switch (-a) question y/n
+command_start_idx = 1      # command starts at which argv index?
 home_path = os.path.expanduser("~")
 
 # Two options for the user to specify they openai api key
 openai.api_key = os.getenv("OPENAI_API_KEY")
-home_path = os.path.expanduser("~")
-openai.api_key_path = home_path+"/.openai.apikey"
+openai.api_key_path = os.path.join(home_path,".openai.apikey")
 
+# Check if the user globally disabled the safety switch
+yolo_safety_off_path = os.path.join(home_path,".yolo-safety-off")
+if os.path.exists(yolo_safety_off_path):
+  yolo_safety_switch = False
+else:
+  yolo_safety_switch = True
 
-# parsing weirdness
+# Parsing weirdness
 if len(sys.argv) < 2:
-  print("Usage Example: yolo [-a] list the current directory information")
+  print("Yolo 0.1 - by @wunderwuzzi23")
+  print()
+  print("Usage: yolo [-a] list the current directory information")
+  print("Argument: -a: Prompt the user before running the command")
+  print()
+  print("Current safety switch setting (~/.yolo-safety-off) is " + str(yolo_safety_switch))
   sys.exit(-1)
 
 # safety switch (no yolo mode)
@@ -45,14 +59,22 @@ if user_prompt == "":
 shell = os.environ.get("SHELL", "powershell.exe")    
 
 # Construct the prompt
-pre_prompt = open("prompt.txt","r").read()
+
+## Find the executing directory (e.g. in case an alias is set)
+## So we can find the prompt.txt file
+yolo_path = os.path.abspath(__file__)
+prompt_path = os.path.dirname(yolo_path)
+
+## Load the prompt and prep it
+prompt_file = os.path.join(prompt_path, "prompt.txt")
+pre_prompt = open(prompt_file,"r").read()
 pre_prompt = pre_prompt.replace("{}", shell)
 prompt = pre_prompt + user_prompt
 
-#make the first line also the system prompt
+## make the first line also the system prompt
 system_prompt = pre_prompt[1]
 
-#be nice and make it a question
+# be nice and make it a question
 if prompt[-1:] != "?":
   prompt+="?"
 
@@ -80,13 +102,13 @@ if res_command.startswith("Sorry, try again") or res_command.startswith("I'm sor
   print(colored("There was an issue: "+res_command, 'red'))
   sys.exit(-1)
 
-#odd corner case, sometimes ChatCompletion returns markdown
+# odd corner case, sometimes ChatCompletion returns markdown
 if res_command.count("```",2):
   print(colored("The proposed command contains markdown, so I thought to not execute the response directly: \n", 'red')+res_command)
   sys.exit(-1)
 
 print("Command: " + colored(res_command, 'blue'))
-if ask == True:
+if yolo_safety_switch == True or ask_flag == True:
   print("Execute the command? Y/n ==> ", end = '')
   yolo = input()
   print()