add os info to prompt

prompt.txt

@@ -1,9 +1,9 @@
-Act as a natural language to {} command translation engine.
+Act as a natural language to {shell} command translation engine on {os}.
 
-You are an expert in {} and translate the question at the end to valid syntax.
+You are an expert in {shell} on {os} and translate the question at the end to valid syntax.
 
 Follow these rules:
-Construct valid {} command that solve the question
+Construct valid {shell} command that solve the question
 Leverage help and man pages to ensure valid syntax and an optimal solution
 Be concise 
 Just show the commands 

yolo.py

@@ -5,97 +5,132 @@
 # Greetings from Seattle! 
 
 import os
+import platform
 import openai
 import sys
 import subprocess
 from termcolor import colored
 from colorama import init
 
-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")
-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")
+def get_yolo_safety_switch_config():
-if os.path.exists(yolo_safety_off_path):
-  yolo_safety_switch = False
-else:
-  yolo_safety_switch = True
   
-# Parsing weirdness
+  home_path = os.path.expanduser("~")
-if len(sys.argv) < 2:
+  yolo_safety_off_path = os.path.join(home_path,".yolo-safety-off")
+  
+  if os.path.exists(yolo_safety_off_path):
+    return False
+  else:
+    return True
+
+# Construct the prompt
+def get_full_prompt(user_prompt, shell):
+
+  ## 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}", shell)
+  pre_prompt = pre_prompt.replace("{os}", get_os_friendly_name())
+  prompt = pre_prompt + user_prompt
+  
+  # be nice and make it a question
+  if prompt[-1:] != "?" and prompt[-1:] != ".":
+    prompt+="?"
+
+  return prompt
+
+def print_usage():
   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)
+def get_os_friendly_name():
-if sys.argv[1] == "-a":
-  ask_flag = True
-  command_start_idx = 2
   
-# to allow easy/natural use we don't require the input to be a 
+  # Get OS Name
-# single string. So, the user can just type yolo what is my name?
+  os_name = platform.system()
-# without having to put the question between ''
-arguments = sys.argv[command_start_idx:]
-user_prompt = " ".join(arguments)
   
-# do we have a prompt from the user?
+  if os_name == "Linux":
-if user_prompt == "":
+      dist_name = platform.freedesktop_os_release()["NAME"]
-    print ("No user prompt specified.")
+      return "Linux/"+dist_name
+  elif os_name == "Windows":
+      return os_name
+  elif os_name == "Darwin":
+     return "Darwin/macOS"
+
+
+if __name__ == "__main__":
+
+  # Get the global safety switch setting (default is True/on) 
+  yolo_safety_switch = get_yolo_safety_switch_config()
+
+  # Unix based SHELL (/bin/bash, /bin/zsh), otherwise assuming it's Windows
+  shell = os.environ.get("SHELL", "powershell.exe") 
+
+  command_start_idx  = 1      # Question starts at which argv index?
+  ask_flag = False           # safety switch -a command line argument
+  yolo = ""                  # user's answer to safety switch (-a) question y/n
+
+  # Two options for the user to specify they openai api key
+  home_path = os.path.expanduser("~")
+  openai.api_key = os.getenv("OPENAI_API_KEY")
+  openai.api_key_path = os.path.join(home_path,".openai.apikey")
+
+  # Parse arguments and make sure we have at least a single word
+  if len(sys.argv) < 2:
+    print_usage()
     sys.exit(-1)
 
-# Get shell info for a better prompt
+  # safety switch via argument -a (local override of global setting)
-# Unix based SHELL (/bin/bash, /bin/zsh), otherwise assuming it's Windows
+  # Force Y/n questions before running the command
-shell = os.environ.get("SHELL", "powershell.exe")    
+  if sys.argv[1] == "-a":
+    ask_flag = True
+    command_start_idx = 2
 
-# Construct the prompt
+  # to allow easy/natural use we don't require the input to be a 
+  # single string. So, the user can just type yolo what is my name?
+  # without having to put the question between ''
+  arguments = sys.argv[command_start_idx:]
+  user_prompt = " ".join(arguments)
 
-## Find the executing directory (e.g. in case an alias is set)
+  # do we have a prompt from the user?
-## So we can find the prompt.txt file
+  if user_prompt == "":
-yolo_path = os.path.abspath(__file__)
+      print ("No user prompt specified.")
-prompt_path = os.path.dirname(yolo_path)
+      sys.exit(-1)
  
-## 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
+  # Load the correct prompt based on Shell and OS and append the user's prompt
-system_prompt = pre_prompt[1]
+  prompt = get_full_prompt(user_prompt, shell)
 
-# be nice and make it a question
+  # Make the first line also the system prompt
-if prompt[-1:] != "?":
+  system_prompt = prompt[1]
-  prompt+="?"
 
-response = openai.ChatCompletion.create(
+  # Call the ChatGPT API
-  model="gpt-3.5-turbo",
+  response = openai.ChatCompletion.create(
-  messages=[
+    model="gpt-3.5-turbo",
-      {"role": "system", "content": system_prompt},
+    messages=[
-      {"role": "user", "content": prompt}
+        {"role": "system", "content": system_prompt},
-  ],
+        {"role": "user", "content": prompt}
-  temperature=0,
+    ],
-  max_tokens=100,
+    temperature=0,
-  # top_p=1,
+    max_tokens=100,
-  # frequency_penalty=0.2,
+    # top_p=1,
-  # presence_penalty=0
+    # frequency_penalty=0.2,
-)
+    # presence_penalty=0
+  )
 
 #print (response)
 
 res_command = response.choices[0].message.content.strip()
 
-#enable color output on Windows using colorama
+#Enable color output on Windows using colorama
 init() 
 
 if res_command.startswith("Sorry, try again") or res_command.startswith("I'm sorry"):
@@ -104,7 +139,7 @@ if res_command.startswith("Sorry, try again") or res_command.startswith("I'm sor
 
 # 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)
+  print(colored("The proposed command contains markdown, so I did not execute the response directly: \n", 'red')+res_command)
   sys.exit(-1)
 
 print("Command: " + colored(res_command, 'blue'))
@@ -119,4 +154,3 @@ if yolo == "Y" or yolo == "":
   else: 
     # Unix: /bin/bash /bin/zsh: uses -c both Ubuntu and macOS should work, others might not
     subprocess.run([shell, "-c", res_command], shell=False)
-