A chatgpt terminal parser

chatgpt-terminal.py

import os
import openai
import argparse
import re


def remove_markdown_codeblocks(text):
    pattern = r'```(?:[a-zA-Z]+\n)?(.+?)```|`(.+?)`'
    return re.sub(pattern, lambda match: match.group(1) or match.group(2), text, flags=re.DOTALL)


parser = argparse.ArgumentParser(
    prog="chatgpt-terminal.py",
    description="This is a command line tool that uses the OpenAI GPT-3 API to generate terminal commands.",
)
parser.add_argument(
    "prompt", type=str, help="The prompt to use for the AI to generate a command from."
)
parser.add_argument(
    "--api-key",
    type=str,
    help="The OpenAI API key to use. If not provided, the API key will be read from the OPENAI_API_KEY environment variable.",
    default=os.environ.get("OPENAI_API_KEY"),
)
parser.add_argument(
    "--terminal",
    type=str,
    help="The terminal to use. If not provided, the terminal will be read from the TERMINAL environment variable.",
    default="bash",
)


args = parser.parse_args()

openai.api_key = args.api_key

prompt = args.prompt.strip()


response = openai.ChatCompletion.create(
    model="gpt-3.5-turbo",
    messages=[
        {
            "role": "system",
            "content": f"""
                You are a AI tool that generates terminal commands. You are given a task to generate a terminal command for {args.terminal}. YOU WILL ONLY RESPONS WITH THE COMMAND ITSELF. NO OTHER TEXT IS ALLOWED.
            """,
        },
        {
            "role": "user",
            "content": prompt,
        },
    ],
)

result_command = response["choices"][0]["message"]["content"].strip()
result_command = remove_markdown_codeblocks(result_command)

print("Generated command: ")
print(result_command)
print("Run command? (y/n)")

while True:
    answer = input()
    if answer == "y":
        os.system(result_command)
        break
    elif answer == "n":
        break
    else:
        print("Please answer with y or n.")