ChatGPT cli quickie

ask.py

#!/usr/bin/env python

# Place this file in your path
# Mark it executable (chmod +x ask)
# You will need openai lib installed (pip install openai)
# You will also need to set the OPENAI_API_KEY environment variable

import os
import sys
from openai import OpenAI

API_KEY = os.environ.get("OPENAI_API_KEY")

if API_KEY is None:
    raise ValueError("API_KEY is not set")

client = OpenAI(api_key=API_KEY)


def ask(msg) -> None:
    messages = [
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": msg},
    ]
    stream = client.chat.completions.create(
        model="gpt-3.5-turbo",
        messages=messages,
        stream=True,
    )

    column_width = os.get_terminal_size().columns
    line_length = 0
    for chunk in stream:
        content = chunk.choices[0].delta.content
        if content is not None:
            words = content.split()
            for word in words:
                if (line_length + len(word)) > column_width - 1:
                    line_length = 0
                    print()
                print(word, end=" ")
                line_length += len(word) + 1
    print("\n")


if __name__ == "__main__":
    if len(sys.argv) < 2:
        print('Usage: ask "Your question in quotes."')
        exit()
    params = sys.argv[1:]
    ask(" ".join(params))