Chat with OpenAI's GPT models directly from the command line
"""
llm-term info file
"""

__author__ = "Justin Flannery"
__email__ = "juftin@juftin.com"
__application__ = "llm-term"
__version__ = "0.9.1"

llm_term/__about__.py
juftin/llm-term
"""
llm-term
"""

from llm_term.__about__ import __application__, __author__, __email__, __version__

__all__ = [
    "__application__",
    "__version__",
    "__author__",
    "__email__",
]

llm_term/__init__.py
juftin/llm-term
"""
Command-line interface
"""

from llm_term.cli import cli as main  # pragma: no cover

if __name__ == "__main__":
    main()

llm_term/__main__.py
juftin/llm-term
"""
Base classes and types for llm_term
"""

from typing import Any, TypedDict

from rich.console import RenderableType
from rich.padding import Padding


class NoPadding(Padding):
    """
    Empty Renderable
    """

    def __init__(self, renderable: RenderableType, **kwargs: Any) -> None:
        """
        Create an empty Padding
        """
        _ = kwargs
        pad = (0, 0, 0, 0)
        super().__init__(renderable, pad=pad)


class Message(TypedDict):
    """
    Message
    """

    role: str
    content: str

llm_term/base.py
juftin/llm-term
"""
LLM CLI
"""

from __future__ import annotations

import click
import rich.traceback
from rich.console import Console

from llm_term.__about__ import __application__, __version__
from llm_term.utils import (
    chat_session,
    check_credentials,
    print_header,
    setup_system_message,
)

rich.traceback.install(show_locals=True)


@click.command()
@click.version_option(version=__version__, prog_name=__application__)
@click.option(
    "--model",
    "-m",
    help="The GPT model to use",
    envvar="OPENAI_MODEL",
    show_envvar=True,
    default="gpt-3.5-turbo",
    type=click.STRING,
    show_default=True,
)
@click.argument(
    "chat",
    nargs=-1,
    type=click.STRING,
    default=None,
)
@click.option(
    "--system",
    "-s",
    help="The system message to use",
    envvar="OPENAI_SYSTEM_MESSAGE",
    show_envvar=True,
    default=None,
    type=click.STRING,
)
@click.option(
    "--api-key",
    "-k",
    help="The OpenAI API key",
    envvar="OPENAI_API_KEY",
    show_envvar=True,
    type=click.STRING,
)
@click.option(
    "--stream/--no-stream",
    help="Stream the response",
    envvar="OPENAI_STREAM",
    show_envvar=True,
    default=True,
    type=click.BOOL,
)
@click.option(
    "--console",
    "-c",
    help="The console width to use - defaults to auto-detect",
    type=click.INT,
    default=None,
)
@click.option(
    "--border/--no-border",
    help="Use a border for returned responses - disable for copy/paste",
    default=True,
    type=click.BOOL,
)
@click.option(
    "--avatar",
    "-a",
    help="The avatar to use",
    type=click.STRING,
    default="🤓",
    show_default=True,
)
def cli(
    model: str,
    chat: tuple[str, ...],
    system: str | None,
    api_key: str,
    stream: bool,
    console: int,
    border: bool,
    avatar: str,
) -> None:
    """
    llm-term is a command line interface for OpenAI's Chat API.
    """
    rich_console: Console = Console(width=console)
    chat_message = " ".join(chat)
    try:
        print_header(console=rich_console, model=model)
        check_credentials(api_key=api_key)
        system_message = setup_system_message(model=model, message=system)
        chat_session(
            console=rich_console,
            system_message=system_message,
            model=model,
            stream=stream,
            panel=border,
            chat_message=chat_message,
            avatar=avatar,
        )
    except KeyboardInterrupt as ki:
        raise click.exceptions.Exit(code=0) from ki


if __name__ == "__main__":
    cli()

llm_term/cli.py
juftin/llm-term
"""
Helper functions for the CLI
"""

from __future__ import annotations

from datetime import datetime, timezone
from pathlib import Path
from textwrap import dedent
from typing import Any, Generator

import click
import openai
from prompt_toolkit import PromptSession
from prompt_toolkit.auto_suggest import AutoSuggestFromHistory
from prompt_toolkit.history import FileHistory
from rich.columns import Columns
from rich.console import Console
from rich.live import Live
from rich.markdown import Markdown
from rich.panel import Panel
from rich.spinner import Spinner

from llm_term.__about__ import __application__, __version__
from llm_term.base import Message, NoPadding


