diff --git a/yolo.py b/yolo.py index 92551a4..67bf8b0 100755 --- a/yolo.py +++ b/yolo.py @@ -7,14 +7,16 @@ Sources: import os import platform -import sys import subprocess +import sys -import openai -import dotenv + +import argparse import distro -import yaml +import dotenv +import openai import pyperclip +import yaml from termcolor import colored from colorama import init @@ -96,16 +98,10 @@ def get_full_prompt(user_prompt, shell): return prompt -def print_usage(): +def print_config(config): """ - Print usage information. + Print config information. """ - print("Yolo v0.2.1 - by @wunderwuzzi23") - print() - 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() - print("Current configuration per yolo.yaml:") print("* Model : " + str(config["model"])) print("* Temperature : " + str(config["temperature"])) @@ -230,38 +226,32 @@ def main(): """ Defined starting point of source code. """ + parser = argparse.ArgumentParser( + description='AI bot that translates your question to a command.' + ) + parser.add_argument('text', nargs='+', + help='A sequence of strings') + parser.add_argument("-s", "--safety", action='store_true', + help='Enable safety mode (only useful when safety is off)') + parser.add_argument("-c", "--config", action='store_true', + help='Print current configuration') + args = parser.parse_args() + + user_prompt = " ".join(args.text) + ask_flag = args.safety + config = read_yaml_config() set_openai_api_key(config) - # set_api_key(config) + if (args.config): + print_config(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 - - # Parse arguments and make sure we have at least a single word - if len(sys.argv) < 2: - print_usage() - sys.exit(-1) - - # 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) - #Enable color output on Windows using colorama init() - + res_command = call_open_ai(config, shell, user_prompt) check_for_issue(res_command) check_for_markdown(res_command) @@ -270,4 +260,4 @@ def main(): evaluate_input(user_input, res_command) if __name__ == "__main__": - main() \ No newline at end of file + main()