diff --git a/yolo.py b/yolo.py index 6911b85..d2a6640 100755 --- a/yolo.py +++ b/yolo.py @@ -2,14 +2,14 @@ # MIT License # Copyright (c) 2023 wunderwuzzi23 -# Greetings from Seattle! +# Greetings from Seattle! import os import platform import openai import sys import subprocess -import dotenv +import dotenv import distro import yaml import pyperclip @@ -42,7 +42,7 @@ def get_full_prompt(user_prompt, shell): 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+="?" @@ -64,10 +64,10 @@ def print_usage(): def get_os_friendly_name(): - + # Get OS Name os_name = platform.system() - + if os_name == "Linux": return "Linux/"+distro.name(pretty=True) elif os_name == "Windows": @@ -85,17 +85,17 @@ def set_api_key(): # or do `export OPENAI_API_KEY=` before use dotenv.load_dotenv() openai.api_key = os.getenv("OPENAI_API_KEY") - + #2. Place a ".openai.apikey" in the home directory that holds the line: # # Note: This options will likely be removed in the future 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") #3. Final option is the key might be in the yolo.yaml config file # openai_apikey: - if not openai.api_key: + if not openai.api_key: openai.api_key = config["openai_api_key"] if __name__ == "__main__": @@ -104,7 +104,7 @@ if __name__ == "__main__": set_api_key() # Unix based SHELL (/bin/bash, /bin/zsh), otherwise assuming it's Windows - shell = os.environ.get("SHELL", "powershell.exe") + 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 @@ -121,7 +121,7 @@ if __name__ == "__main__": ask_flag = True command_start_idx = 2 - # To allow easy/natural use we don't require the input to be a + # 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:] @@ -132,7 +132,7 @@ def call_open_ai(query): if query == "": 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(query, shell) @@ -150,12 +150,12 @@ def call_open_ai(query): temperature=config["temperature"], max_tokens=config["max_tokens"], ) - + return response.choices[0].message.content.strip() #Enable color output on Windows using colorama -init() +init() def check_for_issue(response): prefixes = ("sorry", "i'm sorry", "the question is not clear", "i'm", "i am") @@ -183,19 +183,19 @@ def prompt_user_input(response): prompt_text = "Execute command? [Y]es [n]o [m]odify ==> " print(prompt_text, end = '') user_input = input() - return user_input - + return user_input + if config["safety"] == False: return "Y" def evaluate_input(user_input, command): if user_input.upper() == "Y" or user_input == "": if shell == "powershell.exe": - subprocess.run([shell, "/c", command], shell=False) - else: + subprocess.run([shell, "/c", command], shell=False) + else: # Unix: /bin/bash /bin/zsh: uses -c both Ubuntu and macOS should work, others might not subprocess.run([shell, "-c", command], shell=False) - + if user_input.upper() == "M": print("Modify prompt: ", end = '') modded_query = input() @@ -205,14 +205,14 @@ def evaluate_input(user_input, command): modded_user_input = prompt_user_input(modded_response) print() evaluate_input(modded_user_input, modded_response) - + if user_input.upper() == "C": if os.name == "posix" and missing_posix_display(): return - pyperclip.copy(command) + pyperclip.copy(command) 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_markdown(res_command) user_input = prompt_user_input(res_command)