gpt

#!/usr/bin/env python

import ast
import openai
import sys
import argparse
import os
import configparser
import subprocess
from tempfile import NamedTemporaryFile
import time

home_dir = os.path.expanduser("~")
dotfile_path = os.path.join(home_dir, '.openai_config')
gpt_dir = os.path.join(home_dir, '.gpt')

config = configparser.ConfigParser()
if os.path.exists(dotfile_path):
    config.read(dotfile_path)
else:
    print("config file required at ~/.openai_config")
    exit(1)

openai.api_key = config['DEFAULT']['API_KEY']

commands = {
    "git_commit": ['git', 'status', '-v'],
    # Add more commands here
}

def read_prompts():
    prompts_file = os.path.join(home_dir, 'prompts.example')
    with open(prompts_file, "r") as file:
        content = file.read()
        prompts_dict = ast.literal_eval(content)
    return prompts_dict

prompts = read_prompts()

def get_input_or_file_input():
    parser = argparse.ArgumentParser()
    group = parser.add_mutually_exclusive_group()
    group.add_argument("--file", nargs='+', help="Provide one or multiple filenames to use as input")
    group.add_argument("--command-output", choices=list(commands.keys()), help="Choose a command to run and use its output")
    parser.add_argument("--prompt", type=int, help="Select a prompt by index")
    parser.add_argument("--model", help="model to use for the chat completion", default=config['DEFAULT']['MODEL'])
    parser.add_argument("--temperature", help="temperature to use for the chat completion", default=config['DEFAULT']['TEMPERATURE'], type=float)
    parser.add_argument("--confirm-send", choices=['y', 'n'], help="Automatically confirm sending the message without prompt")
    parser.add_argument("--model-selection", choices=['3', '4'], help="Choose the GPT model to use without prompt")
    parser.add_argument("--list-saved", action="store_true", help="List all saved input-output files and exit")
    parser.add_argument("--output", help="Provide a custom path for saving the output")
    args = parser.parse_args()

    if args.list_saved:
        selected_file = list_saved_files()
        if selected_file is not None:
            args.file = [os.path.join(gpt_dir, selected_file)]  # set the selected file as the --file argument
        else:
            sys.exit()

    if args.file:
        content = ""
        for file_name in args.file:
            with open(file_name, "r") as file:
                file_content = file.read()
                content += f"### {file_name}\n\n```\n{file_content}\n```\n\n"

    
    elif args.command_output:
        content = subprocess.check_output(commands[args.command_output]).decode()
    else:
        content = ""

    return content, args



def list_saved_files():
    if not os.path.exists(gpt_dir):
        print("No saved files found.")
        return

    saved_files = os.listdir(gpt_dir)
    input_files = [f for f in saved_files if f.endswith("_input.txt")]
    input_files.sort()

    if not input_files:
        print("No saved input files found.")
        return

    print("Saved files:")
    for i, input_file in enumerate(input_files, start=1):
        output_file = input_file.replace("_input.txt", "_output.txt")
        print(f"{i}. {input_file} -> {output_file}")

    selection = int(input("Enter the number of your selection: ")) - 1
    if selection < 0 or selection >= len(input_files):
        print("Invalid selection.")
        return

    selected_input_file = input_files[selection]
    return selected_input_file  # returning the selected file name

def select_prompt():
    prompt_names = list(prompts.keys())
    print("Select a prompt:")
    for i, name in enumerate(prompt_names):
        print(f"{i}. {name}")
    selection = int(input("Enter the number of your selection: "))
    return prompts[prompt_names[selection]]

def edit_message_in_vim(message):
    with NamedTemporaryFile(suffix=".tmp") as tf:
        tf.write(message.encode())
        tf.flush()
        subprocess.call(["lvim", tf.name])
        tf.seek(0)
        edited_message = tf.read().decode()
    return edited_message

def save_input_output(input_content, output_content, output_path=None):
    timestamp = time.strftime("%Y%m%d-%H%M%S")
    input_filename = f"{timestamp}_input.txt"

    if output_path is None:  # If no custom output path is provided, use the default
        output_filename = f"{timestamp}_output.txt"
    else:
        output_filename = output_path

    if not os.path.exists(gpt_dir):
        os.makedirs(gpt_dir)

    with open(os.path.join(gpt_dir, input_filename), "w") as input_file:
        input_file.write(input_content)

    # If a custom output path is provided, just write to it
    # Otherwise, write to the default location
    if output_path is not None:
        with open(output_filename, "w") as output_file:
            output_file.write(output_content)
        print(f"Output saved to: {output_filename}")
    else:
        with open(os.path.join(gpt_dir, output_filename), "w") as output_file:
            output_file.write(output_content)
        print(f"Output saved to: {os.path.join(gpt_dir, output_filename)}")

    print(f"Input saved to: {os.path.join(gpt_dir, input_filename)}")


def main():
    content, args = get_input_or_file_input()

    if content is None:
        print("No input provided.")
        sys.exit()

    if args.prompt is not None:
        prompt_names = list(prompts.keys())
        if args.prompt < 0 or args.prompt >= len(prompt_names):
            print("Invalid prompt selection.")
            sys.exit()
        prompt = prompts[prompt_names[args.prompt]]
    else:
        prompt = select_prompt()

    message_content = prompt + content

    print("\nHere is the message to be sent:\n")
    print("\n#######################################################################\n")
    print(message_content)
    print("\n#######################################################################\n")


    while True:  # A loop so the user can choose to edit multiple times.
        if args.confirm_send is None:
            confirm = input("\nDo you want to send this message? [y/n/edit]: ")
        else:
            confirm = args.confirm_send

        if confirm.lower() == 'edit':
            message_content = edit_message_in_vim(message_content)
            continue  # Skip the rest of the loop and start from the beginning.
        elif confirm.lower() == 'n':
            print("Operation cancelled.")
            sys.exit()

        if args.model_selection is None:
            model_choice = input("\nDo you want to use GPT-3 or GPT-4? [3/4]: ")
        else:
            model_choice = args.model_selection

        if model_choice == '4':
            args.model = 'gpt-4'
        elif model_choice == '3':
            args.model = 'gpt-3.5-turbo-16k'
        else:
            print("Invalid model selection. Defaulting to GPT-3.")
            args.model = 'gpt-3.5-turbo-16k'

        chat_completion = openai.ChatCompletion.create(
            model=args.model,
            messages=[{"role": "user", "content": message_content}],
            temperature=args.temperature
        )
        output_content = chat_completion.choices[0].message.content  # type: ignore
        print(output_content)

        save_input_output(message_content, output_content, args.output)

        if args.confirm_send:
            break
        else:
            satisfied = input("\nx to exit, any other key to continue [x]: ")
            if satisfied.lower() == 'x':
                break
    

if __name__ == "__main__":
    main()