Merge pull request #1 from wunderwuzzi23/wunderwuzzi23/refactor
add os info to prompt
This commit is contained in:
wunderwuzzi23
GitHub
+3
-3
@@ -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
|
||||
|
||||
@@ -5,80 +5,115 @@
|
||||
# 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():
|
||||
|
||||
# Parsing weirdness
|
||||
if len(sys.argv) < 2:
|
||||
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
|
||||
|
||||
# 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))
|
||||
|
||||
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"
|
||||
|
||||
|
||||
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)
|
||||
|
||||
# safety switch (no yolo mode)
|
||||
if sys.argv[1] == "-a":
|
||||
# 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
|
||||
|
||||
# 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)
|
||||
# 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 == "":
|
||||
# do we have a prompt from the user?
|
||||
if user_prompt == "":
|
||||
print ("No user prompt specified.")
|
||||
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")
|
||||
|
||||
# Construct the prompt
|
||||
# Load the correct prompt based on Shell and OS and append the user's prompt
|
||||
prompt = 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)
|
||||
# Make the first line also the system prompt
|
||||
system_prompt = prompt[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
|
||||
system_prompt = pre_prompt[1]
|
||||
|
||||
# be nice and make it a question
|
||||
if prompt[-1:] != "?":
|
||||
prompt+="?"
|
||||
|
||||
response = openai.ChatCompletion.create(
|
||||
# Call the ChatGPT API
|
||||
response = openai.ChatCompletion.create(
|
||||
model="gpt-3.5-turbo",
|
||||
messages=[
|
||||
{"role": "system", "content": system_prompt},
|
||||
@@ -89,13 +124,13 @@ response = openai.ChatCompletion.create(
|
||||
# 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'))
|
||||
@@ -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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user