feat: add tokenlens integration for language model usage tracking (#1161)
This commit is contained in:
Nicklas Scharpff
GitHub
parent
7327c9a18a
commit
1e539b65b3
8 changed files with 363 additions and 6 deletions
|
|
@ -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();
|
||||
|
|
|
|||
|
|
@ -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<string>('');
|
||||
const [usage, setUsage] = useState<LanguageModelUsage | undefined>(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}
|
||||
/>
|
||||
)}
|
||||
</div>
|
||||
|
|
|
|||
294
components/elements/context.tsx
Normal file
294
components/elements/context.tsx
Normal file
|
|
@ -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 (
|
||||
<svg
|
||||
aria-label={`${formatPercent(percent)} of model context used`}
|
||||
height="20"
|
||||
role="img"
|
||||
style={{ color: 'currentcolor' }}
|
||||
viewBox={`0 0 ${ICON_VIEWBOX} ${ICON_VIEWBOX}`}
|
||||
width="20"
|
||||
>
|
||||
<circle
|
||||
cx={ICON_CENTER}
|
||||
cy={ICON_CENTER}
|
||||
fill="none"
|
||||
opacity="0.25"
|
||||
r={radius}
|
||||
stroke="currentColor"
|
||||
strokeWidth={ICON_STROKE_WIDTH}
|
||||
/>
|
||||
<circle
|
||||
cx={ICON_CENTER}
|
||||
cy={ICON_CENTER}
|
||||
fill="none"
|
||||
opacity="0.7"
|
||||
r={radius}
|
||||
stroke="currentColor"
|
||||
strokeDasharray={`${circumference} ${circumference}`}
|
||||
strokeDashoffset={dashOffset}
|
||||
strokeLinecap="round"
|
||||
strokeWidth={ICON_STROKE_WIDTH}
|
||||
transform={`rotate(-90 ${ICON_CENTER} ${ICON_CENTER})`}
|
||||
/>
|
||||
</svg>
|
||||
);
|
||||
};
|
||||
|
||||
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 (
|
||||
<HoverCard closeDelay={100} openDelay={100}>
|
||||
<HoverCardTrigger asChild>
|
||||
<button
|
||||
className={cn(
|
||||
'inline-flex select-none items-center gap-2 rounded-md px-2.5 py-1 text-sm',
|
||||
'bg-background text-foreground',
|
||||
)}
|
||||
type="button"
|
||||
{...props}
|
||||
>
|
||||
<span className="font-medium text-muted-foreground">
|
||||
{displayPct}
|
||||
</span>
|
||||
<ContextIcon percent={usedPercent} />
|
||||
</button>
|
||||
</HoverCardTrigger>
|
||||
<HoverCardContent align="center" className="w-fit p-3">
|
||||
<div className="min-w-[240px] space-y-2">
|
||||
<p className="text-center text-sm">
|
||||
{displayPct} • {used} / {total} tokens
|
||||
{costText ? (
|
||||
<span className="ml-1 text-muted-foreground">• {costText}</span>
|
||||
) : null}
|
||||
</p>
|
||||
{true && (
|
||||
<div className="space-y-2">
|
||||
<div className="h-2 w-full overflow-hidden rounded-full bg-muted">
|
||||
<div
|
||||
className="h-full"
|
||||
style={{
|
||||
width: w(segCacheR),
|
||||
background: 'var(--chart-2)',
|
||||
opacity: 0.9,
|
||||
}}
|
||||
/>
|
||||
<div
|
||||
className="h-full"
|
||||
style={{
|
||||
width: w(segCacheW),
|
||||
background: 'var(--chart-4)',
|
||||
opacity: 0.9,
|
||||
}}
|
||||
/>
|
||||
<div
|
||||
className="h-full"
|
||||
style={{
|
||||
width: w(segInput),
|
||||
background: 'var(--chart-1)',
|
||||
opacity: 0.9,
|
||||
}}
|
||||
/>
|
||||
<div
|
||||
className="h-full"
|
||||
style={{
|
||||
width: w(segOutput),
|
||||
background: 'var(--chart-3)',
|
||||
opacity: 0.9,
|
||||
}}
|
||||
/>
|
||||
</div>
|
||||
<div className="grid grid-cols-2 gap-x-4 gap-y-1 text-xs">
|
||||
<div className="flex items-center justify-between">
|
||||
<span className="flex items-center gap-2 text-muted-foreground">
|
||||
<span className="inline-block size-2 rounded-sm bg-chart-1" />
|
||||
Cache Hits
|
||||
</span>
|
||||
<span>{fmtOrUnknown(uBreakdown.cacheReads)}</span>
|
||||
</div>
|
||||
<div className="flex items-center justify-between">
|
||||
<span className="flex items-center gap-2 text-muted-foreground">
|
||||
<span className="inline-block size-2 rounded-sm bg-chart-2" />
|
||||
Cache Writes
|
||||
</span>
|
||||
<span>{fmtOrUnknown(uBreakdown.cacheWrites)}</span>
|
||||
</div>
|
||||
<div className="flex items-center justify-between">
|
||||
<span className="flex items-center gap-2 text-muted-foreground">
|
||||
<span className="inline-block size-2 rounded-sm bg-chart-3" />
|
||||
Input
|
||||
</span>
|
||||
<span>{formatTokens(uNorm.input)}</span>
|
||||
</div>
|
||||
<div className="flex items-center justify-between">
|
||||
<span className="flex items-center gap-2 text-muted-foreground">
|
||||
<span className="inline-block size-2 rounded-sm bg-chart-4" />
|
||||
Output
|
||||
</span>
|
||||
<span>{formatTokens(uNorm.output)}</span>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
)}
|
||||
{showBreakdown && (
|
||||
<div className="mt-1 space-y-1">
|
||||
<div className="flex items-center justify-between text-xs">
|
||||
<span className="text-muted-foreground">Cache Hits</span>
|
||||
<span>{fmtOrUnknown(uBreakdown.cacheReads)}</span>
|
||||
</div>
|
||||
<div className="flex items-center justify-between text-xs">
|
||||
<span className="text-muted-foreground">Cache Writes</span>
|
||||
<span>{fmtOrUnknown(uBreakdown.cacheWrites)}</span>
|
||||
</div>
|
||||
<div className="flex items-center justify-between text-xs">
|
||||
<span className="text-muted-foreground">Input</span>
|
||||
<span>{formatTokens(uNorm.input)}</span>
|
||||
</div>
|
||||
<div className="flex items-center justify-between text-xs">
|
||||
<span className="text-muted-foreground">Output</span>
|
||||
<span>{formatTokens(uNorm.output)}</span>
|
||||
</div>
|
||||
{costText && (
|
||||
<>
|
||||
<Separator className="mt-1" />
|
||||
<div className="flex items-center justify-between pt-1 text-xs">
|
||||
<span className="text-muted-foreground">Total cost</span>
|
||||
<span>{costText}</span>
|
||||
</div>
|
||||
</>
|
||||
)}
|
||||
</div>
|
||||
)}
|
||||
</div>
|
||||
</HoverCardContent>
|
||||
</HoverCard>
|
||||
);
|
||||
};
|
||||
|
|
@ -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<HTMLTextAreaElement>(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<HTMLInputElement>) => {
|
||||
const files = Array.from(event.target.files || []);
|
||||
|
|
@ -323,8 +359,13 @@ function PureMultimodalInput({
|
|||
/>
|
||||
<PromptInputToolbar className="px-3 py-2 !border-t-0 !border-top-0 shadow-none dark:!border-transparent dark:border-0">
|
||||
<PromptInputTools className="gap-2">
|
||||
<AttachmentsButton fileInputRef={fileInputRef} status={status} selectedModelId={selectedModelId} />
|
||||
<AttachmentsButton
|
||||
fileInputRef={fileInputRef}
|
||||
status={status}
|
||||
selectedModelId={selectedModelId}
|
||||
/>
|
||||
<ModelSelectorCompact selectedModelId={selectedModelId} />
|
||||
<Context {...contextProps} />
|
||||
</PromptInputTools>
|
||||
{status === 'submitted' ? (
|
||||
<StopButton stop={stop} setMessages={setMessages} />
|
||||
|
|
@ -367,7 +408,7 @@ function PureAttachmentsButton({
|
|||
selectedModelId: string;
|
||||
}) {
|
||||
const isReasoningModel = selectedModelId === 'chat-model-reasoning';
|
||||
|
||||
|
||||
return (
|
||||
<Button
|
||||
data-testid="attachments-button"
|
||||
|
|
|
|||
|
|
@ -25,7 +25,7 @@ export const myProvider = isTestEnvironment
|
|||
languageModels: {
|
||||
'chat-model': gateway.languageModel('xai/grok-2-vision-1212'),
|
||||
'chat-model-reasoning': wrapLanguageModel({
|
||||
model: gateway.languageModel('xai/grok-3-mini-beta'),
|
||||
model: gateway.languageModel('xai/grok-3-mini'),
|
||||
middleware: extractReasoningMiddleware({ tagName: 'think' }),
|
||||
}),
|
||||
'title-model': gateway.languageModel('xai/grok-2-1212'),
|
||||
|
|
|
|||
|
|
@ -3,7 +3,7 @@ import type { getWeather } from './ai/tools/get-weather';
|
|||
import type { createDocument } from './ai/tools/create-document';
|
||||
import type { updateDocument } from './ai/tools/update-document';
|
||||
import type { requestSuggestions } from './ai/tools/request-suggestions';
|
||||
import type { InferUITool, UIMessage } from 'ai';
|
||||
import type { InferUITool, LanguageModelUsage, UIMessage } from 'ai';
|
||||
|
||||
import type { ArtifactKind } from '@/components/artifact';
|
||||
import type { Suggestion } from './db/schema';
|
||||
|
|
@ -42,6 +42,7 @@ export type CustomUIDataTypes = {
|
|||
kind: ArtifactKind;
|
||||
clear: null;
|
||||
finish: null;
|
||||
usage: LanguageModelUsage;
|
||||
};
|
||||
|
||||
export type ChatMessage = UIMessage<
|
||||
|
|
|
|||
|
|
@ -95,6 +95,7 @@
|
|||
"swr": "^2.2.5",
|
||||
"tailwind-merge": "^2.5.2",
|
||||
"tailwindcss-animate": "^1.0.7",
|
||||
"tokenlens": "^1.0.0",
|
||||
"use-stick-to-bottom": "^1.1.1",
|
||||
"usehooks-ts": "^3.1.0",
|
||||
"zod": "^3.25.76"
|
||||
|
|
|
|||
8
pnpm-lock.yaml
generated
8
pnpm-lock.yaml
generated
|
|
@ -236,6 +236,9 @@ importers:
|
|||
tailwindcss-animate:
|
||||
specifier: ^1.0.7
|
||||
version: 1.0.7(tailwindcss@3.4.17)
|
||||
tokenlens:
|
||||
specifier: ^1.0.0
|
||||
version: 1.0.0
|
||||
use-stick-to-bottom:
|
||||
specifier: ^1.1.1
|
||||
version: 1.1.1(react@19.0.0-rc-45804af1-20241021)
|
||||
|
|
@ -4709,6 +4712,9 @@ packages:
|
|||
resolution: {integrity: sha512-65P7iz6X5yEr1cwcgvQxbbIw7Uk3gOy5dIdtZ4rDveLqhrdJP+Li/Hx6tyK0NEb+2GCyneCMJiGqrADCSNk8sQ==}
|
||||
engines: {node: '>=8.0'}
|
||||
|
||||
tokenlens@1.0.0:
|
||||
resolution: {integrity: sha512-IslxKcdsPQv4yTUgzeAe9UoEuvXUW6BZyKvKc4N9D7IWnyUBvSTJ3hq8Wb13fM8CkgaS/0k28vpJIQiVQJrpVA==}
|
||||
|
||||
trim-lines@3.0.1:
|
||||
resolution: {integrity: sha512-kRj8B+YHZCc9kQYdWfJB2/oUl9rA99qbowYYBtr4ui4mZyAQ2JpvVBd/6U2YloATfqBhBTSMhTpgBHtU0Mf3Rg==}
|
||||
|
||||
|
|
@ -9769,6 +9775,8 @@ snapshots:
|
|||
dependencies:
|
||||
is-number: 7.0.0
|
||||
|
||||
tokenlens@1.0.0: {}
|
||||
|
||||
trim-lines@3.0.1: {}
|
||||
|
||||
trough@2.2.0: {}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue