chatcost.py

"""What a chat request actually costs, in tokens.

The API bills for the whole serialised conversation, not the words you typed.
Every message is wrapped in ChatML scaffolding — <|im_start|>, the role name,
<|im_sep|>, <|im_end|> — and the request ends with a partial header inviting the
assistant to speak. That works out to an exact rule, which verify_rule() checks
against the tokenizer's own chat encoder rather than trusting it:

    billed = sum(content tokens) + 4 per message + 3

The 4 is <|im_start|>, the role, <|im_sep|> and <|im_end|>. The 3 is the trailing
<|im_start|>assistant<|im_sep|> that primes the reply.

This is the ChatML layout used by the GPT-4/4o family. Other providers wrap
messages differently; the fact that there IS a wrapper is universal, the exact
count is not.
"""

import json
import random

from cjsload import Reference

PER_MESSAGE = 4
PER_REQUEST = 3

MODEL = "gpt-4o"


class Chat:
    def __init__(self, encoding="o200k_base"):
        self.ref = Reference(encoding)

    def count(self, text):
        return len(self.ref.ids(text))

    def encode_chat(self, messages):
        self.ref.ctx.set("_m", json.dumps(messages))
        return json.loads(self.ref.ctx.eval(
            f'JSON.stringify(Array.from(M.encodeChat(JSON.parse(_m), "{MODEL}")))'))

    def billed(self, messages):
        """Apply the rule rather than the encoder — same answer, and it is the
        rule the browser calculator uses."""
        content = sum(self.count(m["content"]) for m in messages)
        return content + PER_MESSAGE * len(messages) + PER_REQUEST

    def conversation_cost(self, system, turns, user_tokens, reply_tokens):
        """Cumulative tokens sent across a whole conversation.

        Turn n re-sends everything before it, so the input side grows linearly
        per turn and the running total grows with the square of the turn count.
        """
        sys_tokens = self.count(system)
        rows = []
        cumulative = 0
        history = sys_tokens + PER_MESSAGE          # the system message
        messages = 1
        for turn in range(1, turns + 1):
            history += user_tokens + PER_MESSAGE    # this turn's user message
            messages += 1
            sent = history + PER_REQUEST
            cumulative += sent
            rows.append({
                "turn": turn,
                "sent": sent,
                "cumulative": cumulative,
                "system_share": sys_tokens,
                "resent": sent - (user_tokens + PER_MESSAGE + PER_REQUEST),
            })
            history += reply_tokens + PER_MESSAGE   # the assistant's reply
            messages += 1
        return rows


def verify_rule(chat, trials=200, seed=5):
    """The arithmetic must match the tokenizer's own chat encoder exactly."""
    rng = random.Random(seed)
    words = ["hello", "please", "summarise", "the", "document", "café", "🍓",
             "for me", "in Japanese", "こんにちは", "x = 1", "", " ", "\n\n",
             "a much longer sentence that goes on for a while without stopping"]
    roles = ["system", "user", "assistant"]
    bad = []
    for _ in range(trials):
        n = rng.randint(1, 6)
        msgs = [{"role": rng.choice(roles),
                 "content": " ".join(rng.choice(words)
                                     for _ in range(rng.randint(1, 8)))}
                for _ in range(n)]
        try:
            actual = len(chat.encode_chat(msgs))
        except Exception as e:
            bad.append((msgs, f"encoder failed: {e}"))
            continue
        predicted = chat.billed(msgs)
        if actual != predicted:
            bad.append((msgs, f"rule {predicted} != encoder {actual}"))
    return bad


if __name__ == "__main__":
    chat = Chat()
    bad = verify_rule(chat)
    if bad:
        print(f"FAIL  chat rule disagrees with the encoder on {len(bad)} of 200")
        for msgs, why in bad[:3]:
            print(f"  {why}: {msgs}")
        raise SystemExit(1)
    print("PASS  chat cost rule matches encodeChat on 200 random conversations")

    demo = [{"role": "system", "content": "You are a helpful assistant."},
            {"role": "user", "content": "What is the capital of France?"}]
    content = sum(chat.count(m["content"]) for m in demo)
    print(f"      example: {content} tokens of content -> "
          f"{chat.billed(demo)} billed")