Estimate the cost of a model call

"""
What a model call will cost, before making it, from figures you provide.

Rung N0. A provider bills tokens, publishes a price per million of them, and
the arithmetic between the two is a multiplication. What makes this worth a
page is not the multiplication: it is everything this code refuses to guess.

It never names a price. Prices change every quarter and differ per model, per
region and per contract; the two prices are parameters, and every amount that
comes back is in whatever unit you passed them in. A snippet that carried a
price would be wrong before it was read.

It says whether the tokens were counted or estimated. Counting them needs the
tokeniser your provider uses; estimating them from a character count needs a
ratio, and a ratio is not a constant of nature — it depends on the language,
on the punctuation, on whether the text is prose or JSON. Both are allowed
here, and the report says which one was used, because an estimate presented as
a count is the failure this page exists to avoid.

And it answers the question that decides: `break_even_items` is the number of
calls at which the bill reaches the one-off cost of the thing the call would
replace. That number, not the price per million, is what says whether the call
is worth making.

Rounding happens once, at the end. The cost of one call is a small number — a
few ten-thousandths — and rounding it before multiplying by the number of calls
multiplies the rounding with it: at two decimals, a hundred thousand calls that
cost 36 came back as 0,00. The arithmetic is carried at a finer scale than the
one asked for, the total is computed there, and only what is rendered is
rounded. It is the discipline of the neighbouring entry on summing a column,
applied to the entry that cites it.
"""

from __future__ import annotations

import re

# A price or a ratio as it is written: digits, optionally a decimal point.
# Nothing becomes a float here — money is counted in integers.
NUMBER = re.compile(r"\A[0-9]+(?:\.[0-9]+)?\Z")

# The unit the published prices use. It is a parameter because it is a
# convention, not a law: some providers quote per thousand tokens.
PER = 1_000_000

# How many decimals the amounts come back with, unless the prices carry more.
DECIMALS = 6

# Decimals the arithmetic keeps beyond the ones asked for. Nine is far more
# than any price and any token count needs — the exact value of a call at a
# price per million has at most six decimals more than the price — so the
# amounts below are exact, and rounding happens once, on the way out.
WORKING_DECIMALS = 9


