OpenAI powered AI Streaming CLI in just a few lines of code

aicli.py

import os
from datetime import datetime, timezone
from pathlib import Path

import openai
import rich.traceback
from prompt_toolkit import PromptSession
from prompt_toolkit.history import FileHistory
from rich.console import Console
from rich.live import Live
from rich.markdown import Markdown

rich.traceback.install(show_locals=True)

__version__ = "0.2.0"
console = Console(width=120)

console.print(
    f"aicli - OpenAI powered AI CLI v{__version__}", style="green bold", highlight=False
)

try:
    openai.api_key = os.environ["OPENAI_API_KEY"]
except KeyError:
    console.print("You must set the OPENAI_API_KEY environment variable", style="red")
    exit(1)

now_utc = datetime.now(timezone.utc)
setup = f"""
Help the user by responding to their request, the output should be concise and always written in markdown.
The current date and time is {datetime.now()} {now_utc.astimezone().tzinfo.tzname(now_utc)}.
"""
messages = [{"role": "system", "content": setup}]

history = Path().home() / ".openai-prompt-history.txt"
session = PromptSession(history=FileHistory(history))
while True:
    try:
        text = session.prompt("➤ ")
    except (KeyboardInterrupt, EOFError):
        break
    if not text:
        continue
    messages.append({"role": "user", "content": text})
    try:
        response = openai.ChatCompletion.create(
            model="gpt-4", messages=messages, stream=True
        )
    except KeyboardInterrupt:
        break
    complete_message = ""
    markdown = Markdown(complete_message)
    with Live(markdown, refresh_per_second=15, console=console) as live:
        for chunk in response:
            if chunk["choices"][0]["finish_reason"] is not None:
                break
            chunk_text = chunk["choices"][0]["delta"].get("content", "")
            complete_message += chunk_text
            live.update(Markdown(complete_message))
    messages.append({"role": "assistant", "content": complete_message})