Move counter to plugin. Use chat context to get messages

This commit is contained in:
Cohee 2023-10-21 14:23:56 +03:00
parent 63ecca1fe2
commit 85de505553
2 changed files with 18 additions and 18 deletions

View File

@ -1,6 +1,7 @@
import { callPopup, main_api } from "../../../script.js";
import { getContext } from "../../extensions.js";
import { getTokenizerModel } from "../../tokenizers.js";
import { registerSlashCommand } from "../../slash-commands.js";
import { getTokenCount, getTokenizerModel } from "../../tokenizers.js";
async function doTokenCounter() {
const selectedTokenizer = main_api == 'openai'
@ -29,6 +30,20 @@ async function doTokenCounter() {
callPopup(dialog, 'text');
}
function doCount() {
// get all of the messages in the chat
const context = getContext();
const messages = context.chat.filter(x => x.mes && !x.is_system).map(x => x.mes);
//concat all the messages into a single string
const allMessages = messages.join(' ');
console.debug('All messages:', allMessages);
//toastr success with the token count of the chat
toastr.success(`Token count: ${getTokenCount(allMessages)}`);
}
jQuery(() => {
const buttonHtml = `
<div id="token_counter" class="list-group-item flex-container flexGap5">
@ -37,4 +52,5 @@ jQuery(() => {
</div>`;
$('#extensionsMenu').prepend(buttonHtml);
$('#token_counter').on('click', doTokenCounter);
registerSlashCommand('count', doCount, [], ' counts the number of tokens in the current chat', true, false);
});

View File

@ -28,7 +28,7 @@ import {
} from "./instruct-mode.js";
import { registerSlashCommand } from "./slash-commands.js";
import { tokenizers, getTokenCount } from "./tokenizers.js";
import { tokenizers } from "./tokenizers.js";
import { countOccurrences, debounce, delay, isOdd, resetScrollHeight, sortMoments, timestampToMoment } from "./utils.js";
@ -1554,21 +1554,6 @@ function doResetPanels() {
$("#movingUIreset").trigger('click');
}
function doCount() {
// get all of the messages in the chat
const messages = $('#chat .mes');
//concat all the messages into a single string
const allMessages = messages.toArray().map(x => $(x).find('.mes_text').text()).join(' ');
console.log(allMessages);
//toastr success with the token count of the chat
toastr.success(`Token count: ${getTokenCount(allMessages)}`);
}
function setAvgBG() {
const bgimg = new Image();
bgimg.src = $('#bg1')
@ -2411,5 +2396,4 @@ $(document).ready(() => {
registerSlashCommand('cut', doMesCut, [], '<span class="monospace">(number)</span> cuts the specified message from the chat', true, true);
registerSlashCommand('resetpanels', doResetPanels, ['resetui'], ' resets UI panels to original state.', true, true);
registerSlashCommand('bgcol', setAvgBG, [], ' WIP test of auto-bg avg coloring', true, true);
registerSlashCommand('count', doCount, [], ' counts the number of tokens in the current chat', true, false);
});