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 ff58ada6bf
commit b8618b0c91
8 changed files with 485 additions and 1 deletions
+193 -1
View File
@@ -6,7 +6,7 @@ from fastapi.testclient import TestClient
from sqlalchemy import select
from lembas.db.models import Connection, Model
from lembas.services import tokens
from lembas.services import metrics, tokens
from lembas.services.crypto import encrypt
from lembas.services.llm.openai_client import chunk_usage, context_from
@@ -258,3 +258,195 @@ def test_a_request_estimate_includes_the_tools_array():
assert tokens.estimate_request(payload) > tokens.estimate_request(
{"messages": payload["messages"]}
)
# --- The Metrics object -------------------------------------------------------
def test_the_percentage_is_zero_when_nobody_said_how_big_the_window_is():
"""Unknown must stay tellable from small. A percentage of an unknown total
is a made-up number in a place people trust numbers."""
assert metrics.Metrics(context_tokens=5000, context_limit=0).percent == 0
assert metrics.Metrics(context_tokens=5000, context_limit=10000).percent == 50
def test_pressure_marks_the_bar_only_when_it_is_worth_noticing():
assert metrics.Metrics(context_tokens=50, context_limit=100).pressure == ""
assert metrics.Metrics(context_tokens=85, context_limit=100).pressure == "warning"
assert metrics.Metrics(context_tokens=96, context_limit=100).pressure == "danger"
assert metrics.Metrics(context_tokens=0, context_limit=0).pressure == ""
def test_speed_needs_both_a_count_and_a_clock():
assert metrics.Metrics(completion_tokens=100, elapsed_ms=2000).tokens_per_second == 50.0
assert metrics.Metrics(completion_tokens=100, elapsed_ms=0).tokens_per_second == 0.0
assert metrics.Metrics(completion_tokens=0, elapsed_ms=2000).tokens_per_second == 0.0
def test_metrics_survive_a_round_trip_through_the_row():
original = metrics.Metrics(
prompt_tokens=100,
completion_tokens=20,
total_tokens=120,
context_tokens=120,
context_limit=8192,
estimated=True,
elapsed_ms=1500,
rounds=2,
)
assert metrics.from_message(metrics.to_json(original)) == original
def test_a_row_with_no_usage_reads_as_nothing_rather_than_failing():
assert metrics.from_message(None).has_anything is False
assert metrics.from_message({}).has_anything is False
assert metrics.from_message({"prompt_tokens": "lots"}).prompt_tokens == 0
# --- Through a generation -----------------------------------------------------
def _chunks(*frames: str) -> str:
return "".join(f"data: {frame}\n\n" for frame in frames) + "data: [DONE]\n\n"
async def test_usage_is_summed_across_tool_rounds(db, registered, make_chat, monkeypatch):
"""Prompt and completion are what the reply cost; context_tokens is what the
window holds. A three-round reply pays for its prompt three times and only
ever occupies the window once."""
from lembas.services import generation as generation_service
generation = generation_service.Generation(chat_id="c", message_id="m")
for prompt_tokens, completion_tokens in ((100, 10), (250, 20)):
counts = {"prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens}
generation.prompt_tokens += counts["prompt_tokens"]
generation.completion_tokens += counts["completion_tokens"]
generation.context_tokens = prompt_tokens + completion_tokens
assert generation.prompt_tokens == 350
assert generation.completion_tokens == 30
assert generation.context_tokens == 270
async def test_a_reply_records_what_it_cost(client: TestClient, db, registered, mock_http):
import httpx
from lembas.db.models import Message
def handler(_request: httpx.Request) -> httpx.Response:
if _request.url.path.endswith("/models"):
return httpx.Response(200, json={"data": [{"id": "m", "context_length": 1000}]})
return httpx.Response(
200,
text=_chunks(
'{"choices":[{"delta":{"content":"Waybread."}}]}',
'{"choices":[],"usage":{"prompt_tokens":120,"completion_tokens":8}}',
),
)
mock_http(handler)
client.post(
"/admin/connections",
data={"name": "Local", "base_url": "http://x.test", "api_key": ""},
follow_redirects=False,
)
client.post("/api/chats/start", data={"content": "what is lembas?"})
reply = db.scalar(select(Message).where(Message.role == "assistant"))
# Following the stream is what runs the generation to completion.
with client.stream(
"GET", f"/api/chats/{reply.chat_id}/messages/{reply.id}/stream"
) as response:
list(response.iter_lines())
db.refresh(reply)
assert reply.usage_json["prompt_tokens"] == 120
assert reply.usage_json["completion_tokens"] == 8
assert reply.usage_json["context_limit"] == 1000
assert reply.usage_json["estimated"] is False
assert reply.usage_json["elapsed_ms"] >= 0
def test_the_chips_render_from_a_stored_row(client: TestClient, db, registered, make_chat):
from lembas.db.models import Message
_model(db)
chat_id = make_chat()
db.add(
Message(
chat_id=chat_id,
role="assistant",
content="Waybread.",
usage_json={
"prompt_tokens": 100,
"completion_tokens": 20,
"total_tokens": 120,
"context_tokens": 120,
"context_limit": 1000,
"estimated": False,
"elapsed_ms": 2000,
"rounds": 1,
},
)
)
db.commit()
page = client.get(f"/chat/{chat_id}").text
assert "120 tokens" in page
assert "12%" in page
assert "10.0 tok/s" in page
assert "~" not in page.split("msg__metrics")[1][:400]
def test_an_estimated_row_says_so(client: TestClient, db, registered, make_chat):
from lembas.db.models import Message
_model(db)
chat_id = make_chat()
db.add(
Message(
chat_id=chat_id,
role="assistant",
content="Waybread.",
usage_json={
"prompt_tokens": 100,
"completion_tokens": 20,
"total_tokens": 120,
"context_tokens": 120,
"context_limit": 1000,
"estimated": True,
"elapsed_ms": 2000,
"rounds": 1,
},
)
)
db.commit()
page = client.get(f"/chat/{chat_id}").text
assert "~120 tokens" in page
assert "reports no token counts" in page
def test_no_context_length_means_no_percentage(client: TestClient, db, registered, make_chat):
from lembas.db.models import Message
_model(db)
chat_id = make_chat()
db.add(
Message(
chat_id=chat_id,
role="assistant",
content="Waybread.",
usage_json={
"prompt_tokens": 100,
"completion_tokens": 20,
"total_tokens": 120,
"context_tokens": 120,
"context_limit": 0,
"estimated": False,
"elapsed_ms": 2000,
"rounds": 1,
},
)
)
db.commit()
page = client.get(f"/chat/{chat_id}").text
assert "120 tokens" in page
assert "metric--context" not in page