Add reasoning tokens to token count.

This commit is contained in:
Cohee
2025-01-26 18:29:31 +02:00
parent a7516937f7
commit bfedf20db5

View File

@ -3195,7 +3195,8 @@ class StreamingProcessor {
this.#updateMessageBlockVisibility(); this.#updateMessageBlockVisibility();
const currentTime = new Date(); const currentTime = new Date();
// Don't waste time calculating token count for streaming // Don't waste time calculating token count for streaming
const currentTokenCount = isFinal && power_user.message_token_count_enabled ? getTokenCount(processedText, 0) : 0; const tokenCountText = (this.reasoning || '') + processedText;
const currentTokenCount = isFinal && power_user.message_token_count_enabled ? getTokenCount(tokenCountText, 0) : 0;
const timePassed = formatGenerationTimer(this.timeStarted, currentTime, currentTokenCount); const timePassed = formatGenerationTimer(this.timeStarted, currentTime, currentTokenCount);
chat[messageId]['mes'] = processedText; chat[messageId]['mes'] = processedText;
chat[messageId]['gen_started'] = this.timeStarted; chat[messageId]['gen_started'] = this.timeStarted;
@ -5936,7 +5937,8 @@ export async function saveReply(type, getMessage, fromStreaming, title, swipes,
chat[chat.length - 1]['extra']['model'] = getGeneratingModel(); chat[chat.length - 1]['extra']['model'] = getGeneratingModel();
chat[chat.length - 1]['extra']['reasoning'] = reasoning; chat[chat.length - 1]['extra']['reasoning'] = reasoning;
if (power_user.message_token_count_enabled) { if (power_user.message_token_count_enabled) {
chat[chat.length - 1]['extra']['token_count'] = await getTokenCountAsync(chat[chat.length - 1]['mes'], 0); const tokenCountText = (reasoning || '') + chat[chat.length - 1]['mes'];
chat[chat.length - 1]['extra']['token_count'] = await getTokenCountAsync(tokenCountText, 0);
} }
const chat_id = (chat.length - 1); const chat_id = (chat.length - 1);
await eventSource.emit(event_types.MESSAGE_RECEIVED, chat_id); await eventSource.emit(event_types.MESSAGE_RECEIVED, chat_id);
@ -5957,7 +5959,8 @@ export async function saveReply(type, getMessage, fromStreaming, title, swipes,
chat[chat.length - 1]['extra']['model'] = getGeneratingModel(); chat[chat.length - 1]['extra']['model'] = getGeneratingModel();
chat[chat.length - 1]['extra']['reasoning'] += reasoning; chat[chat.length - 1]['extra']['reasoning'] += reasoning;
if (power_user.message_token_count_enabled) { if (power_user.message_token_count_enabled) {
chat[chat.length - 1]['extra']['token_count'] = await getTokenCountAsync(chat[chat.length - 1]['mes'], 0); const tokenCountText = (reasoning || '') + chat[chat.length - 1]['mes'];
chat[chat.length - 1]['extra']['token_count'] = await getTokenCountAsync(tokenCountText, 0);
} }
const chat_id = (chat.length - 1); const chat_id = (chat.length - 1);
await eventSource.emit(event_types.MESSAGE_RECEIVED, chat_id); await eventSource.emit(event_types.MESSAGE_RECEIVED, chat_id);
@ -5975,7 +5978,8 @@ export async function saveReply(type, getMessage, fromStreaming, title, swipes,
chat[chat.length - 1]['extra']['model'] = getGeneratingModel(); chat[chat.length - 1]['extra']['model'] = getGeneratingModel();
chat[chat.length - 1]['extra']['reasoning'] += reasoning; chat[chat.length - 1]['extra']['reasoning'] += reasoning;
if (power_user.message_token_count_enabled) { if (power_user.message_token_count_enabled) {
chat[chat.length - 1]['extra']['token_count'] = await getTokenCountAsync(chat[chat.length - 1]['mes'], 0); const tokenCountText = (reasoning || '') + chat[chat.length - 1]['mes'];
chat[chat.length - 1]['extra']['token_count'] = await getTokenCountAsync(tokenCountText, 0);
} }
const chat_id = (chat.length - 1); const chat_id = (chat.length - 1);
await eventSource.emit(event_types.MESSAGE_RECEIVED, chat_id); await eventSource.emit(event_types.MESSAGE_RECEIVED, chat_id);
@ -6001,7 +6005,8 @@ export async function saveReply(type, getMessage, fromStreaming, title, swipes,
chat[chat.length - 1]['gen_finished'] = generationFinished; chat[chat.length - 1]['gen_finished'] = generationFinished;
if (power_user.message_token_count_enabled) { if (power_user.message_token_count_enabled) {
chat[chat.length - 1]['extra']['token_count'] = await getTokenCountAsync(chat[chat.length - 1]['mes'], 0); const tokenCountText = (reasoning || '') + chat[chat.length - 1]['mes'];
chat[chat.length - 1]['extra']['token_count'] = await getTokenCountAsync(tokenCountText, 0);
} }
if (selected_group) { if (selected_group) {
@ -8544,7 +8549,8 @@ function swipe_left() { // when we swipe left..but no generation.
} }
const swipeMessage = $('#chat').find(`[mesid="${chat.length - 1}"]`); const swipeMessage = $('#chat').find(`[mesid="${chat.length - 1}"]`);
const tokenCount = await getTokenCountAsync(chat[chat.length - 1].mes, 0); const tokenCountText = (chat[chat.length - 1]?.extra?.reasoning || '') + chat[chat.length - 1].mes;
const tokenCount = await getTokenCountAsync(tokenCountText, 0);
chat[chat.length - 1]['extra']['token_count'] = tokenCount; chat[chat.length - 1]['extra']['token_count'] = tokenCount;
swipeMessage.find('.tokenCounterDisplay').text(`${tokenCount}t`); swipeMessage.find('.tokenCounterDisplay').text(`${tokenCount}t`);
} }
@ -8719,7 +8725,8 @@ const swipe_right = () => {
chat[chat.length - 1].extra = {}; chat[chat.length - 1].extra = {};
} }
const tokenCount = await getTokenCountAsync(chat[chat.length - 1].mes, 0); const tokenCountText = (chat[chat.length - 1]?.extra?.reasoning || '') + chat[chat.length - 1].mes;
const tokenCount = await getTokenCountAsync(tokenCountText, 0);
chat[chat.length - 1]['extra']['token_count'] = tokenCount; chat[chat.length - 1]['extra']['token_count'] = tokenCount;
swipeMessage.find('.tokenCounterDisplay').text(`${tokenCount}t`); swipeMessage.find('.tokenCounterDisplay').text(`${tokenCount}t`);
} }
@ -9521,7 +9528,8 @@ function addDebugFunctions() {
message.extra = {}; message.extra = {};
} }
message.extra.token_count = await getTokenCountAsync(message.mes, 0); const tokenCountText = (message?.extra?.reasoning || '') + message.mes;
message.extra.token_count = await getTokenCountAsync(tokenCountText, 0);
} }
await saveChatConditional(); await saveChatConditional();