Show what a reply cost, live and afterwards

Tokens, how full the context is, and tokens per second -- as chips under
each assistant bubble, updating while the reply streams and still there when
it finishes.

The numbers come from one Metrics object built either from the generation
still being written or from the row it left behind. That is the point rather
than tidiness: the finished bubble is re-rendered from the database the
instant the stream ends, so two code paths would make the figures visibly
jump at exactly the moment someone is watching them. Here the only thing
that changes is that an estimate may become exact.

Message.usage_json has existed and been dead since the schema was written.
It is the store.

Two counts that look like one. prompt and completion are summed across tool
rounds -- what the reply cost. context_tokens is overwritten each round with
that round's prompt plus completion -- what the window actually holds. A
three-round reply pays for its prompt three times and only ever occupies the
window once, so a single number would be wrong for one of the two questions.

Generation gains started_at as a field rather than a local in _run, because
_follow is a different function that sees only the Generation and otherwise
has nothing to compute a live speed against. It also carries a prompt
estimate taken before the first chunk, since real usage arrives in one chunk
at the very end and a percentage that appears only after the reply is
useless.

Everything is marked with a tilde when the endpoint reported nothing, and
the percentage is simply absent when no context length is set: unknown has
to stay tellable from small, and a percentage of an unknown total is a
made-up number in a place people trust numbers.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Jaroslav Beneš
2026-08-01 00:40:44 +02:00
parent 9e2caeac48
commit e185edc9e1
8 changed files with 485 additions and 1 deletions
+136
View File
@@ -0,0 +1,136 @@
"""What a reply cost, how fast it arrived, and how full the window is.
One shape, built either from a generation still being written or from the row
it left behind. That matters more than it looks: the finished bubble is
re-rendered from the database the instant the stream ends, so if the live
numbers and the stored ones came from different code they would visibly jump at
exactly the moment the reader is looking at them. Here the only thing that
changes when a reply finishes is that an estimate may become exact.
Nothing here is authoritative about tokens. `estimated` says which kind of
number this is, and every surface that shows one has to say so too.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Any
from lembas.services import tokens
# Where the context bar changes colour. Not thresholds anyone tunes: they mark
# "worth noticing" and "about to be a problem", and the second is deliberately
# below the default compaction threshold so the warning arrives first.
WARNING_AT = 80
DANGER_AT = 95
@dataclass(frozen=True)
class Metrics:
"""Token counts and timing for one reply."""
prompt_tokens: int = 0
completion_tokens: int = 0
total_tokens: int = 0
# What the window holds after this turn: the last round's prompt plus its
# completion. Distinct from prompt+completion summed over tool rounds, which
# is what the reply *cost* -- a three-round reply pays for its prompt three
# times but only ever occupies the window once.
context_tokens: int = 0
context_limit: int = 0
estimated: bool = False
elapsed_ms: int = 0
rounds: int = 1
@property
def percent(self) -> int:
"""How full the window is, or 0 when nobody has said how big it is."""
if self.context_limit <= 0 or self.context_tokens <= 0:
return 0
return min(100, round(self.context_tokens * 100 / self.context_limit))
@property
def tokens_per_second(self) -> float:
if self.elapsed_ms <= 0 or self.completion_tokens <= 0:
return 0.0
return self.completion_tokens / (self.elapsed_ms / 1000)
@property
def pressure(self) -> str:
""""", "warning" or "danger" -- the class the context chip takes."""
percent = self.percent
if not percent:
return ""
if percent >= DANGER_AT:
return "danger"
if percent >= WARNING_AT:
return "warning"
return ""
@property
def has_anything(self) -> bool:
return bool(self.total_tokens or self.completion_tokens or self.elapsed_ms)
def from_generation(generation: Any) -> Metrics:
"""Metrics for a reply still being written.
Usage arrives in a single chunk at the very end, so mid-stream there is
nothing to report and everything is estimated. The counts stop being
estimates the moment that chunk lands, which is usually a beat before the
bubble is replaced.
"""
import time
completion = generation.completion_tokens or tokens.estimate(
generation.text + generation.thinking
)
prompt = generation.prompt_tokens or generation.prompt_estimate
elapsed = generation.elapsed_ms or (
int((time.monotonic() - generation.started_at) * 1000) if generation.started_at else 0
)
return Metrics(
prompt_tokens=prompt,
completion_tokens=completion,
total_tokens=prompt + completion,
context_tokens=generation.context_tokens or (prompt + completion),
context_limit=generation.context_limit,
estimated=not (generation.prompt_tokens and generation.completion_tokens),
elapsed_ms=elapsed,
rounds=max(1, generation.rounds),
)
def from_message(usage_json: dict[str, Any] | None) -> Metrics:
"""Metrics for a finished reply, read back off the row."""
stored = usage_json or {}
def _int(key: str) -> int:
value = stored.get(key)
return int(value) if isinstance(value, (int, float)) and not isinstance(value, bool) else 0
return Metrics(
prompt_tokens=_int("prompt_tokens"),
completion_tokens=_int("completion_tokens"),
total_tokens=_int("total_tokens"),
context_tokens=_int("context_tokens"),
context_limit=_int("context_limit"),
estimated=bool(stored.get("estimated")),
elapsed_ms=_int("elapsed_ms"),
rounds=max(1, _int("rounds")),
)
def to_json(metrics: Metrics) -> dict[str, Any]:
"""The shape stored in Message.usage_json."""
return {
"prompt_tokens": metrics.prompt_tokens,
"completion_tokens": metrics.completion_tokens,
"total_tokens": metrics.total_tokens,
"context_tokens": metrics.context_tokens,
"context_limit": metrics.context_limit,
"estimated": metrics.estimated,
"elapsed_ms": metrics.elapsed_ms,
"rounds": metrics.rounds,
}