yolo-ai-cmdbot/yolo.py

"""
TODO: Description

Sources:
— https://github.com/wunderwuzzi23/yolo-ai-cmdbot
"""

import os
import platform
import subprocess
import sys


import argparse
import distro
import dotenv
import openai
import pyperclip
import yaml

from termcolor import colored
from colorama import init

CONFIG_FILE = "yolo.yaml"

def read_yaml_config() -> any:
    """
    Read the configuration file from the executing directory.

    This function determines the execution folder (which may vary if an alias is set) in order to
    find the configuration file. It reads the file and returns its content in a Python data
    structure.

    Returns:
        The content of the configuration file. Could be dictionary, list, etc. depending on
        the YAML file structure.
    """
    yolo_path = os.path.abspath(__file__)
    prompt_path = os.path.dirname(yolo_path)

    config_file = os.path.join(prompt_path, CONFIG_FILE)
    with open(config_file, 'r') as file:
        return yaml.safe_load(file)

def set_openai_api_key(config):
    """
    Set the OpenAI API key by attempting several methods.

    This function first tries to grab the OpenAI API key from environment variables,
    if not found, it then looks for the key in the `.openai.apikey` in the home directory,
    and lastly, it will look in the provided config dictionary. It sets the `openai.api_key`
    with the retrieved key.

    Parameters:
        config (dict): A dictionary containing configuration values.
                       It may contain `openai_api_key` as one of the keys.
    """
    dotenv.load_dotenv()

    # Method 1: Read API key from environment variable
    # The user can set their OpenAI API key by creating a ".env" file in the same
    # directory as this script or by exporting it to their environment variables.
    # The file or environment variable should contain the line `OPENAI_API_KEY="<yourkey>"`.
    config["openai_api_key"] = os.getenv("OPENAI_API_KEY")

    # Method 2: Read API key from a file in the home directory
    # The user can also place a file named ".openai.apikey" in their home directory,
    # which includes the API key in raw format. This method might be deprecated in future versions.
    if not openai.api_key:  # Check this to avoid potential "invalid filepath" error.
        home_path = os.path.expanduser("~")
        openai.api_key_path = os.path.join(home_path, ".openai.apikey")

    # Method 3: Read API key from the provided config dictionary
    # The final method to set the API key is by providing it in the 'config' dictionary under the
    # key 'openai_api_key'. For instance, in a `yolo.yaml` config file, it would appear as
    # `openai_apikey: <yourkey>`.
    if not openai.api_key:
        openai.api_key = config["openai_api_key"]

# Construct the prompt
def get_full_prompt(user_prompt, shell):
    """
    Construct the prompt, finding the executing directiory so we can find the prompt.txty 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_config(config):
    """
    Print config information.
    """
    print("Current configuration per yolo.yaml:")
    print("* Model        : " + str(config["model"]))
    print("* Temperature  : " + str(config["temperature"]))
    print("* Max. Tokens  : " + str(config["max_tokens"]))
    print("* Safety       : " + str(bool(config["safety"])))


def get_os_friendly_name():
    """
    Get the name of the operating system.
    """
    os_name = platform.system()

    if os_name == "Linux":
        return "Linux/"+distro.name(pretty=True)
    elif os_name == "Windows":
        return os_name
    elif os_name == "Darwin":
        return "Darwin/macOS"
    else:
        return os_name

def call_open_ai(config, shell, query):
    """
    Do we have a prompt from the user?
    """
    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)

    # Make the first line also the system prompt
    system_prompt = prompt[1]
    #print(prompt)

    # Call the ChatGPT API
    response = openai.ChatCompletion.create(
        model=config["model"],
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": prompt}
        ],
        temperature=config["temperature"],
        max_tokens=config["max_tokens"],
    )

    return response.choices[0].message.content.strip()

def check_for_issue(response):
    """
    Todo
    """
    prefixes = ("sorry", "i'm sorry", "the question is not clear", "i'm", "i am")
    if response.lower().startswith(prefixes):
        print(colored("There was an issue: "+response, 'red'))
        sys.exit(-1)

def check_for_markdown(response):
    """
    Odd corner case, sometimes ChatCompletion returns markdown
    """
    if response.count("```",2):
        print(colored("The proposed command contains markdown, so I did not execute the response directly: \n", 'red')+response)
        sys.exit(-1)

def missing_posix_display():
    """
    Todo
    """
    display = subprocess.check_output("echo $DISPLAY", shell=True)
    return display == b'\n'

def prompt_user_input(config, response):
    """
    Todo
    """
    print("Command: " + colored(response, 'blue'))
    #print(config["safety"])

    if bool(config["safety"]) == True or ask_flag == True:
        prompt_text = "Execute command? [Y]es [n]o [m]odify [c]opy to clipboard ==> "
        if os.name == "posix" and missing_posix_display():
            prompt_text =  "Execute command? [Y]es [n]o [m]odify ==> "
        print(prompt_text, end = '')
        user_input = input()
        return user_input

    if config["safety"] == False:
        return "Y"

def evaluate_input(user_input, command):
    """
    Todo
    """
    if user_input.upper() == "Y" or user_input == "":
        if shell == "powershell.exe":
            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()
        modded_response = call_open_ai(modded_query)
        check_for_issue(modded_response)
        check_for_markdown(modded_response)
        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)
        print("Copied command to clipboard.")


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)

    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")

    #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)
    user_input = prompt_user_input(config, res_command)
    print()
    evaluate_input(user_input, res_command)

if __name__ == "__main__":
    main()