Simple ChatGPT terminal app

chatgpt.py

#!/usr/bin/python3.8
# Author: Teshan Liyanage <[email protected]>

import readline
import openai

openai.api_key = "<get from https://platform.openai.com/account/api-keys>"

# ANSI colors
c_b, c_g, c_gr, c_nc = "\033[94m", "\033[32m", "\033[30m", "\033[0m"
default_role = "You are a helpful assistant"


def chat_complete(whoami=default_role, model="gpt-3.5-turbo"):
    if whoami is None:
        whoami = input(c_g + f"Who am I? ['{default_role}']\n>> " + c_nc) or default_role

    completion = openai.ChatCompletion()
    chat_log = [{"role": "system", "content": whoami}]
    while (ip := input("\n> ")) != "q":
        if len(ip) < 2:
            continue
        chat_log.append({"role": "user", "content": ip})
        resp = completion.create(model=model, messages=chat_log)
        c = resp.choices[0]
        print("\n" + c_b + c.message.content + c_nc)

        u = resp.usage
        print(c_gr + f"\n{c.finish_reason} | #tokens q:{u.prompt_tokens} a:{u.completion_tokens} tot:{u.total_tokens}" + c_nc)

        chat_log.append({"role": "assistant", "content": c.message.content})


def parse_args():
    from argparse import ArgumentParser

    ap = ArgumentParser("ChatGPT in terminal")
    ap.add_argument('-w', '--whoami', nargs='?', default=default_role, help='role')

    return ap.parse_args()


if __name__ == "__main__":
    args = parse_args()
    try:
        chat_complete(whoami=args.whoami)
    except (KeyboardInterrupt, EOFError):
        print()