def print_header(console: Console, model: str) -> None:
    """
    Print the header
    """
    console.print(
        Panel(
            f"[bold red]{__application__} "
            f"is a command line interface for OpenAI's Chat API[/bold red]",
            title=f"[blue bold]{__application__} v{__version__}[/blue bold]",
            subtitle=f"[yellow bold]{model}[/yellow bold]",
            style="green bold",
            expand=False,
        ),
    )
    console.print("")


def check_credentials(api_key: str) -> None:
    """
    Check the credentials
    """
    if api_key is not None:
        openai.api_key = api_key
    else:
        msg = (
            "You must set the OPENAI_API_KEY environment variable "
            "or set the `--api-key` / `-k` option"
        )
        raise click.ClickException(msg)


def setup_system_message(model: str, message: str | None = None) -> Message:
    """
    Set up the system message
    """
    setup = f"""
        You are a helpful AI assistant named {__application__}, built by OpenAI and based on
        the model {model}. Help the user by responding to their request, the output should
        be concise and always written in markdown. Ensure that all code blocks have the correct
        language tag.

        The current UTC date and time at the start of this conversation is
        {datetime.now(tz=timezone.utc).isoformat()}
        """
    system_message = message or dedent(setup).strip()
    return Message(role="system", content=system_message)


def chat_session(
    console: Console,
    system_message: Message,
    model: str,
    stream: bool,
    panel: bool,
    chat_message: str,
    avatar: str,
) -> None:
    """
    Chat session with ChatGPT
    """
    history_file = Path().home() / ".llm-term-history.txt"
    history = FileHistory(str(history_file))
    session: PromptSession = PromptSession(history=history, erase_when_done=True)
    messages: list[Message] = [system_message]
    if chat_message.strip() != "":
        messages.append(Message(role="user", content=chat_message))
    message_counter = 0
    while True:
        if message_counter == 0 and len(messages) == 2:  # noqa: PLR2004
            console.print(f"{avatar}: {chat_message}")
            history.append_string(chat_message)
            if panel is False:
                console.print("")
                console.rule()
        else:
            text = session.prompt(f"{avatar}: ", auto_suggest=AutoSuggestFromHistory())
            if not text:
                continue
            console.print(f"{avatar}: {text}")
            if panel is False:
                console.print("")
                console.rule()
            messages.append(Message(role="user", content=text))
        console.print("")
        streamed_message = print_response(
            console=console,
            messages=messages,
            model=model,
            stream=stream,
            panel=panel,
        )
        if panel is False:
            console.print("")
            console.rule()
        messages.append(streamed_message)
        console.print("")
        message_counter += 1


def print_response(
    console: Console, messages: list[Message], model: str, stream: bool, panel: bool
) -> Message:
    """
    Stream the response
    """
    panel_class = Panel if panel is True else NoPadding
    if stream is False:
        with Live(Spinner("aesthetic"), refresh_per_second=15, console=console, transient=True):
            response = openai.ChatCompletion.create(model=model, messages=messages, stream=False)
            complete_response = response["choices"][0]["message"]["content"]
            console.print(panel_class(Markdown(complete_response), title="🤖", title_align="left"))
        message = Message(role="assistant", content=complete_response)
    else:
        response = openai.ChatCompletion.create(model=model, messages=messages, stream=True)
        message = render_streamed_response(response=response, console=console, panel=panel)
    return message


def render_streamed_response(
    response: Generator[dict[str, Any], None, None], console: Console, panel: bool
) -> Message:
    """
    Render the streamed response and a spinner
    """
    panel_class = Panel if panel is True else NoPadding
    complete_message = ""
    rich_response = Columns(
        [
            panel_class(Markdown(complete_message), title="🤖", title_align="left"),
            Spinner("aesthetic"),
        ]
    )
    with Live(
        rich_response,
        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
            updated_response = Columns(
                [
                    panel_class(Markdown(complete_message), title="🤖", title_align="left"),
                    Spinner("aesthetic"),
                ]
            )
            live.update(updated_response)
        live.update(panel_class(Markdown(complete_message), title="🤖", title_align="left"))
    return Message(role="assistant", content=complete_message)

llm_term/utils.py
juftin/llm-term
"""
Shared fixtures for tests.
"""

import pytest
from click.testing import CliRunner


@pytest.fixture
def runner() -> CliRunner:
    """
    Fixture for invoking command-line interfaces.
    """
    return CliRunner()

tests/conftest.py
juftin/llm-term
"""
Test cases for the cli.
"""

from click.testing import CliRunner

from llm_term.cli import cli


def test_main_succeeds(runner: CliRunner) -> None:
    """
    It exits with a status code of zero.
    """
    result = runner.invoke(cli, args=["--help"])
    assert result.exit_code == 0

tests/test_cli.py
juftin/llm-term