diff --git a/app/(chat)/api/chat/route.ts b/app/(chat)/api/chat/route.ts index 80c0ad3..7a62f5a 100644 --- a/app/(chat)/api/chat/route.ts +++ b/app/(chat)/api/chat/route.ts @@ -2,7 +2,6 @@ import { convertToModelMessages, createUIMessageStream, JsonToSseTransformStream, - type LanguageModelUsage, smoothStream, stepCountIs, streamText, @@ -39,11 +38,32 @@ import { ChatSDKError } from '@/lib/errors'; import type { ChatMessage } from '@/lib/types'; import type { ChatModel } from '@/lib/ai/models'; import type { VisibilityType } from '@/components/visibility-selector'; +import { unstable_cache as cache } from 'next/cache'; +import { fetchModels } from 'tokenlens/fetch'; +import { getUsage } from 'tokenlens/helpers'; +import type { ModelCatalog } from 'tokenlens/core'; +import type { AppUsage } from '@/lib/usage'; export const maxDuration = 60; let globalStreamContext: ResumableStreamContext | null = null; +const getTokenlensCatalog = cache( + async (): Promise => { + try { + return await fetchModels(); + } catch (err) { + console.warn( + 'TokenLens: catalog fetch failed, using default catalog', + err, + ); + return undefined; // tokenlens helpers will fall back to defaultCatalog + } + }, + ['tokenlens-catalog'], + { revalidate: 24 * 60 * 60 }, // 24 hours +); + export function getStreamContext() { if (!globalStreamContext) { try { @@ -151,7 +171,7 @@ export async function POST(request: Request) { const streamId = generateUUID(); await createStreamId({ streamId, chatId: id }); - let finalUsage: LanguageModelUsage | undefined; + let finalMergedUsage: AppUsage | undefined; const stream = createUIMessageStream({ execute: ({ writer: dataStream }) => { @@ -183,9 +203,31 @@ export async function POST(request: Request) { isEnabled: isProductionEnvironment, functionId: 'stream-text', }, - onFinish: ({ usage }) => { - finalUsage = usage; - dataStream.write({ type: 'data-usage', data: usage }); + onFinish: async ({ usage }) => { + try { + const providers = await getTokenlensCatalog(); + const modelId = + myProvider.languageModel(selectedChatModel).modelId; + if (!modelId) { + finalMergedUsage = usage; + dataStream.write({ type: 'data-usage', data: finalMergedUsage }); + return; + } + + if (!providers) { + finalMergedUsage = usage; + dataStream.write({ type: 'data-usage', data: finalMergedUsage }); + return; + } + + const summary = getUsage({ modelId, usage, providers }); + finalMergedUsage = { ...usage, ...summary, modelId } as AppUsage; + dataStream.write({ type: 'data-usage', data: finalMergedUsage }); + } catch (err) { + console.warn('TokenLens enrichment failed', err); + finalMergedUsage = usage; + dataStream.write({ type: 'data-usage', data: finalMergedUsage }); + } }, }); @@ -210,11 +252,11 @@ export async function POST(request: Request) { })), }); - if (finalUsage) { + if (finalMergedUsage) { try { await updateChatLastContextById({ chatId: id, - context: finalUsage, + context: finalMergedUsage, }); } catch (err) { console.warn('Unable to persist last usage for chat', id, err); diff --git a/components/chat.tsx b/components/chat.tsx index 5e6f3fb..86d51a8 100644 --- a/components/chat.tsx +++ b/components/chat.tsx @@ -1,6 +1,6 @@ 'use client'; -import { DefaultChatTransport, type LanguageModelUsage } from 'ai'; +import { DefaultChatTransport } from 'ai'; import { useChat } from '@ai-sdk/react'; import { useEffect, useState } from 'react'; import useSWR, { useSWRConfig } from 'swr'; @@ -21,6 +21,7 @@ import { useChatVisibility } from '@/hooks/use-chat-visibility'; import { useAutoResume } from '@/hooks/use-auto-resume'; import { ChatSDKError } from '@/lib/errors'; import type { Attachment, ChatMessage } from '@/lib/types'; +import type { AppUsage } from '@/lib/usage'; import { useDataStream } from './data-stream-provider'; import { AlertDialog, @@ -50,7 +51,7 @@ export function Chat({ isReadonly: boolean; session: Session; autoResume: boolean; - initialLastContext?: LanguageModelUsage; + initialLastContext?: AppUsage; }) { const { visibilityType } = useChatVisibility({ chatId: id, @@ -61,9 +62,7 @@ export function Chat({ const { setDataStream } = useDataStream(); const [input, setInput] = useState(''); - const [usage, setUsage] = useState( - initialLastContext, - ); + const [usage, setUsage] = useState(initialLastContext); const [showCreditCardAlert, setShowCreditCardAlert] = useState(false); const { @@ -96,9 +95,7 @@ export function Chat({ }), onData: (dataPart) => { setDataStream((ds) => (ds ? [...ds, dataPart] : [])); - if (dataPart.type === 'data-usage') { - setUsage(dataPart.data); - } + if (dataPart.type === 'data-usage') setUsage(dataPart.data); }, onFinish: () => { mutate(unstable_serialize(getChatHistoryPaginationKey)); diff --git a/components/elements/context.tsx b/components/elements/context.tsx index 72ea552..d290172 100644 --- a/components/elements/context.tsx +++ b/components/elements/context.tsx @@ -7,20 +7,13 @@ import { } from '@/components/ui/dropdown-menu'; 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'; import { Progress } from '@/components/ui/progress'; +import type { AppUsage } from '@/lib/usage'; 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; + usage?: AppUsage; }; const THOUSAND = 1000; @@ -34,57 +27,6 @@ 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}`; -}; - -const formatUSDFixed = (value?: number, decimals = 5) => { - if (value === undefined || !Number.isFinite(value)) return undefined; - return `$${Number(value).toFixed(decimals)}`; -}; - type ContextIconProps = { percent: number; // 0 - 100 }; @@ -96,7 +38,7 @@ export const ContextIcon = ({ percent }: ContextIconProps) => { return ( { ); }; -function TokensWithCost({ - tokens, - costText, -}: { - tokens?: number; - costText?: string; -}) { - return ( - - {tokens === undefined ? '—' : formatTokens(tokens)} - {costText ? ( - • {costText} - ) : null} - - ); -} function InfoRow({ label, @@ -156,126 +82,36 @@ function InfoRow({ costText?: string; }) { return ( -
+
{label} - +
+ + {tokens === undefined ? '—' : tokens.toLocaleString()} + + {costText !== undefined && costText !== null && !isNaN(parseFloat(costText)) && ( + + ${parseFloat(costText).toFixed(6)} + + )} +
); } -export const Context = ({ - className, - maxTokens, - usedTokens, - usage, - modelId, - ...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, - ); - - // used percent and used tokens to display (demo-aware) - const displayUsedTokens = safeUsed; - const usedPercent = - safeMax > 0 - ? Math.min( - PERCENT_MAX, - Math.max(0, (displayUsedTokens / safeMax) * PERCENT_MAX), - ) - : 0; - - const displayPct = formatPercent(Math.round(usedPercent * 10) / 10); - - const used = formatTokens(displayUsedTokens); - const total = formatTokens(safeMax); - - const uNorm = normalizeUsage(usage); - const uBreakdown = breakdownTokens(usage); - - const hasUsage = - !!usage && - ((uNorm.input ?? 0) > 0 || - (uNorm.output ?? 0) > 0 || - (uBreakdown.cacheReads ?? 0) > 0 || - (uBreakdown.cacheWrites ?? 0) > 0 || - (uBreakdown.reasoningTokens ?? 0) > 0); - - // Values to render in rows (demo or real) - const displayInput = uNorm.input; - const displayOutput = uNorm.output; - - // Per-segment costs - const inputCostText = modelId - ? formatUSDFixed( - estimateCost({ - modelId, - usage: { input: displayInput ?? 0, output: 0 }, - }).inputUSD, - ) - : undefined; - const outputCostText = modelId - ? formatUSDFixed( - estimateCost({ - modelId, - usage: { input: 0, output: displayOutput ?? 0 }, - }).outputUSD, - ) - : undefined; - // Not supported by tokenlens pricing hints; leave undefined so no bullet is shown - const cacheReadsTokens = uBreakdown.cacheReads ?? 0; - const cacheWritesTokens = uBreakdown.cacheWrites ?? 0; - const cacheReadsCostText = - modelId && cacheReadsTokens > 0 - ? formatUSDFixed( - estimateCost({ - modelId, - // Cast to any to support extended pricing fields provided by tokenlens - usage: { cacheReads: cacheReadsTokens } as any, - }).totalUSD, - ) - : undefined; - const cacheWritesCostText = - modelId && cacheWritesTokens > 0 - ? formatUSDFixed( - estimateCost({ - modelId, - usage: { cacheWrites: cacheWritesTokens } as any, - }).totalUSD, - ) - : undefined; - - const reasoningTokens = uBreakdown.reasoningTokens ?? 0; - let reasoningCostText: string | undefined; - if (modelId && reasoningTokens > 0) { - const est = estimateCost({ - modelId, - usage: { reasoningTokens }, - }).totalUSD; - // TokenLens does not provide reasoning pricing for some models. Show em dash when unknown. - reasoningCostText = - est && Number.isFinite(est) && est > 0 ? formatUSDFixed(est) : '—'; - } - - const costUSD = modelId - ? estimateCost({ - modelId, - usage: { input: displayInput ?? 0, output: displayOutput ?? 0 }, - }).totalUSD - : undefined; - const costText = formatUSDFixed(costUSD); - - const fmtOrUnknown = (n?: number) => - n === undefined ? '—' : formatTokens(n); - +export const Context = ({ className, usage, ...props }: ContextProps) => { + const used = usage?.totalTokens ?? 0; + const max = + usage?.context?.totalMax ?? + usage?.context?.combinedMax ?? + usage?.context?.inputMax; + const hasMax = typeof max === 'number' && Number.isFinite(max) && max > 0; + const usedPercent = hasMax ? Math.min(100, (used / max) * 100) : 0; return ( - +
-
- {displayPct} - {used} / {total} tokens +
+ {usedPercent.toFixed(1)}% + + {hasMax ? `${used} / ${max} tokens` : `${used} tokens`} +
- {hasUsage && uBreakdown.cacheReads && uBreakdown.cacheReads > 0 && ( + {usage?.cachedInputTokens && usage.cachedInputTokens > 0 && ( )} - {hasUsage && - uBreakdown.cacheWrites && - uBreakdown.cacheWrites > 0 && ( - - )} 0 ? reasoningTokens : undefined} - costText={reasoningCostText} + tokens={ + usage?.reasoningTokens && usage.reasoningTokens > 0 + ? usage.reasoningTokens + : undefined + } + costText={usage?.costUSD?.reasoningUSD?.toString()} /> - {costText && ( + {usage?.costUSD?.totalUSD !== undefined && ( <>
Total cost - {costText} +
+ + + {!isNaN(parseFloat(usage.costUSD.totalUSD.toString())) + ? `$${parseFloat(usage.costUSD.totalUSD.toString()).toFixed(6)}` + : '—' + } + +
)} @@ -343,4 +184,4 @@ export const Context = ({ ); -}; \ No newline at end of file +}; diff --git a/components/multimodal-input.tsx b/components/multimodal-input.tsx index d207346..2c8c8e9 100644 --- a/components/multimodal-input.tsx +++ b/components/multimodal-input.tsx @@ -1,6 +1,6 @@ 'use client'; -import type { LanguageModelUsage, UIMessage } from 'ai'; +import type { UIMessage } from 'ai'; import { useRef, useEffect, @@ -15,7 +15,13 @@ import { import { toast } from 'sonner'; import { useLocalStorage, useWindowSize } from 'usehooks-ts'; -import { ArrowUpIcon, PaperclipIcon, CpuIcon, StopIcon, ChevronDownIcon } from './icons'; +import { + ArrowUpIcon, + PaperclipIcon, + CpuIcon, + StopIcon, + ChevronDownIcon, +} from './icons'; import { PreviewAttachment } from './preview-attachment'; import { Button } from './ui/button'; import { SuggestedActions } from './suggested-actions'; @@ -37,10 +43,10 @@ import { ArrowDown } from 'lucide-react'; import { useScrollToBottom } from '@/hooks/use-scroll-to-bottom'; import type { VisibilityType } from './visibility-selector'; import type { Attachment, ChatMessage } from '@/lib/types'; +import type { AppUsage } from '@/lib/usage'; 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'; @@ -73,7 +79,7 @@ function PureMultimodalInput({ className?: string; selectedVisibilityType: VisibilityType; selectedModelId: string; - usage?: LanguageModelUsage; + usage?: AppUsage; }) { const textareaRef = useRef(null); const { width } = useWindowSize(); @@ -193,29 +199,11 @@ function PureMultimodalInput({ 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, }), - [contextMax, usedTokens, usage, modelResolver], + [usage], ); const handleFileChange = useCallback( @@ -253,7 +241,7 @@ function PureMultimodalInput({ }, [status, scrollToBottom]); return ( -
+
{!isAtBottom && (