feat: show token usage and TPS stats below each assistant message

This commit is contained in:
2026-06-29 13:08:47 +02:00
parent f977c484b0
commit 3fcc933999

View File

@@ -1,10 +1,19 @@
import { useState, useEffect, useRef } from 'react'
import { getToken } from '../lib/api'
interface Stats {
promptTokens: number
completionTokens: number
totalTokens: number
tps: number // completion tokens / elapsed seconds
elapsedMs: number
}
interface Message {
role: 'user' | 'assistant'
content: string
pending?: boolean
stats?: Stats
}
interface ModelOption { id: string }
@@ -47,6 +56,9 @@ export default function Chat() {
abortRef.current = abort
setStreaming(true)
const startMs = performance.now()
let firstTokenMs: number | null = null
try {
const res = await fetch('/api/chat/completions', {
method: 'POST',
@@ -58,6 +70,7 @@ export default function Chat() {
body: JSON.stringify({
model,
stream: true,
stream_options: { include_usage: true },
messages: next.map(m => ({ role: m.role, content: m.content })),
}),
})
@@ -68,6 +81,7 @@ export default function Chat() {
const decoder = new TextDecoder()
let buf = ''
let reply = ''
let usage: { prompt_tokens: number; completion_tokens: number; total_tokens: number } | null = null
while (true) {
const { done, value } = await reader.read()
@@ -78,16 +92,35 @@ export default function Chat() {
for (const line of lines) {
if (!line.startsWith('data:')) continue
const data = line.slice(5).trim()
if (data === '[DONE]') break
if (data === '[DONE]') continue
try {
const chunk = JSON.parse(data)
// Capture usage from the final chunk (stream_options.include_usage)
if (chunk.usage) usage = chunk.usage
const delta = chunk.choices?.[0]?.delta?.content ?? chunk.choices?.[0]?.delta?.reasoning_content ?? ''
if (delta) { reply += delta; setMessages([...next, { role: 'assistant', content: reply }]) }
if (delta) {
if (firstTokenMs === null) firstTokenMs = performance.now()
reply += delta
setMessages([...next, { role: 'assistant', content: reply }])
}
} catch {}
}
}
setMessages([...next, { role: 'assistant', content: reply || '(empty response)' }])
const endMs = performance.now()
const elapsedMs = endMs - (firstTokenMs ?? startMs)
const completionTokens = usage?.completion_tokens ?? 0
const tps = elapsedMs > 0 && completionTokens > 0 ? (completionTokens / (elapsedMs / 1000)) : 0
const stats: Stats | undefined = usage ? {
promptTokens: usage.prompt_tokens,
completionTokens: usage.completion_tokens,
totalTokens: usage.total_tokens,
tps,
elapsedMs: endMs - startMs,
} : undefined
setMessages([...next, { role: 'assistant', content: reply || '(empty response)', stats }])
} catch (e: any) {
if (e.name !== 'AbortError') { setError(e.message); setMessages(next) }
else setMessages(prev => prev.map((m, i) => i === prev.length - 1 ? { ...m, pending: false } : m))
@@ -141,19 +174,21 @@ export default function Chat() {
display: 'flex', alignItems: 'center', justifyContent: 'center',
fontSize: 11, fontWeight: 600, color: 'rgba(237,237,240,0.7)',
}}>{m.role === 'user' ? 'U' : 'AI'}</div>
<div style={{
maxWidth: '80%',
background: m.role === 'user' ? 'rgba(124,58,237,0.15)' : '#17171c',
border: `1px solid ${m.role === 'user' ? 'rgba(124,58,237,0.25)' : 'rgba(255,255,255,0.07)'}`,
borderRadius: m.role === 'user' ? '12px 4px 12px 12px' : '4px 12px 12px 12px',
padding: '10px 14px', fontSize: 13.5, lineHeight: 1.6, color: '#ededf0',
whiteSpace: 'pre-wrap', wordBreak: 'break-word',
}}>
{m.content || (m.pending
? <span style={{ display: 'inline-flex', gap: 4, alignItems: 'center' }}>
{[0,1,2].map(j => <span key={j} style={{ width: 5, height: 5, borderRadius: '50%', background: 'rgba(237,237,240,0.3)', display: 'inline-block' }}/>)}
</span>
: '')}
<div style={{ maxWidth: '80%', display: 'flex', flexDirection: 'column', gap: 5 }}>
<div style={{
background: m.role === 'user' ? 'rgba(124,58,237,0.15)' : '#17171c',
border: `1px solid ${m.role === 'user' ? 'rgba(124,58,237,0.25)' : 'rgba(255,255,255,0.07)'}`,
borderRadius: m.role === 'user' ? '12px 4px 12px 12px' : '4px 12px 12px 12px',
padding: '10px 14px', fontSize: 13.5, lineHeight: 1.6, color: '#ededf0',
whiteSpace: 'pre-wrap', wordBreak: 'break-word',
}}>
{m.content || (m.pending
? <span style={{ display: 'inline-flex', gap: 4, alignItems: 'center' }}>
{[0,1,2].map(j => <span key={j} style={{ width: 5, height: 5, borderRadius: '50%', background: 'rgba(237,237,240,0.3)', display: 'inline-block' }}/>)}
</span>
: '')}
</div>
{m.stats && <StatsBar stats={m.stats} />}
</div>
</div>
))}
@@ -197,3 +232,23 @@ export default function Chat() {
</div>
)
}
function StatsBar({ stats }: { stats: Stats }) {
const c = 'rgba(237,237,240,0.3)'
const items = [
{ label: 'prompt', value: stats.promptTokens.toLocaleString() + ' tok' },
{ label: 'completion', value: stats.completionTokens.toLocaleString() + ' tok' },
{ label: 'total', value: stats.totalTokens.toLocaleString() + ' tok' },
...(stats.tps > 0 ? [{ label: 'speed', value: stats.tps.toFixed(1) + ' tok/s' }] : []),
{ label: 'time', value: (stats.elapsedMs / 1000).toFixed(2) + 's' },
]
return (
<div style={{ display: 'flex', gap: 10, flexWrap: 'wrap' }}>
{items.map(({ label, value }) => (
<span key={label} style={{ fontSize: 11, fontFamily: "'DM Mono',monospace", color: c }}>
<span style={{ color: 'rgba(237,237,240,0.18)' }}>{label} </span>{value}
</span>
))}
</div>
)
}