llm

#!/usr/local/bin/python3
# -*- coding: utf-8 -*-

# Author: Doug Rudolph (https://github.com/11)
# Date: Jan. 5, 2025
# Description: LLM CLI

# Setup:
# You will need create these environment variables for each service you'd like to work with
#   - `OPENAI_API_KEY='<your-openai-key>'`
#   - `ANTHROPIC_API_KEY='<your-anthropic-key>'`
#   - `GEMINI_API_KEY='<your-gemini-key>'`


import os
import sys
import json

from typing import List
from pathlib import Path
from urllib import request, error
from argparse import ArgumentParser, BooleanOptionalAction


#
# UTILITY FUNCTIONS
#


def find_and_replace(prompt: str, files: List[Path] | None) -> str:
    """
    Find and replace variables in prompt with data from file,
    otherwise append uninjected file text to end of prompt
    """
    # Initialize files to empty list if None
    files = files or []

    # Find and replace text from file into prompt
    file_data = [
        (idx, file.stem, file.name, read_file(file))
        for idx, file in enumerate(files)
    ]

    found_placeholders = [False] * len(file_data)
    for (idx, stem, _, content) in file_data:
        name_placeholder = f'${stem.upper()}'
        name_placeholder2 = f'${stem.lower()}'
        index_placeholder = f'${idx+1}'
        found_placeholders[idx] = name_placeholder in prompt or name_placeholder2 in prompt or index_placeholder in prompt

        prompt = prompt.replace(name_placeholder, content.strip())
        prompt = prompt.replace(name_placeholder2, content.strip())
        prompt = prompt.replace(index_placeholder, content.strip())

    # If no placeholders found in prompt, append file data to end of file
    for idx, found in enumerate(found_placeholders):
        if found:
            continue

        name = file_data[idx][2]
        data = file_data[idx][3]
        prompt += f'\n{name}:\n{data}'

    return prompt


def print_inputs(prompt: str, model: str, output: str, tokens: int, entropy: float, role: str, files: List[Path] | None, verbose: bool = False):
    params = locals()
    params['prompt'] = prompt.split('\n')

    print("\nInputs:")
    for arg_name, arg_value in params.items():
        if arg_name == 'prompt':
            print(f'  {arg_name}: {"\n      ".join(arg_value)}')
        else:
            print(f'  {arg_name}: {arg_value}')
    print()


def create_argparser() -> ArgumentParser:
    parser = ArgumentParser(prog='llm', description='LLM CLI')
    subparsers = parser.add_subparsers(dest='command', required=True)

    # env command
    env_cmd = subparsers.add_parser('env', help='List found LLM environment variables')

    env_cmd.add_argument(
        '-V',
        '--verbose',
        type=bool,
        default=False,
        action=BooleanOptionalAction,
        help='Enable verbose output'
    )

    # chat command and its arguments
    prompt_cmd = subparsers.add_parser('chat', help='Send a prompt to an LLM model')

    prompt_cmd.add_argument(
        '-o',
        '--output',
        type=str,
        choices=['json', 'text'],
        default='text',
        help='Output format - defaults to `text`'
    )

    prompt_cmd.add_argument(
        '-t',
        '--tokens',
        type=int,
        default=2048,
        help='Max amount of tokens- defaults to `2048`'
    )

    prompt_cmd.add_argument(
        '-e',
        '--entropy',
        metavar='{0.1...2.0}',
        type=float,
        choices=[.1, .2, .3, .4, .5, .6, .7, .8, .9, 1.0, 1.1, 1.2, 1.3, 1.4, 1.5, 1.6, 1.7, 1.8, 1.9, 2.0],
        default=0.7,
        help='Model\'s  temperature - defaults to `0.7`'
    )

    prompt_cmd.add_argument(
        '-r',
        '--role',
        type=str,
        choices=['user', 'system'],
        default='user',
        help='Model role - defaults to `user`'
    )

    prompt_cmd.add_argument(
        '-V',
        '--verbose',
        type=bool,
        default=False,
        action=BooleanOptionalAction,
        help='Enable verbose output'
    )

    prompt_cmd.add_argument(
        '-f',
        '--files',
        type=Path,
        nargs='*',
        help='Optional file paths to read and append to prompt. Use $1, $2, etc. in prompt to specify file placement - otherwise files are appended in order at the end of the prompt.'
    )

    prompt_cmd.add_argument(
        '-p',
        '--prompt',
        type=str,
        required=True,
        help='Input prompt to send to model'
    )

    prompt_cmd.add_argument(
        '-m',
        '--model',
        type=str,
        metavar='',
        choices=[
            # openai text models
            'gpt-4',
            'gpt-4o-mini',
            'gpt-4o',

            # anthropic text models
            'claude-3-5-sonnet-latest',
            'claude-3-5-haiku-latest',
            'claude-3-opus-latest',

            # gemini text models
            'gemini-1.5-flash',
            'gemini-1.5-flash-8b',
            'gemini-1.5-pro',
        ],
        default='gpt-4o-mini',
        help='Defaults to gpt-4o-mini. OpenAI: gpt-4, gpt-4o-mini, gpt-4o. Anthropic: claude-3-5-sonnet-latest, claude-3-5-haiku-latest, claude-3-opus-latest. Gemini: gemini-1.5-flash, gemini-1.5-flash-8b, gemini-1.5-pro'
    )

    return parser


