diff --git a/src/lib/components/chat/Chat.svelte b/src/lib/components/chat/Chat.svelte index a912960bff..2a47d4f926 100644 --- a/src/lib/components/chat/Chat.svelte +++ b/src/lib/components/chat/Chat.svelte @@ -61,7 +61,8 @@ processDetails, removeAllDetails, getCodeBlockContents, - displayFileHandler + displayFileHandler, + getUsageTokenCount } from '$lib/utils'; import { AudioQueue } from '$lib/utils/audio'; import { createTemporaryChatId, isTemporaryChatId } from '$lib/utils/chatId'; @@ -287,13 +288,10 @@ for (let idx = activeMessages.length - 1; idx >= 0; idx -= 1) { const usage = activeMessages[idx]?.usage ?? activeMessages[idx]?.info?.usage; - const inputTokens = usage?.input_tokens ?? usage?.prompt_tokens; - if (inputTokens) { + const usageTokens = getUsageTokenCount(usage); + if (usageTokens) { hasUsageCheckpoint = true; - estimatedTokens = - Number(inputTokens || 0) + - Number(usage.output_tokens ?? usage.completion_tokens ?? 0) + - estimateMessagesTokens(activeMessages.slice(idx + 1)); + estimatedTokens = usageTokens + estimateMessagesTokens(activeMessages.slice(idx + 1)); break; } } diff --git a/src/lib/components/chat/MessageInput.svelte b/src/lib/components/chat/MessageInput.svelte index cee85324ed..3eeb6173a8 100644 --- a/src/lib/components/chat/MessageInput.svelte +++ b/src/lib/components/chat/MessageInput.svelte @@ -52,6 +52,7 @@ getCurrentDateTime, getFormattedDate, getFormattedTime, + getUsageTokenCount, getUserPosition, getUserTimezone, getWeekday @@ -482,13 +483,10 @@ for (let idx = activeMessages.length - 1; idx >= 0; idx -= 1) { const usage = activeMessages[idx]?.usage ?? activeMessages[idx]?.info?.usage; - const inputTokens = usage?.input_tokens ?? usage?.prompt_tokens; - if (inputTokens) { + const usageTokens = getUsageTokenCount(usage); + if (usageTokens) { hasUsageCheckpoint = true; - estimatedTokens = - Number(inputTokens || 0) + - Number(usage.output_tokens ?? usage.completion_tokens ?? 0) + - estimateMessagesTokens(activeMessages.slice(idx + 1)); + estimatedTokens = usageTokens + estimateMessagesTokens(activeMessages.slice(idx + 1)); break; } } diff --git a/src/lib/utils/index.ts b/src/lib/utils/index.ts index fa727dfda1..0de6c55308 100644 --- a/src/lib/utils/index.ts +++ b/src/lib/utils/index.ts @@ -1417,6 +1417,31 @@ export const createMessagesList = (history, messageId) => { return list.reverse(); }; +const toTokenCount = (value: unknown) => { + const parsed = Number(value || 0); + return Number.isFinite(parsed) ? Math.trunc(parsed) : 0; +}; + +export const getUsageTokenCount = (usage?: Record | null) => { + if (!usage) { + return 0; + } + + let promptTokens = toTokenCount(usage.prompt_tokens || usage.prompt_eval_count || 0); + if (!promptTokens && (usage.prompt_n != null || usage.cache_n != null)) { + promptTokens = toTokenCount(usage.prompt_n || 0) + toTokenCount(usage.cache_n || 0); + } + if (!promptTokens) { + promptTokens = toTokenCount(usage.input_tokens || 0); + } + + const completionTokens = toTokenCount( + usage.completion_tokens || usage.output_tokens || usage.eval_count || usage.predicted_n || 0 + ); + + return promptTokens + completionTokens; +}; + export const formatFileSize = (size) => { if (size == null) return 'Unknown size'; if (typeof size !== 'number' || size < 0) return 'Invalid size';