From 22f42018be1b5c7a14cca64ed8d9b126ee6ce30a Mon Sep 17 00:00:00 2001 From: wunderwuzzi23 Date: Sun, 5 Mar 2023 15:56:22 -0800 Subject: [PATCH] add os info to prompt --- prompt.txt | 6 +- yolo.py | 166 ++++++++++++++++++++++++++++++++--------------------- 2 files changed, 103 insertions(+), 69 deletions(-) diff --git a/prompt.txt b/prompt.txt index c7a58d2..93787f4 100644 --- a/prompt.txt +++ b/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 diff --git a/yolo.py b/yolo.py index 34f59ec..0333f45 100755 --- a/yolo.py +++ b/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") -if os.path.exists(yolo_safety_off_path): - yolo_safety_switch = False -else: - yolo_safety_switch = True +def get_yolo_safety_switch_config(): + + home_path = os.path.expanduser("~") + 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 -# Parsing weirdness -if len(sys.argv) < 2: +# 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) -if sys.argv[1] == "-a": - ask_flag = True - command_start_idx = 2 +def get_os_friendly_name(): + + # Get OS Name + os_name = platform.system() + + if os_name == "Linux": + dist_name = platform.freedesktop_os_release()["NAME"] + return "Linux/"+dist_name + elif os_name == "Windows": + return os_name + elif os_name == "Darwin": + return "Darwin/macOS" -# 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) -# do we have a prompt from the user? -if user_prompt == "": - print ("No user prompt specified.") +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 -# Unix based SHELL (/bin/bash, /bin/zsh), otherwise assuming it's Windows -shell = os.environ.get("SHELL", "powershell.exe") + # safety switch via argument -a (local override of global setting) + # Force Y/n questions before running the command + 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) -## So we can find the prompt.txt file -yolo_path = os.path.abspath(__file__) -prompt_path = os.path.dirname(yolo_path) + # do we have a prompt from the user? + if user_prompt == "": + print ("No user prompt specified.") + sys.exit(-1) + + + # Load the correct prompt based on Shell and OS and append the user's prompt + prompt = get_full_prompt(user_prompt, shell) -## 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 + system_prompt = prompt[1] -## make the first line also the system prompt -system_prompt = pre_prompt[1] - -# be nice and make it a question -if prompt[-1:] != "?": - prompt+="?" - -response = openai.ChatCompletion.create( - model="gpt-3.5-turbo", - messages=[ - {"role": "system", "content": system_prompt}, - {"role": "user", "content": prompt} - ], - temperature=0, - max_tokens=100, - # top_p=1, - # frequency_penalty=0.2, - # presence_penalty=0 -) + # Call the ChatGPT API + response = openai.ChatCompletion.create( + model="gpt-3.5-turbo", + messages=[ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": prompt} + ], + temperature=0, + max_tokens=100, + # top_p=1, + # 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')) @@ -118,5 +153,4 @@ if yolo == "Y" or yolo == "": subprocess.run([shell, "/c", res_command], shell=False) 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) - + subprocess.run([shell, "-c", res_command], shell=False) \ No newline at end of file