get openai to generate shell commands 4 u :)

openai_shell.py

#!/usr/bin/env python3
import sys
import openai
import os
import time
import threading
import itertools
from colorama import Fore as C, Style as S, init as cinit

cinit()
openai.api_key = os.getenv("OPENAI_API_KEY")

EXAMPLES_CONTEXT = "Linux bash command to accomplish the task"

EXAMPLES = [
    ("Get the content of rieqf.txt",
     "cat rieqf.txt"),
    ("Get the last 5 lines of foo.txt",
     "tail -n 5 foo.txt"),
    ("Find files ending in \"log\" in the root directory",
     "find / -name \"*.log\""),
    ("Find \".docx\" files in the directory \"/home/arx/documents\"",
     "find \"/home/arx/documents/\" -name \"*.docx\""),
    ("Look up the IP address 204.48.17.207",
     "nslookup 204.48.17.207"),
    ("Convert g5gaoa4.png to a JPEG",
     "convert g5gaoa4.png g5gao4.jpg"),
    ("Create a git branch foobie-bletch",
     "git checkout -b foobie-bletch")
]

MODEL = "davinci"

prompt = ' '.join(sys.argv[1:]).strip()
if prompt == '':
    print(f'{C.RED}No prompt entered!{C.RESET}')
    exit(1)

done = False
def animate():
    for c in itertools.cycle(['|', '/', '-', '\\']):
        if done:
            break
        sys.stdout.write(f'\r{C.CYAN}[{C.BLUE}{c}{C.CYAN}]{C.YELLOW} Waiting for OpenAI Response...')
        sys.stdout.flush()
        time.sleep(0.1)

t = threading.Thread(target=animate)
t.start()

results = openai.Answer.create(
    model=MODEL,
    question=prompt,
    examples_context=EXAMPLES_CONTEXT,
    examples=EXAMPLES,
    max_tokens=100,
    documents=[]
)
done = True
print()
def cleanup(answer):
    if answer.endswith('\n---'):
        answer = answer[:-4]
    return answer

if results:
    if results["answers"] and len(results["answers"])>0:
        cleaned = [cleanup(answer) for answer in results["answers"]]
        for answer in cleaned:
            print(f"{C.CYAN} \u2022 {S.BRIGHT}{C.GREEN}{answer}{S.RESET_ALL}")
        exit(1)

print(f'{C.RED}Error! No results =({C.RESET}')
exit(1)