mirror of
https://github.com/open-jarvis/OpenJarvis.git
synced 2026-07-30 02:42:16 +00:00
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:
@@ -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` });
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user