From 1e539b65b30aab844791e71d25da6e2e703ed0ea Mon Sep 17 00:00:00 2001 From: Nicklas Scharpff <64584979+xn1cklas@users.noreply.github.com> Date: Mon, 8 Sep 2025 16:56:10 +0200 Subject: [PATCH] feat: add tokenlens integration for language model usage tracking (#1161) --- app/(chat)/api/chat/route.ts | 7 + components/chat.tsx | 7 +- components/elements/context.tsx | 294 ++++++++++++++++++++++++++++++++ components/multimodal-input.tsx | 47 ++++- lib/ai/providers.ts | 2 +- lib/types.ts | 3 +- package.json | 1 + pnpm-lock.yaml | 8 + 8 files changed, 363 insertions(+), 6 deletions(-) create mode 100644 components/elements/context.tsx diff --git a/app/(chat)/api/chat/route.ts b/app/(chat)/api/chat/route.ts index 8e16929..002b9a4 100644 --- a/app/(chat)/api/chat/route.ts +++ b/app/(chat)/api/chat/route.ts @@ -2,6 +2,7 @@ import { convertToModelMessages, createUIMessageStream, JsonToSseTransformStream, + LanguageModelUsage, smoothStream, stepCountIs, streamText, @@ -149,6 +150,8 @@ export async function POST(request: Request) { const streamId = generateUUID(); await createStreamId({ streamId, chatId: id }); + let finalUsage: LanguageModelUsage | undefined; + const stream = createUIMessageStream({ execute: ({ writer: dataStream }) => { const result = streamText({ @@ -179,6 +182,10 @@ export async function POST(request: Request) { isEnabled: isProductionEnvironment, functionId: 'stream-text', }, + onFinish: ({ usage }) => { + finalUsage = usage; + dataStream.write({ type: 'data-usage', data: usage }); + }, }); result.consumeStream(); diff --git a/components/chat.tsx b/components/chat.tsx index f62a703..7c36512 100644 --- a/components/chat.tsx +++ b/components/chat.tsx @@ -1,6 +1,6 @@ 'use client'; -import { DefaultChatTransport } from 'ai'; +import { DefaultChatTransport, LanguageModelUsage } from 'ai'; import { useChat } from '@ai-sdk/react'; import { useEffect, useState } from 'react'; import useSWR, { useSWRConfig } from 'swr'; @@ -49,6 +49,7 @@ export function Chat({ const { setDataStream } = useDataStream(); const [input, setInput] = useState(''); + const [usage, setUsage] = useState(undefined); const { messages, @@ -80,6 +81,9 @@ export function Chat({ }), onData: (dataPart) => { setDataStream((ds) => (ds ? [...ds, dataPart] : [])); + if (dataPart.type === 'data-usage') { + setUsage(dataPart.data); + } }, onFinish: () => { mutate(unstable_serialize(getChatHistoryPaginationKey)); @@ -163,6 +167,7 @@ export function Chat({ sendMessage={sendMessage} selectedVisibilityType={visibilityType} selectedModelId={initialChatModel} + usage={usage} /> )} diff --git a/components/elements/context.tsx b/components/elements/context.tsx new file mode 100644 index 0000000..ea938a8 --- /dev/null +++ b/components/elements/context.tsx @@ -0,0 +1,294 @@ +'use client'; + +import { + HoverCard, + HoverCardContent, + HoverCardTrigger, +} from '@/components/ui/hover-card'; +import { cn } from '@/lib/utils'; +import type { ComponentProps } from 'react'; +import type { LanguageModelUsage } from 'ai'; +import { breakdownTokens, estimateCost, normalizeUsage } from 'tokenlens'; +import { Separator } from '@/components/ui/separator'; + +export type ContextProps = ComponentProps<'button'> & { + /** Total context window size in tokens */ + maxTokens: number; + /** Tokens used so far */ + usedTokens: number; + /** Optional full usage payload to enable breakdown view */ + usage?: LanguageModelUsage | undefined; + /** Optional model id (canonical or alias) to compute cost */ + modelId?: string; + /** Show token breakdown and optional cost inside hover */ + showBreakdown?: boolean; +}; + +const THOUSAND = 1000; +const MILLION = 1_000_000; +const BILLION = 1_000_000_000; +const PERCENT_MAX = 100; + +// Lucide CircleIcon geometry +const ICON_VIEWBOX = 24; +const ICON_CENTER = 12; +const ICON_RADIUS = 10; +const ICON_STROKE_WIDTH = 2; + +const formatTokens = (tokens?: number) => { + if (tokens === undefined) { + return; + } + if (!Number.isFinite(tokens)) { + return; + } + const abs = Math.abs(tokens); + if (abs < THOUSAND) { + return `${tokens}`; + } + if (abs < MILLION) { + return `${(tokens / THOUSAND).toFixed(1)}K`; + } + if (abs < BILLION) { + return `${(tokens / MILLION).toFixed(1)}M`; + } + return `${(tokens / BILLION).toFixed(1)}B`; +}; + +const formatPercent = (value: number) => { + if (!Number.isFinite(value)) { + return '0%'; + } + const rounded = Math.round(value * 10) / 10; + return Number.isInteger(rounded) + ? `${Math.trunc(rounded)}%` + : `${rounded.toFixed(1)}%`; +}; + +const formatUSD = (value?: number) => { + if (value === undefined || !Number.isFinite(value)) return undefined; + const abs = Math.abs(value); + // Finer precision for very small amounts common in LLM pricing + let decimals = 2; + if (abs < 0.001) decimals = 5; + else if (abs < 0.01) decimals = 4; + else if (abs < 0.1) decimals = 3; + else if (abs < 10) decimals = 2; + else decimals = 1; + const text = value.toFixed(decimals); + // Trim trailing zeros/decimal if not needed (e.g., 1.2300 -> 1.23, 2.0 -> 2) + const trimmed = text.replace(/\.0+$/, '').replace(/(\.\d*?)0+$/, '$1'); + return `$${trimmed}`; +}; + +type ContextIconProps = { + percent: number; // 0 - 100 +}; + +export const ContextIcon = ({ percent }: ContextIconProps) => { + const radius = ICON_RADIUS; + const circumference = 2 * Math.PI * radius; + const dashOffset = circumference * (1 - percent / PERCENT_MAX); + + return ( + + + + + ); +}; + +export const Context = ({ + className, + maxTokens, + usedTokens, + usage, + modelId, + showBreakdown, + ...props +}: ContextProps) => { + const safeMax = Math.max(0, Number.isFinite(maxTokens) ? maxTokens : 0); + const safeUsed = Math.min( + Math.max(0, Number.isFinite(usedTokens) ? usedTokens : 0), + safeMax, + ); + const usedPercent = + safeMax > 0 + ? Math.min(PERCENT_MAX, Math.max(0, (safeUsed / safeMax) * PERCENT_MAX)) + : 0; + + const displayPct = formatPercent(Math.round(usedPercent * 10) / 10); + + const used = formatTokens(safeUsed); + const total = formatTokens(safeMax); + + const uNorm = normalizeUsage(usage as any); + const uBreakdown = breakdownTokens(usage as any); + const costUSD = modelId + ? estimateCost({ modelId, usage: uNorm }).totalUSD + : undefined; + const costText = formatUSD(costUSD); + + const segInput = Math.max(0, uNorm.input ?? 0); + const segOutput = Math.max(0, uNorm.output ?? 0); + const segCacheR = Math.max(0, uBreakdown.cacheReads ?? 0); + const segCacheW = Math.max(0, uBreakdown.cacheWrites ?? 0); + const denom = safeMax > 0 ? safeMax : 1; + const w = (n: number) => + `${Math.min(100, Math.max(0, (n / denom) * 100)).toFixed(2)}%`; + const fmtOrUnknown = (n?: number) => + n === undefined ? '—' : formatTokens(n); + return ( + + + + + +
+

+ {displayPct} • {used} / {total} tokens + {costText ? ( + • {costText} + ) : null} +

+ {true && ( +
+
+
+
+
+
+
+
+
+ + + Cache Hits + + {fmtOrUnknown(uBreakdown.cacheReads)} +
+
+ + + Cache Writes + + {fmtOrUnknown(uBreakdown.cacheWrites)} +
+
+ + + Input + + {formatTokens(uNorm.input)} +
+
+ + + Output + + {formatTokens(uNorm.output)} +
+
+
+ )} + {showBreakdown && ( +
+
+ Cache Hits + {fmtOrUnknown(uBreakdown.cacheReads)} +
+
+ Cache Writes + {fmtOrUnknown(uBreakdown.cacheWrites)} +
+
+ Input + {formatTokens(uNorm.input)} +
+
+ Output + {formatTokens(uNorm.output)} +
+ {costText && ( + <> + +
+ Total cost + {costText} +
+ + )} +
+ )} +
+ + + ); +}; diff --git a/components/multimodal-input.tsx b/components/multimodal-input.tsx index e314e43..8477670 100644 --- a/components/multimodal-input.tsx +++ b/components/multimodal-input.tsx @@ -1,6 +1,6 @@ 'use client'; -import type { UIMessage } from 'ai'; +import type { LanguageModelUsage, UIMessage } from 'ai'; import { useRef, useEffect, @@ -10,6 +10,7 @@ import { type SetStateAction, type ChangeEvent, memo, + useMemo, } from 'react'; import { toast } from 'sonner'; import { useLocalStorage, useWindowSize } from 'usehooks-ts'; @@ -39,6 +40,9 @@ import type { Attachment, ChatMessage } from '@/lib/types'; import { chatModels } from '@/lib/ai/models'; import { saveChatModelAsCookie } from '@/app/(chat)/actions'; import { startTransition } from 'react'; +import { getContextWindow, normalizeUsage } from 'tokenlens'; +import { Context } from './elements/context'; +import { myProvider } from '@/lib/ai/providers'; function PureMultimodalInput({ chatId, @@ -54,6 +58,7 @@ function PureMultimodalInput({ className, selectedVisibilityType, selectedModelId, + usage, }: { chatId: string; input: string; @@ -68,6 +73,7 @@ function PureMultimodalInput({ className?: string; selectedVisibilityType: VisibilityType; selectedModelId: string; + usage?: LanguageModelUsage; }) { const textareaRef = useRef(null); const { width } = useWindowSize(); @@ -183,6 +189,36 @@ function PureMultimodalInput({ } }; + const modelResolver = useMemo(() => { + return myProvider.languageModel(selectedModelId); + }, [selectedModelId]); + + const contextMax = useMemo(() => { + // Resolve from selected model; stable across chunks. + const cw = getContextWindow(modelResolver.modelId); + return cw.combinedMax ?? cw.inputMax ?? 0; + }, [modelResolver]); + + const usedTokens = useMemo(() => { + // Prefer explicit usage data part captured via onData + if (!usage) return 0; // update only when final usage arrives + const n = normalizeUsage(usage); + return typeof n.total === 'number' + ? n.total + : (n.input ?? 0) + (n.output ?? 0); + }, [usage]); + + const contextProps = useMemo( + () => ({ + maxTokens: contextMax, + usedTokens, + usage, + modelId: modelResolver.modelId, + showBreakdown: true as const, + }), + [contextMax, usedTokens, usage, modelResolver], + ); + const handleFileChange = useCallback( async (event: ChangeEvent) => { const files = Array.from(event.target.files || []); @@ -323,8 +359,13 @@ function PureMultimodalInput({ /> - + + {status === 'submitted' ? ( @@ -367,7 +408,7 @@ function PureAttachmentsButton({ selectedModelId: string; }) { const isReasoningModel = selectedModelId === 'chat-model-reasoning'; - + return (