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
+55
View File
@@ -25,10 +25,13 @@ from datetime import UTC, datetime, timedelta
from lembas.db.models import ROLE_ASSISTANT, ROLE_USER, Chat, Message, User
from lembas.db.session import session_scope
from lembas.services import chat as chat_service
from lembas.services import metrics as metrics_service
from lembas.services import prompts as prompts_service
from lembas.services import tokens
from lembas.services import tools as tools_service
from lembas.services.llm.openai_client import (
LLMError,
chunk_usage,
delta_reasoning,
delta_text,
delta_tool_calls,
@@ -64,6 +67,26 @@ class Generation:
# live as the model works and kept on the message afterwards.
tool_events: list[dict] = field(default_factory=list)
# --- What it cost --------------------------------------------------------
# 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, because a three-round reply pays for its prompt three times
# but only ever occupies the window once.
prompt_tokens: int = 0
completion_tokens: int = 0
context_tokens: int = 0
context_limit: int = 0
# Filled from the assembled request before the first chunk, so a follower
# has a percentage to show while the reply is still being written -- real
# usage only arrives in a single chunk at the very end.
prompt_estimate: int = 0
rounds: int = 0
# time.monotonic() at the start. A field rather than a local in `_run`
# because `_follow` is a different function that sees only this object, and
# without it there is nothing to compute a live tokens/second against.
started_at: float = 0.0
elapsed_ms: int = 0
error: str = ""
stopped: bool = False
done: bool = False
@@ -183,6 +206,7 @@ async def _run(generation: Generation) -> None:
"""
splitter = ReasoningSplitter()
started = time.monotonic()
generation.started_at = started
reasoning_started: float | None = None
question = ""
endpoint = model_id = None
@@ -212,13 +236,30 @@ async def _run(generation: Generation) -> None:
title_prompt = prompts_service.resolve(db, "task.title")
tool_context = tools_service.context_for(db, owner, chat)
model = chat_service.model_for(db, chat)
generation.context_limit = model.context_length if model is not None else 0
generation.prompt_estimate = tokens.estimate_request(payload)
for round_number in range(tools_service.MAX_ROUNDS + 1):
generation.rounds = round_number + 1
accumulator = tools_service.ToolCallAccumulator()
# Text the model produced in *this* round, needed separately from
# generation.content when echoing the assistant turn back.
round_text: list[str] = []
async for chunk in stream_chat(endpoint, payload):
counts = chunk_usage(chunk)
if counts is not None:
generation.prompt_tokens += counts.get("prompt_tokens", 0)
generation.completion_tokens += counts.get("completion_tokens", 0)
# Overwritten, not summed: this round's prompt already
# contains every earlier round.
generation.context_tokens = counts.get("prompt_tokens", 0) + counts.get(
"completion_tokens", 0
)
generation.touch()
thought = delta_reasoning(chunk)
if thought:
if reasoning_started is None:
@@ -307,6 +348,17 @@ async def _run(generation: Generation) -> None:
if reasoning_started is not None and not generation.reasoning_ms:
generation.reasoning_ms = int((time.monotonic() - reasoning_started) * 1000)
generation.elapsed_ms = int((time.monotonic() - started) * 1000)
if not generation.completion_tokens:
# The endpoint reported nothing, so fall back to the estimate. Marked
# as such everywhere it is shown -- four characters to a token is
# wrong enough on code and CJK to be worth saying out loud.
generation.completion_tokens = tokens.estimate(
generation.text + generation.thinking
)
generation.prompt_tokens = generation.prompt_estimate
generation.context_tokens = generation.prompt_tokens + generation.completion_tokens
# Naming the chat is a second, short completion, so it has to happen
# here rather than in the synchronous persist step below. Best-effort:
# a chat title is never worth surfacing an error for.
@@ -379,6 +431,9 @@ def _persist(generation: Generation, title: str, elapsed: float) -> None:
message.reasoning = generation.thinking
message.reasoning_ms = generation.reasoning_ms
message.tool_calls_json = generation.tool_events
message.usage_json = metrics_service.to_json(
metrics_service.from_generation(generation)
)
message.error = generation.error
message.stopped = generation.stopped
message.complete = True