feat: display thinking tokens in X-Ray footer and stream usage data

- Ollama engine: capture eval_count and prompt_eval_count from the
  streaming final chunk (includes thinking/reasoning tokens).
- Routes: include usage data in the SSE finish chunk so the frontend
  gets accurate token counts even in streaming mode.
- X-Ray footer: show estimated thinking tokens when completion_tokens
  significantly exceeds visible output (e.g. "1695 generated · 11 prompt
  · ~1680 thinking").
This commit is contained in:
Jon Saad-Falcon
2026-03-14 14:55:42 -07:00
parent 34000aea01
commit 3cf4bbae30
3 changed files with 42 additions and 4 deletions
+13 -1
View File
@@ -35,7 +35,19 @@ export function XRayFooter({ usage, telemetry }: Props) {
rows.push({ label: 'Engine', value: `${telemetry.engine}${modelDetail ? ` (${modelDetail})` : ''}` });
}
if (usage) {
rows.push({ label: 'Tokens', value: `${usage.completion_tokens} generated \u00B7 ${usage.prompt_tokens} prompt` });
const tokenParts = [`${usage.completion_tokens} generated`, `${usage.prompt_tokens} prompt`];
// Estimate thinking tokens: if total generated >> visible output, the
// difference is internal reasoning (e.g. Qwen3.5 thinking mode).
if (telemetry?.tokens_per_sec && telemetry.total_ms && usage.completion_tokens > 50) {
const visibleEstimate = Math.ceil(
(telemetry.total_ms / 1000) * telemetry.tokens_per_sec,
);
if (usage.completion_tokens > visibleEstimate * 1.5) {
const thinking = usage.completion_tokens - visibleEstimate;
tokenParts.push(`~${thinking} thinking`);
}
}
rows.push({ label: 'Tokens', value: tokenParts.join(' \u00B7 ') });
}
if (telemetry?.tokens_per_sec) {
rows.push({ label: 'Speed', value: `${Math.round(telemetry.tokens_per_sec)} tok/s` });
+11
View File
@@ -41,6 +41,8 @@ class OllamaEngine(InferenceEngine):
host = env_host or self._DEFAULT_HOST
self._host = host.rstrip("/")
self._client = httpx.Client(base_url=self._host, timeout=timeout)
# Last stream usage — captured from Ollama's final chunk
self._last_stream_usage: Dict[str, int] = {}
def generate(
self,
@@ -174,6 +176,15 @@ class OllamaEngine(InferenceEngine):
if content:
yield content
if chunk.get("done", False):
# Capture usage from final chunk
self._last_stream_usage = {
"prompt_tokens": chunk.get("prompt_eval_count", 0),
"completion_tokens": chunk.get("eval_count", 0),
"total_tokens": (
chunk.get("prompt_eval_count", 0)
+ chunk.get("eval_count", 0)
),
}
break
except (httpx.ConnectError, httpx.TimeoutException) as exc:
raise EngineConnectionError(
+18 -3
View File
@@ -222,8 +222,9 @@ async def _handle_stream(engine, model: str, req: ChatCompletionRequest):
yield "data: [DONE]\n\n"
return
# Send finish chunk
finish_chunk = ChatCompletionChunk(
# Send finish chunk with usage data if available
import json as _json
finish_data = ChatCompletionChunk(
id=chunk_id,
model=model,
choices=[StreamChoice(
@@ -231,7 +232,21 @@ async def _handle_stream(engine, model: str, req: ChatCompletionRequest):
finish_reason="stop",
)],
)
yield f"data: {finish_chunk.model_dump_json()}\n\n"
finish_dict = _json.loads(finish_data.model_dump_json())
# Pull usage from the engine if it tracked it during streaming
raw_engine = engine
# Unwrap InstrumentedEngine if present
if hasattr(raw_engine, "_inner"):
raw_engine = raw_engine._inner
# Unwrap MultiEngine if present
if hasattr(raw_engine, "_engine_for"):
raw_engine = raw_engine._engine_for(model)
stream_usage = getattr(raw_engine, "_last_stream_usage", None)
if isinstance(stream_usage, dict) and stream_usage.get("total_tokens", 0) > 0:
finish_dict["usage"] = stream_usage
yield f"data: {_json.dumps(finish_dict)}\n\n"
yield "data: [DONE]\n\n"
return StreamingResponse(