Tokens, Context and Cost: Measure Before You Spend

"""Tokens, context and cost: count before you send, price before you run, fit before you overflow."""
import json

import tiktoken
from openai import OpenAI

CLIENT = OpenAI()
MODEL = "meta/llama-3.1-8b-instruct"
ENC = tiktoken.get_encoding("cl100k_base")
PRICES = json.load(open("prices.json"))["models"]


# ---------- Step 1: count ----------
def count_tokens(text):
    """How many tokens `text` is, with the cl100k_base tokenizer."""
    # TODO (Step 1): ENC.encode(text) turns text into a list of token ids.
    # Return how many there are.
    raise NotImplementedError("Step 1: write count_tokens()")


def chars_per_token(text):
    """Average characters per token, rounded to 2 places (0.0 for empty text)."""
    # TODO (Step 1): len(text) / count_tokens(text), rounded to 2 places.
    # Return 0.0 for empty text (dividing by zero crashes).
    raise NotImplementedError("Step 1: write chars_per_token()")


# ---------- Step 2: your count vs the provider's receipt ----------
def measure(text, client=None):
    """Send `text` as one user message (max_tokens=1) and compare counts.
    Return {"local": your count, "reported": usage.prompt_tokens, "overhead": reported - local}."""
    client = client or CLIENT
    raise NotImplementedError("measure() arrives in Step 2")


# ---------- Step 3: price a job before you run it ----------
def call_cost(prompt_tokens, completion_tokens, model):
    """US dollars for one call, from PRICES (per 1M tokens)."""
    raise NotImplementedError("call_cost() arrives in Step 3")


def quote_job(texts, instructions, model, output_tokens_each, overhead_each=40):
    """Estimate a job that sends `instructions` + one text per call.
    Return {"calls", "prompt_tokens", "completion_tokens", "usd"} (usd rounded to 4 places)."""
    raise NotImplementedError("quote_job() arrives in Step 3")


# ---------- Step 4: make it fit ----------
def truncate_tokens(text, max_tokens):
    """The longest start of `text` that is at most max_tokens tokens (cut on tokens, not characters)."""
    raise NotImplementedError("truncate_tokens() arrives in Step 4")


def estimate_messages(messages):
    """Our rule of thumb for a whole request: each message costs its content + 4, plus 3 for the reply."""
    return sum(count_tokens(m["content"]) + 4 for m in messages) + 3


def fit_history(messages, budget):
    """Drop the oldest messages after the system message until estimate_messages() <= budget.
    Never drop the system message or the last message. Return a new list."""
    raise NotImplementedError("fit_history() arrives in Step 4")


# ---------- Step 5: notice when the reply was cut off ----------
def complete(prompt, max_tokens, client=None):
    """One call. Return {"text", "finish_reason", "truncated"} where truncated means the
    model stopped because it hit max_tokens, not because it was done."""
    client = client or CLIENT
    raise NotImplementedError("complete() arrives in Step 5")


def reply_budget(messages, context_window, wanted):
    """The prompt and the reply share one context window. Return how many tokens the reply
    may use: `wanted`, or less if the window is nearly full. Raise ValueError if nothing is left."""
    raise NotImplementedError("reply_budget() arrives in Step 5")