feat: add tokenlens integration for language model usage tracking (#1161)

2025-09-08 16:56:10 +02:00 · 2025-09-08 16:56:10 +02:00 · 1e539b65b3
commit 1e539b65b3
parent 7327c9a18a
8 changed files with 363 additions and 6 deletions
--- 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();
--- 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<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>
--- a/components/elements/context.tsx
+++ 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 (
+    <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>
+  );
+};
--- 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<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} />
--- a/lib/ai/providers.ts
+++ b/lib/ai/providers.ts
@ -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'),
--- a/lib/types.ts
+++ b/lib/types.ts
@ -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<
--- a/package.json
+++ b/package.json
@ -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"
--- a/pnpm-lock.yaml
+++ b/pnpm-lock.yaml
@ -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: {}