Estimer le coût d’un appel de modèle

"""
Ce qu'un appel de modèle coûtera, avant de le faire, à partir de vos chiffres.

Niveau N0. Un fournisseur facture des jetons, publie un prix par million
d'entre eux, et l'arithmétique entre les deux est une multiplication. Ce qui
vaut une fiche n'est pas la multiplication : c'est tout ce que ce code refuse
de deviner.

Il ne nomme jamais un prix. Les tarifs changent tous les trimestres et
diffèrent selon le modèle, la région et le contrat ; les deux prix sont des
paramètres, et tout montant rendu est dans l'unité dans laquelle vous les
avez passés. Un extrait qui porterait un prix serait faux avant d'être lu.

Il dit si les jetons ont été comptés ou estimés. Les compter demande le
tokeniseur de votre fournisseur ; les estimer depuis un nombre de caractères
demande un rapport, et un rapport n'est pas une constante de la nature : il
dépend de la langue, de la ponctuation, du fait que le texte soit de la prose
ou du JSON. Les deux sont permis ici, et le rapport dit lequel a servi, parce
qu'une estimation présentée comme un comptage est exactement la faute que
cette fiche existe pour éviter.

Et il répond à la question qui tranche : `break_even_items` est le nombre
d'appels à partir duquel la facture atteint le coût unique de ce que l'appel
remplacerait. C'est ce nombre-là, et non le prix par million, qui dit si
l'appel vaut la peine.

L'arrondi n'a lieu qu'une fois, à la fin. Le coût d'un appel est un petit
nombre — quelques dix-millièmes — et l'arrondir avant de le multiplier par le
nombre d'appels multiplie l'arrondi avec lui : à deux décimales, cent mille
appels qui coûtent 36 revenaient à 0,00. Le calcul est mené à une échelle
plus fine que celle demandée, le total y est calculé, et seul ce qui est
affiché est arrondi. C'est la discipline de la fiche voisine sur la somme
d'une colonne, appliquée à la fiche qui la cite.
"""

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}

/**
 * Ce qu'un appel de modèle coûtera, avant de le faire, à partir de vos chiffres.
 *
 * Niveau N0. Un fournisseur facture des jetons, publie un prix par million
 * d'entre eux, et l'arithmétique entre les deux est une multiplication. Ce qui
 * vaut une fiche n'est pas la multiplication : c'est tout ce que ce code refuse
 * de deviner.
 *
 * Il ne nomme jamais un prix. Les tarifs changent tous les trimestres et
 * diffèrent selon le modèle, la région et le contrat ; les deux prix sont des
 * paramètres, et tout montant rendu est dans l'unité dans laquelle vous les
 * avez passés. Un extrait qui porterait un prix serait faux avant d'être lu.
 *
 * Il dit si les jetons ont été comptés ou estimés. Les compter demande le
 * tokeniseur de votre fournisseur ; les estimer depuis un nombre de caractères
 * demande un rapport, et un rapport n'est pas une constante de la nature : il
 * dépend de la langue, de la ponctuation, du fait que le texte soit de la prose
 * ou du JSON. Les deux sont permis ici, et le rapport dit lequel a servi, parce
 * qu'une estimation présentée comme un comptage est exactement la faute que
 * cette fiche existe pour éviter.
 *
 * Et il répond à la question qui tranche : `break_even_items` est le nombre
 * d'appels à partir duquel la facture atteint le coût unique de ce que l'appel
 * remplacerait. C'est ce nombre-là, et non le prix par million, qui dit si
 * l'appel vaut la peine.
 *
 * L'arrondi n'a lieu qu'une fois, à la fin. Le coût d'un appel est un petit
 * nombre — quelques dix-millièmes — et l'arrondir avant de le multiplier par le
 * nombre d'appels multiplie l'arrondi avec lui : à deux décimales, cent mille
 * appels qui coûtent 36 revenaient à 0,00. Le calcul est mené à une échelle
 * plus fine que celle demandée, le total y est calculé, et seul ce qui est
 * affiché est arrondi. C'est la discipline de la fiche voisine sur la somme
 * d'une colonne, appliquée à la fiche qui la cite.
 */

// 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 };
}