Fixed bad indentation
Please enter the commit message for your changes. Lines starting
This commit is contained in:
Heiko Joerg Schick
@@ -18,199 +18,193 @@ from termcolor import colored
|
|||||||
from colorama import init
|
from colorama import init
|
||||||
|
|
||||||
def read_config() -> any:
|
def read_config() -> any:
|
||||||
|
## Find the executing directory (e.g. in case an alias is set)
|
||||||
|
## So we can find the config file
|
||||||
|
yolo_path = os.path.abspath(__file__)
|
||||||
|
prompt_path = os.path.dirname(yolo_path)
|
||||||
|
|
||||||
## Find the executing directory (e.g. in case an alias is set)
|
config_file = os.path.join(prompt_path, "yolo.yaml")
|
||||||
## So we can find the config file
|
with open(config_file, 'r') as file:
|
||||||
yolo_path = os.path.abspath(__file__)
|
return yaml.safe_load(file)
|
||||||
prompt_path = os.path.dirname(yolo_path)
|
|
||||||
|
|
||||||
config_file = os.path.join(prompt_path, "yolo.yaml")
|
|
||||||
with open(config_file, 'r') as file:
|
|
||||||
return yaml.safe_load(file)
|
|
||||||
|
|
||||||
# Construct the prompt
|
# Construct the prompt
|
||||||
def get_full_prompt(user_prompt, shell):
|
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)
|
||||||
|
|
||||||
## Find the executing directory (e.g. in case an alias is set)
|
## Load the prompt and prep it
|
||||||
## So we can find the prompt.txt file
|
prompt_file = os.path.join(prompt_path, "prompt.txt")
|
||||||
yolo_path = os.path.abspath(__file__)
|
pre_prompt = open(prompt_file,"r").read()
|
||||||
prompt_path = os.path.dirname(yolo_path)
|
pre_prompt = pre_prompt.replace("{shell}", shell)
|
||||||
|
pre_prompt = pre_prompt.replace("{os}", get_os_friendly_name())
|
||||||
|
prompt = pre_prompt + user_prompt
|
||||||
|
|
||||||
## Load the prompt and prep it
|
# be nice and make it a question
|
||||||
prompt_file = os.path.join(prompt_path, "prompt.txt")
|
if prompt[-1:] != "?" and prompt[-1:] != ".":
|
||||||
pre_prompt = open(prompt_file,"r").read()
|
prompt+="?"
|
||||||
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
|
return prompt
|
||||||
if prompt[-1:] != "?" and prompt[-1:] != ".":
|
|
||||||
prompt+="?"
|
|
||||||
|
|
||||||
return prompt
|
|
||||||
|
|
||||||
def print_usage():
|
def print_usage():
|
||||||
print("Yolo v0.2.1 - by @wunderwuzzi23")
|
print("Yolo v0.2.1 - by @wunderwuzzi23")
|
||||||
print()
|
print()
|
||||||
print("Usage: yolo [-a] list the current directory information")
|
print("Usage: yolo [-a] list the current directory information")
|
||||||
print("Argument: -a: Prompt the user before running the command (only useful when safety is off)")
|
print("Argument: -a: Prompt the user before running the command (only useful when safety is off)")
|
||||||
print()
|
print()
|
||||||
|
|
||||||
print("Current configuration per yolo.yaml:")
|
print("Current configuration per yolo.yaml:")
|
||||||
print("* Model : " + str(config["model"]))
|
print("* Model : " + str(config["model"]))
|
||||||
print("* Temperature : " + str(config["temperature"]))
|
print("* Temperature : " + str(config["temperature"]))
|
||||||
print("* Max. Tokens : " + str(config["max_tokens"]))
|
print("* Max. Tokens : " + str(config["max_tokens"]))
|
||||||
print("* Safety : " + str(bool(config["safety"])))
|
print("* Safety : " + str(bool(config["safety"])))
|
||||||
|
|
||||||
|
|
||||||
def get_os_friendly_name():
|
def get_os_friendly_name():
|
||||||
|
# Get OS Name
|
||||||
|
os_name = platform.system()
|
||||||
|
|
||||||
# Get OS Name
|
if os_name == "Linux":
|
||||||
os_name = platform.system()
|
return "Linux/"+distro.name(pretty=True)
|
||||||
|
elif os_name == "Windows":
|
||||||
if os_name == "Linux":
|
return os_name
|
||||||
return "Linux/"+distro.name(pretty=True)
|
elif os_name == "Darwin":
|
||||||
elif os_name == "Windows":
|
return "Darwin/macOS"
|
||||||
return os_name
|
else:
|
||||||
elif os_name == "Darwin":
|
return os_name
|
||||||
return "Darwin/macOS"
|
|
||||||
else:
|
|
||||||
return os_name
|
|
||||||
|
|
||||||
|
|
||||||
def set_api_key():
|
def set_api_key():
|
||||||
# Two options for the user to specify they openai api key.
|
# Two options for the user to specify they openai api key.
|
||||||
#1. Place a ".env" file in same directory as this with the line:
|
#1. Place a ".env" file in same directory as this with the line:
|
||||||
# OPENAI_API_KEY="<yourkey>"
|
# OPENAI_API_KEY="<yourkey>"
|
||||||
# or do `export OPENAI_API_KEY=<yourkey>` before use
|
# or do `export OPENAI_API_KEY=<yourkey>` before use
|
||||||
dotenv.load_dotenv()
|
dotenv.load_dotenv()
|
||||||
openai.api_key = os.getenv("OPENAI_API_KEY")
|
openai.api_key = os.getenv("OPENAI_API_KEY")
|
||||||
|
|
||||||
#2. Place a ".openai.apikey" in the home directory that holds the line:
|
#2. Place a ".openai.apikey" in the home directory that holds the line:
|
||||||
# <yourkey>
|
# <yourkey>
|
||||||
# Note: This options will likely be removed in the future
|
# Note: This options will likely be removed in the future
|
||||||
if not openai.api_key: #If statement to avoid "invalid filepath" error
|
if not openai.api_key: #If statement to avoid "invalid filepath" error
|
||||||
home_path = os.path.expanduser("~")
|
home_path = os.path.expanduser("~")
|
||||||
openai.api_key_path = os.path.join(home_path,".openai.apikey")
|
openai.api_key_path = os.path.join(home_path,".openai.apikey")
|
||||||
|
|
||||||
#3. Final option is the key might be in the yolo.yaml config file
|
#3. Final option is the key might be in the yolo.yaml config file
|
||||||
# openai_apikey: <yourkey>
|
# openai_apikey: <yourkey>
|
||||||
if not openai.api_key:
|
if not openai.api_key:
|
||||||
openai.api_key = config["openai_api_key"]
|
openai.api_key = config["openai_api_key"]
|
||||||
|
|
||||||
if __name__ == "__main__":
|
if __name__ == "__main__":
|
||||||
|
config = read_config()
|
||||||
|
set_api_key()
|
||||||
|
|
||||||
config = read_config()
|
# Unix based SHELL (/bin/bash, /bin/zsh), otherwise assuming it's Windows
|
||||||
set_api_key()
|
shell = os.environ.get("SHELL", "powershell.exe")
|
||||||
|
|
||||||
# Unix based SHELL (/bin/bash, /bin/zsh), otherwise assuming it's Windows
|
command_start_idx = 1 # Question starts at which argv index?
|
||||||
shell = os.environ.get("SHELL", "powershell.exe")
|
ask_flag = False # safety switch -a command line argument
|
||||||
|
yolo = "" # user's answer to safety switch (-a) question y/n
|
||||||
|
|
||||||
command_start_idx = 1 # Question starts at which argv index?
|
# Parse arguments and make sure we have at least a single word
|
||||||
ask_flag = False # safety switch -a command line argument
|
if len(sys.argv) < 2:
|
||||||
yolo = "" # user's answer to safety switch (-a) question y/n
|
print_usage()
|
||||||
|
sys.exit(-1)
|
||||||
|
|
||||||
# Parse arguments and make sure we have at least a single word
|
# Safety switch via argument -a (local override of global setting)
|
||||||
if len(sys.argv) < 2:
|
# Force Y/n questions before running the command
|
||||||
print_usage()
|
if sys.argv[1] == "-a":
|
||||||
sys.exit(-1)
|
ask_flag = True
|
||||||
|
command_start_idx = 2
|
||||||
|
|
||||||
# Safety switch via argument -a (local override of global setting)
|
# To allow easy/natural use we don't require the input to be a
|
||||||
# Force Y/n questions before running the command
|
# single string. So, the user can just type yolo what is my name?
|
||||||
if sys.argv[1] == "-a":
|
# without having to put the question between ''
|
||||||
ask_flag = True
|
arguments = sys.argv[command_start_idx:]
|
||||||
command_start_idx = 2
|
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)
|
|
||||||
|
|
||||||
def call_open_ai(query):
|
def call_open_ai(query):
|
||||||
# do we have a prompt from the user?
|
# do we have a prompt from the user?
|
||||||
if query == "":
|
if query == "":
|
||||||
print ("No user prompt specified.")
|
print ("No user prompt specified.")
|
||||||
sys.exit(-1)
|
sys.exit(-1)
|
||||||
|
|
||||||
# Load the correct prompt based on Shell and OS and append the user's prompt
|
# Load the correct prompt based on Shell and OS and append the user's prompt
|
||||||
prompt = get_full_prompt(query, shell)
|
prompt = get_full_prompt(query, shell)
|
||||||
|
|
||||||
# Make the first line also the system prompt
|
# Make the first line also the system prompt
|
||||||
system_prompt = prompt[1]
|
system_prompt = prompt[1]
|
||||||
#print(prompt)
|
#print(prompt)
|
||||||
|
|
||||||
# Call the ChatGPT API
|
# Call the ChatGPT API
|
||||||
response = openai.ChatCompletion.create(
|
response = openai.ChatCompletion.create(
|
||||||
model=config["model"],
|
model=config["model"],
|
||||||
messages=[
|
messages=[
|
||||||
{"role": "system", "content": system_prompt},
|
{"role": "system", "content": system_prompt},
|
||||||
{"role": "user", "content": prompt}
|
{"role": "user", "content": prompt}
|
||||||
],
|
],
|
||||||
temperature=config["temperature"],
|
temperature=config["temperature"],
|
||||||
max_tokens=config["max_tokens"],
|
max_tokens=config["max_tokens"],
|
||||||
)
|
)
|
||||||
|
|
||||||
return response.choices[0].message.content.strip()
|
|
||||||
|
|
||||||
|
return response.choices[0].message.content.strip()
|
||||||
|
|
||||||
#Enable color output on Windows using colorama
|
#Enable color output on Windows using colorama
|
||||||
init()
|
init()
|
||||||
|
|
||||||
def check_for_issue(response):
|
def check_for_issue(response):
|
||||||
prefixes = ("sorry", "i'm sorry", "the question is not clear", "i'm", "i am")
|
prefixes = ("sorry", "i'm sorry", "the question is not clear", "i'm", "i am")
|
||||||
if response.lower().startswith(prefixes):
|
if response.lower().startswith(prefixes):
|
||||||
print(colored("There was an issue: "+response, 'red'))
|
print(colored("There was an issue: "+response, 'red'))
|
||||||
sys.exit(-1)
|
sys.exit(-1)
|
||||||
|
|
||||||
def check_for_markdown(response):
|
def check_for_markdown(response):
|
||||||
# odd corner case, sometimes ChatCompletion returns markdown
|
# odd corner case, sometimes ChatCompletion returns markdown
|
||||||
if response.count("```",2):
|
if response.count("```",2):
|
||||||
print(colored("The proposed command contains markdown, so I did not execute the response directly: \n", 'red')+response)
|
print(colored("The proposed command contains markdown, so I did not execute the response directly: \n", 'red')+response)
|
||||||
sys.exit(-1)
|
sys.exit(-1)
|
||||||
|
|
||||||
def missing_posix_display():
|
def missing_posix_display():
|
||||||
display = subprocess.check_output("echo $DISPLAY", shell=True)
|
display = subprocess.check_output("echo $DISPLAY", shell=True)
|
||||||
return display == b'\n'
|
return display == b'\n'
|
||||||
|
|
||||||
def prompt_user_input(response):
|
def prompt_user_input(response):
|
||||||
print("Command: " + colored(response, 'blue'))
|
print("Command: " + colored(response, 'blue'))
|
||||||
#print(config["safety"])
|
#print(config["safety"])
|
||||||
|
|
||||||
if bool(config["safety"]) == True or ask_flag == True:
|
if bool(config["safety"]) == True or ask_flag == True:
|
||||||
prompt_text = "Execute command? [Y]es [n]o [m]odify [c]opy to clipboard ==> "
|
prompt_text = "Execute command? [Y]es [n]o [m]odify [c]opy to clipboard ==> "
|
||||||
if os.name == "posix" and missing_posix_display():
|
if os.name == "posix" and missing_posix_display():
|
||||||
prompt_text = "Execute command? [Y]es [n]o [m]odify ==> "
|
prompt_text = "Execute command? [Y]es [n]o [m]odify ==> "
|
||||||
print(prompt_text, end = '')
|
print(prompt_text, end = '')
|
||||||
user_input = input()
|
user_input = input()
|
||||||
return user_input
|
return user_input
|
||||||
|
|
||||||
if config["safety"] == False:
|
if config["safety"] == False:
|
||||||
return "Y"
|
return "Y"
|
||||||
|
|
||||||
def evaluate_input(user_input, command):
|
def evaluate_input(user_input, command):
|
||||||
if user_input.upper() == "Y" or user_input == "":
|
if user_input.upper() == "Y" or user_input == "":
|
||||||
if shell == "powershell.exe":
|
if shell == "powershell.exe":
|
||||||
subprocess.run([shell, "/c", command], shell=False)
|
subprocess.run([shell, "/c", command], shell=False)
|
||||||
else:
|
else:
|
||||||
# Unix: /bin/bash /bin/zsh: uses -c both Ubuntu and macOS should work, others might not
|
# Unix: /bin/bash /bin/zsh: uses -c both Ubuntu and macOS should work, others might not
|
||||||
subprocess.run([shell, "-c", command], shell=False)
|
subprocess.run([shell, "-c", command], shell=False)
|
||||||
|
|
||||||
if user_input.upper() == "M":
|
if user_input.upper() == "M":
|
||||||
print("Modify prompt: ", end = '')
|
print("Modify prompt: ", end = '')
|
||||||
modded_query = input()
|
modded_query = input()
|
||||||
modded_response = call_open_ai(modded_query)
|
modded_response = call_open_ai(modded_query)
|
||||||
check_for_issue(modded_response)
|
check_for_issue(modded_response)
|
||||||
check_for_markdown(modded_response)
|
check_for_markdown(modded_response)
|
||||||
modded_user_input = prompt_user_input(modded_response)
|
modded_user_input = prompt_user_input(modded_response)
|
||||||
print()
|
print()
|
||||||
evaluate_input(modded_user_input, modded_response)
|
evaluate_input(modded_user_input, modded_response)
|
||||||
|
|
||||||
if user_input.upper() == "C":
|
if user_input.upper() == "C":
|
||||||
if os.name == "posix" and missing_posix_display():
|
if os.name == "posix" and missing_posix_display():
|
||||||
return
|
return
|
||||||
pyperclip.copy(command)
|
pyperclip.copy(command)
|
||||||
print("Copied command to clipboard.")
|
print("Copied command to clipboard.")
|
||||||
|
|
||||||
res_command = call_open_ai(user_prompt)
|
res_command = call_open_ai(user_prompt)
|
||||||
check_for_issue(res_command)
|
check_for_issue(res_command)
|
||||||
|
|||||||
Reference in New Issue
Block a user