def fetch(url: str, method: str = 'POST', headers: dict | None = None, body: dict | None = None):
    """ Make a post HTTP request """

    if body:
        body = json.dumps(body).encode('utf-8')

    req = request.Request(
        url,
        method=method,
        headers=headers or {},
        data=body
    )

    try:
        with request.urlopen(req) as res:
            data = res.read().decode('utf-8')
            try:
                return json.loads(data)
            except json.JSONDecodeError:
                return data
    except error.HTTPError as err:
        return f'HTTP Error: {err.code}, {err.reason}'
    except error.URLError as err:
        return f'URL Error: {err.reason}'


def read_file(file: Path | None) -> str:
    """Read contents of a file and return as string. Returns empty string if file is None."""
    if file is None:
        sys.exit(0)

    if not file.exists():
        print(f'Warning: File not found: {file}', file=sys.stderr)
        sys.exit(0)

    try:
        return file.read_text()
    except Exception as e:
        print(f'Error reading file {file}: {e}', file=sys.stderr)
        sys.exit(0)


#
# LLM REQUEST FUNCTIONS
#


def openai_request(prompt: str, model: str, output: str, tokens: int, entropy: float, role: str, verbose: bool = False):
    OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY', None)
    if OPENAI_API_KEY is None:
        print('Could not find `OPENAI_API_KEY` environment variable', file=sys.stderr)
        sys.exit(0)

    res = fetch(
        'https://api.openai.com/v1/chat/completions',
        method='POST',
        headers={
            'Content-Type': 'application/json',
            'Authorization': f'Bearer {OPENAI_API_KEY}',
        },
        body={
            'model': model,
            'messages': [{ 'role': role, 'content': prompt }],
            'temperature': entropy,
            'max_completion_tokens': tokens,
            'response_format': {
                'type': 'json_object' if output == 'json' else output
            },
        }
    )

    if verbose:
        print(json.dumps(res, sort_keys=True, indent=4))
        return

    # Print response in format specified in command
    msg = res['choices'][0]['message']['content']
    if output == 'json':
        data = json.loads(msg)
        prety_json = json.dumps(data, sort_keys=True, indent=4)
        print(prety_json)
    else:
        print(msg)


#
# CLI COMMAND FUNCTIONS
#


def env(verbose: bool = False):
    """ List out `found` and `missing` LLM API environment variables """
    services = {
        'Anthropic': os.environ.get('ANTHROPIC_API_KEY', None),
        'OpenAI': os.environ.get('OPENAI_API_KEY', None),
        'Gemini': os.environ.get('GEMINI_API_KEY', None),
    }

    padding = max(len(service) for service in services)
    for service, var in services.items():
        if verbose:
            status = f'\033[92m{var}\033[0m' if var is not None else '\033[91mmissing\033[0m'
            print(f'{service:<{padding}}\t{status}')
        else:
            status = '\033[92mfound\033[0m' if var is not None else '\033[91mmissing\033[0m'
            print(f'{service:<{padding}}\t{status}')


def chat(prompt: str, model: str, output: str, tokens: int, entropy: float, role: str, files: List[Path] | None, verbose: bool = False):
    prompt = find_and_replace(prompt, files)
    if verbose:
        params = locals()
        print_inputs(**params)

    match model:
        case 'gpt-4' | 'gpt-4o-mini' | 'gpt-4o':
            params = locals()
            openai_request(prompt, model, output, tokens, entropy, role, verbose)
        case 'claude-3-5-sonnet-latest' | 'claude-3-5-haiku-latest' |'claude-3-opus-latest':
            print('NOT_YET_IMPLEMENTED_ERROR: API call not yet implemented for Anthropic')
        case 'gemini-1.5-flash' |'gemini-1.5-flash-8b' | 'gemini-1.5-pro':
            print('NOT_YET_IMPLEMENTED_ERROR: API call not yet implemented for Gemini')


if __name__ == '__main__':
    try:
        parser = create_argparser()
        args = parser.parse_args()

        match args.command:
            case 'env':
                params = vars(args)
                del params['command']

                env(**params)
            case 'chat':
                params = vars(args)
                del params['command']

                chat(**params)
    except KeyboardInterrupt:
        sys.exit(1)