feat: show token usage and TPS stats below each assistant message
This commit is contained in:
@@ -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>
|
||||
)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user