common.py

import argparse
import json
import time
import warnings
from datetime import datetime
from pathlib import Path
from typing import Any, cast

import openai
from openai.types.chat import ChatCompletion
from ratelimit import limits, sleep_and_retry


class ExchangeLogger:
    def __init__(self) -> None:
        self.file: Path | None = None
        self.print_log = False

    def config(self, file: Path, print_log: bool = False) -> None:
        self.file = file
        self.print_log = print_log

    def log_exchange(self, params: dict[str, Any], response: dict[str, Any]) -> None:
        if self.file is None:
            raise ValueError("Must call config() before logging exchanges.")

        log = {"params": params, "response": response}

        with self.file.open("a") as f:
            json.dump(log, f)
            f.write("\n")

        if self.print_log:
            print(json.dumps(log, indent=2))
            print()


logger = ExchangeLogger()


def get_key(key_file: Path, key_name: str) -> str:
    keys = json.loads(key_file.read_text())
    return cast(str, keys[key_name])


def make_msg(role: str, content: str) -> dict[str, str]:
    return {"role": role, "content": content}


CALLS_PER_MINUTE = 3500  # full plan


def make_chat_request(client: openai.OpenAI, **kwargs: Any) -> ChatCompletion:
    # Ignores (mypy): untyped decorator makes function untyped
    @sleep_and_retry  # type: ignore[misc]
    @limits(calls=CALLS_PER_MINUTE, period=60)  # type: ignore[misc]
    def _make_chat_request(**kwargs: Any) -> ChatCompletion:
        attempts = 0
        while True:
            try:
                response = client.chat.completions.create(**kwargs)
            except Exception as e:
                ts = datetime.now().isoformat()
                print(f'{ts} | Connection error - "{e}" | Attempt {attempts + 1}')
                attempts += 1

                if isinstance(e, openai.APIStatusError) and e.status_code == 429:
                    print("Rate limit exceeded. Waiting 10 seconds.")
                    time.sleep(10)
            else:
                logger.log_exchange(kwargs, response.model_dump())
                return response

    # This cast is necessary because of the sleep_and_retry and limits decorators,
    # which make the function untyped.
    return cast(ChatCompletion, _make_chat_request(**kwargs))


def get_result(response: ChatCompletion) -> str:
    return response.choices[0].message.content or "<empty>"


# Costs as of 2023-11-13 from https://openai.com/pricing
# model: input cost, output cost
MODEL_COSTS = {
    "gpt-3.5-turbo-0613": (  # in: $0.001 / 1K tokens, out: $0.002 / 1K tokens
        0.000001,
        0.000002,
    ),
    "gpt-4-1106-preview": (  # in: $0.03 / 1K tokens, out: $0.06 / 1K tokens
        0.00003,
        0.00006,
    ),
    "gpt-3.5-turbo-0125": (  # in: $0.001 / 1K tokens, out: $0.002 / 1K tokens
        0.000001,
        0.000002,
    ),
    "gpt-4-0125-preview": (  # in: $0.03 / 1K tokens, out: $0.06 / 1K tokens
        0.00003,
        0.00006,
    ),
}


def calculate_cost(model: str, response: ChatCompletion) -> float:
    """Calculate the cost of a GPT response based on the tokens used.

    If model is not not valid, a warning is issued and 0 cost is returned. Same if the
    usage information is not available in the response.

    0 cost is returned instead of raising an exception, because the cost is not that
    important, and it's better not to have it than crash the program.
    """
    if response.usage is None:
        warnings.warn(
            "Usage information not available in GPT response. Returning 0 cost."
        )
        return 0

    if model not in MODEL_COSTS:
        warnings.warn(f"Model {model} not found in costs table. Returning 0 cost.")
        return 0

    input_tokens = response.usage.prompt_tokens
    output_tokens = response.usage.completion_tokens
    cost_input, cost_output = MODEL_COSTS[model]

    return input_tokens * cost_input + output_tokens * cost_output


def log_args(args: argparse.Namespace, path: Path | None) -> None:
    args_dict = vars(args).copy()
    for key, value in args_dict.items():
        if isinstance(value, Path):
            args_dict[key] = str(value)

    out = json.dumps(args_dict, indent=2)
    if path is not None:
        path.parent.mkdir(parents=True, exist_ok=True)
        path.write_text(out)
    else:
        print(out)


def init_argparser(*, prompt: bool = True) -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(
        allow_abbrev=False, formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )
    parser.add_argument("key_file", type=Path, help="Path to JSON file with API keys")
    parser.add_argument("key_name", type=str, help="Name of key to use")
    parser.add_argument(
        "--model", type=str, default=next(iter(MODEL_COSTS)), help="Model to use"
    )
    parser.add_argument(
        "--print-logs", action="store_true", help="Print logs to stdout"
    )
    parser.add_argument(
        "--log-file", type=Path, default="chat_log.jsonl", help="Log file"
    )
    parser.add_argument("--input", "-i", type=Path, help="Input file")
    parser.add_argument("--output", "-o", type=Path, help="Output file for predictions")
    parser.add_argument("--metrics-path", type=Path, help="Path where to save metrics")
    parser.add_argument("--args-path", type=Path, help="Path where to save args")

    if prompt:
        parser.add_argument(
            "--prompt",
            type=int,
            default=0,
            help="User prompt index to use for the chat session",
        )
        parser.add_argument(
            "--sys-prompt",
            type=int,
            default=0,
            help="System prompt index to use for the chat session",
        )

    return parser