diff --git a/src/lembas/api/chats.py b/src/lembas/api/chats.py index db22ee1..0554681 100644 --- a/src/lembas/api/chats.py +++ b/src/lembas/api/chats.py @@ -21,6 +21,7 @@ from lembas.services import audio as audio_service from lembas.services import chat as chat_service from lembas.services import files as files_service from lembas.services import generation as generation_service +from lembas.services import metrics as metrics_service from lembas.services import sse from lembas.services.markdown import escape_text, render_markdown from lembas.web.templating import render, templates @@ -257,6 +258,7 @@ async def _follow(chat_id: str, message_id: str) -> AsyncIterator[str]: yield sse.event("tools", _tool_activity(generation.tool_events)) if generation.content: yield sse.event("render", render_markdown(generation.text)) + yield sse.event("metrics", _metrics_html(generation)) last_frame = time.monotonic() if generation.done: @@ -314,6 +316,18 @@ async def _follow(chat_id: str, message_id: str) -> AsyncIterator[str]: yield sse.event("close", "") +def _metrics_html(generation) -> str: + """The metric chips for a reply still being written. + + Built from the same Metrics object the finished bubble uses, so the numbers + do not jump when the stream ends -- the only thing that changes is that an + estimate may have become exact. + """ + return templates.get_template("chat/_metrics.html").render( + {"metrics": metrics_service.from_generation(generation)} + ) + + def _thread_context(db: DBSession, chat: Chat, user: User) -> dict: """Everything chat/_thread.html needs to render the conversation.""" messages = list( diff --git a/src/lembas/services/generation.py b/src/lembas/services/generation.py index aaf751f..b2913ec 100644 --- a/src/lembas/services/generation.py +++ b/src/lembas/services/generation.py @@ -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 diff --git a/src/lembas/services/metrics.py b/src/lembas/services/metrics.py new file mode 100644 index 0000000..b371f12 --- /dev/null +++ b/src/lembas/services/metrics.py @@ -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, + } diff --git a/src/lembas/web/static/css/chat.css b/src/lembas/web/static/css/chat.css index fc9794b..44f148c 100644 --- a/src/lembas/web/static/css/chat.css +++ b/src/lembas/web/static/css/chat.css @@ -167,6 +167,43 @@ .reasoning__summary::-webkit-details-marker { display: none; } .reasoning__summary:hover { color: var(--ink); background: var(--surface-hover); } +/* --- Metrics --------------------------------------------------------------- + What a reply cost, under the bubble. Quiet by default: it is reference, not + something to read every time. +*/ +.msg__metrics { + display: flex; + align-items: center; + flex-wrap: wrap; + gap: var(--sp-3); + margin-top: var(--sp-2); + font-size: var(--text-xs); + color: var(--ink-faint); + font-variant-numeric: tabular-nums; +} +.msg__metrics:empty { display: none; } + +.metric { display: inline-flex; align-items: center; gap: var(--sp-1); cursor: default; } + +.metric__bar { + display: inline-block; + width: 3rem; + height: 0.3rem; + border-radius: var(--radius-full); + background: var(--surface-active); + overflow: hidden; +} +.metric__fill { + display: block; + height: 100%; + background: var(--ink-faint); + transition: width var(--transition); +} +.metric--context.is-warning { color: var(--warning); } +.metric--context.is-warning .metric__fill { background: var(--warning); } +.metric--context.is-danger { color: var(--danger); } +.metric--context.is-danger .metric__fill { background: var(--danger); } + .reasoning__icon { color: var(--leaf); flex: none; } .reasoning__label { flex: 1; font-style: italic; } diff --git a/src/lembas/web/templates/chat/_message.html b/src/lembas/web/templates/chat/_message.html index cf6e51c..83756df 100644 --- a/src/lembas/web/templates/chat/_message.html +++ b/src/lembas/web/templates/chat/_message.html @@ -120,6 +120,10 @@
+ {# Counts as the reply is written. Everything is an estimate until the + usage chunk lands at the very end, and the chips say so. #} +
{% else %} {# Finished. Same order as the live view above -- thinking, then what it looked up, then the answer -- so a reply does not rearrange itself the @@ -177,6 +181,15 @@ leave an empty box under the file. #} {% endif %} + {% if not streaming and message.role == "assistant" and message.usage_json %} + {# Above the buttons, not among them: the actions row is things you press. #} +
+ {% with metrics = message.usage_json | metrics %} + {% include "chat/_metrics.html" %} + {% endwith %} +
+ {% endif %} + {% if not streaming %}