def estimate_cost(items, tokens_in, tokens_out, price_in, price_out, *,
                  per: int = PER, fixed_alternative=None, decimals: int = DECIMALS) -> dict:
    """
    The tokens a job will use and what it will cost, in your own unit.

    `tokens_in` and `tokens_out` are either a whole number of tokens — counted
    with your provider's tokeniser — or a pair `{"characters": n,
    "characters_per_token": "4"}`, which is an estimate and is reported as one.

    `reason` says why something is missing: the whole report, when an argument
    is unusable; or `break_even_items` alone, when the call costs nothing at
    these prices and no number of calls reaches the alternative's cost. With no
    alternative declared, `break_even_items` is None and `reason` stays None:
    there was no question to answer.
    """
    if not isinstance(items, int) or isinstance(items, bool) or items < 0:
        return _report(None, None, None, "items must be a whole number, zero or more")
    prices = {}
    for name, price in (("in", price_in), ("out", price_out)):
        parsed = _parse(price)
        if parsed is None:
            return _report(None, None, None, f"the {name} price must be written in digits")
        prices[name] = parsed
    if not isinstance(per, int) or per <= 0:
        return _report(None, None, None, "the price unit must be a whole number of tokens")

    counts, sources = {}, []
    for name, value in (("in", tokens_in), ("out", tokens_out)):
        count, source = _tokens(value)
        if count is None:
            return _report(None, None, None, f"the {name} tokens are neither a count nor an estimate")
        counts[name], _ = count, sources.append(source)

    scale = max(decimals, prices["in"][1], prices["out"][1])
    work = scale + WORKING_DECIMALS
    # Everything is computed at `work` decimals, including the total, and only
    # the two rendered amounts are brought back to `scale`.
    per_item = sum(_amount(counts[name], prices[name], per, work) for name in ("in", "out"))
    total = per_item * items

    break_even, reason = None, None
    if fixed_alternative is not None:
        fixed = _parse(fixed_alternative)
        if fixed is None:
            return _report(None, None, None, "the alternative's cost must be written in digits")
        if per_item > 0:
            fixed_scaled = fixed[0] * 10 ** (work - fixed[1])
            # A ceiling: the first call whose bill reaches the fixed cost.
            break_even = -(-fixed_scaled // per_item)
        else:
            reason = ("the call costs nothing at these prices: "
                      "no number of calls reaches the alternative's cost")

    # A count declared at zero does not weigh in the verdict: an estimate of
    # nothing beside a count of nothing is not a mixed report.
    weighing = {source for name, source in zip(("in", "out"), sources) if counts[name] > 0}
    tokens = {"in": counts["in"], "out": counts["out"],
              "total": counts["in"] + counts["out"],
              "source": weighing.pop() if len(weighing) == 1 else
                        ("mixed" if weighing else "counted")}
    cost = {"per_item": _text(_round(per_item, work - scale), scale),
            "total": _text(_round(total, work - scale), scale),
            "decimals": scale}
    return _report(tokens, cost, break_even, reason)


def _round(digits: int, places: int) -> int:
    """A scaled integer brought down by `places` decimals, half away from zero."""
    if places <= 0:
        return digits
    unit = 10 ** places
    return (2 * digits + unit) // (2 * unit)


def _tokens(value):
    """A counted number of tokens, or an estimate from a number of characters."""
    if isinstance(value, int) and not isinstance(value, bool) and value >= 0:
        return value, "counted"
    if isinstance(value, dict):
        characters = value.get("characters")
        ratio = _parse(value.get("characters_per_token"))
        if (isinstance(characters, int) and not isinstance(characters, bool)
                and characters >= 0 and ratio and ratio[0] > 0):
            # A part of a token is billed as a token: the estimate rounds up.
            return -(-characters * 10 ** ratio[1] // ratio[0]), "estimated"
    return None, None


def _amount(tokens: int, price, per: int, scale: int) -> int:
    """What these tokens cost at this price, as an integer of `scale` decimals."""
    digits, places = price
    numerator = tokens * digits * 10 ** (scale - places)
    return (2 * numerator + per) // (2 * per)  # rounded away from zero on a half


def _parse(value):
    """A written number as an integer and its number of decimals."""
    if isinstance(value, int) and not isinstance(value, bool) and value >= 0:
        return value, 0
    if not isinstance(value, str) or not NUMBER.match(value.strip()):
        return None
    whole, _, fraction = value.strip().partition(".")
    return int(f"{whole}{fraction}"), len(fraction)


def _text(digits: int, places: int) -> str:
    body = str(digits).rjust(places + 1, "0")
    return body if places == 0 else f"{body[:-places]}.{body[-places:]}"


def _report(tokens, cost, break_even, reason) -> dict:
    return {"tokens": tokens, "cost": cost, "break_even_items": break_even,
            "reason": reason}

/**
 * What a model call will cost, before making it, from figures you provide.
 *
 * Rung N0. A provider bills tokens, publishes a price per million of them, and
 * the arithmetic between the two is a multiplication. What makes this worth a
 * page is not the multiplication: it is everything this code refuses to guess.
 *
 * It never names a price. Prices change every quarter and differ per model,
 * per region and per contract; the two prices are parameters, and every amount
 * that comes back is in whatever unit you passed them in. A snippet that
 * carried a price would be wrong before it was read.
 *
 * It says whether the tokens were counted or estimated. Counting them needs
 * the tokeniser your provider uses; estimating them from a character count
 * needs a ratio, and a ratio is not a constant of nature — it depends on the
 * language, on the punctuation, on whether the text is prose or JSON. Both are
 * allowed here, and the report says which one was used, because an estimate
 * presented as a count is the failure this page exists to avoid.
 *
 * And it answers the question that decides: `break_even_items` is the number
 * of calls at which the bill reaches the one-off cost of the thing the call
 * would replace. That number, not the price per million, is what says whether
 * the call is worth making.
 */

// A price or a ratio as it is written: digits, optionally a decimal point.
// Nothing becomes a float here — money is counted in integers.
const NUMBER = /^[0-9]+(?:\.[0-9]+)?$/;

// The unit the published prices use. It is a parameter because it is a
// convention, not a law: some providers quote per thousand tokens.
export const PER = 1_000_000;

// How many decimals the amounts come back with, unless the prices carry more.
export const DECIMALS = 6;

// Decimals the arithmetic keeps beyond the ones asked for. Nine is far more
// than any price and any token count needs — the exact value of a call at a
// price per million has at most six decimals more than the price — so the
// amounts below are exact, and rounding happens once, on the way out.
export const WORKING_DECIMALS = 9;

/**
 * The tokens a job will use and what it will cost, in your own unit.
 *
 * `tokensIn` and `tokensOut` are either a whole number of tokens — counted
 * with your provider's tokeniser — or a pair `{characters,
 * characters_per_token}`, which is an estimate and is reported as one.
 *
 * `reason` says why something is missing: the whole report, when an argument
 * is unusable; or `break_even_items` alone, when the call costs nothing at
 * these prices and no number of calls reaches the alternative's cost. With no
 * alternative declared, `break_even_items` is null and `reason` stays null:
 * there was no question to answer.
 */
export function estimateCost(items, tokensIn, tokensOut, priceIn, priceOut, options = {}) {
  const { per = PER, fixedAlternative = null, decimals = DECIMALS } = options;
  if (!Number.isInteger(items) || items < 0) {
    return report(null, null, null, 'items must be a whole number, zero or more');
  }
  const prices = {};
  for (const [name, price] of [['in', priceIn], ['out', priceOut]]) {
    const parsed = parseNumber(price);
    if (parsed === null) {
      return report(null, null, null, `the ${name} price must be written in digits`);
    }
    prices[name] = parsed;
  }
  if (!Number.isInteger(per) || per <= 0) {
    return report(null, null, null, 'the price unit must be a whole number of tokens');
  }

  const counts = {};
  const sources = [];
  for (const [name, value] of [['in', tokensIn], ['out', tokensOut]]) {
    const [count, source] = tokensOf(value);
    if (count === null) {
      return report(null, null, null, `the ${name} tokens are neither a count nor an estimate`);
    }
    counts[name] = count;
    sources.push(source);
  }

  const scale = Math.max(decimals, prices.in[1], prices.out[1]);
  const work = scale + WORKING_DECIMALS;
  // Everything is computed at `work` decimals, including the total, and only
  // the two rendered amounts are brought back to `scale`.
  const perItem = ['in', 'out']
    .reduce((sum, name) => sum + amountOf(counts[name], prices[name], per, work), 0n);
  const total = perItem * BigInt(items);

  let breakEven = null;
  let reason = null;
  if (fixedAlternative !== null && fixedAlternative !== undefined) {
    const fixed = parseNumber(fixedAlternative);
    if (fixed === null) {
      return report(null, null, null, "the alternative's cost must be written in digits");
    }
    if (perItem > 0n) {
      const fixedScaled = fixed[0] * 10n ** BigInt(work - fixed[1]);
      // A ceiling: the first call whose bill reaches the fixed cost.
      breakEven = Number((fixedScaled + perItem - 1n) / perItem);
    } else {
      reason = 'the call costs nothing at these prices: '
        + "no number of calls reaches the alternative's cost";
    }
  }

  // A count declared at zero does not weigh in the verdict: an estimate of
  // nothing beside a count of nothing is not a mixed report.
  const weighing = new Set(['in', 'out']
    .filter((name) => counts[name] > 0)
    .map((name) => sources[name === 'in' ? 0 : 1]));
  const tokens = {
    in: counts.in,
    out: counts.out,
    total: counts.in + counts.out,
    source: weighing.size === 1 ? [...weighing][0] : (weighing.size > 1 ? 'mixed' : 'counted'),
  };
  const cost = {
    per_item: asText(roundDown(perItem, work - scale), scale),
    total: asText(roundDown(total, work - scale), scale),
    decimals: scale,
  };
  return report(tokens, cost, breakEven, reason);
}

/** A scaled integer brought down by `places` decimals, half away from zero. */
function roundDown(digits, places) {
  if (places <= 0) return digits;
  const unit = 10n ** BigInt(places);
  return (2n * digits + unit) / (2n * unit);
}

/** A counted number of tokens, or an estimate from a number of characters. */
function tokensOf(value) {
  if (Number.isInteger(value) && value >= 0) return [value, 'counted'];
  if (value && typeof value === 'object' && !Array.isArray(value)) {
    const characters = value.characters;
    // The key is spelled the same in both languages, so one report can be
    // compared with the other.
    const ratio = parseNumber(value.characters_per_token);
    if (Number.isInteger(characters) && characters >= 0 && ratio && ratio[0] > 0n) {
      // A part of a token is billed as a token: the estimate rounds up.
      const numerator = BigInt(characters) * 10n ** BigInt(ratio[1]);
      return [Number((numerator + ratio[0] - 1n) / ratio[0]), 'estimated'];
    }
  }
  return [null, null];
}

/** What these tokens cost at this price, as an integer of `scale` decimals. */
function amountOf(tokens, price, per, scale) {
  const [digits, places] = price;
  const numerator = BigInt(tokens) * digits * 10n ** BigInt(scale - places);
  const unit = BigInt(per);
  return (2n * numerator + unit) / (2n * unit); // rounded away from zero on a half
}

/** A written number as an integer and its number of decimals. */
function parseNumber(value) {
  if (Number.isInteger(value) && value >= 0) return [BigInt(value), 0];
  if (typeof value !== 'string' || !NUMBER.test(value.trim())) return null;
  const [whole, fraction = ''] = value.trim().split('.');
  return [BigInt(`${whole}${fraction}`), fraction.length];
}

function asText(digits, places) {
  const body = digits.toString().padStart(places + 1, '0');
  return places === 0 ? body : `${body.slice(0, -places)}.${body.slice(-places)}`;
}

function report(tokens, cost, breakEven, reason) {
  return { tokens, cost, break_even_items: breakEven, reason };
}