Add WebLLM source to /summarize command

This commit is contained in:
Cohee 2024-08-13 00:57:32 +03:00
parent 5b3388d0c8
commit 30c19eb8fd
1 changed files with 12 additions and 9 deletions

View File

@ -40,6 +40,7 @@ let inApiCall = false;
/** /**
* Count the number of tokens in the provided text. * Count the number of tokens in the provided text.
* @param {string} text Text to count tokens for * @param {string} text Text to count tokens for
* @param {number} padding Number of additional tokens to add to the count
* @returns {Promise<number>} Number of tokens in the text * @returns {Promise<number>} Number of tokens in the text
*/ */
async function countSourceTokens(text, padding = 0) { async function countSourceTokens(text, padding = 0) {
@ -389,16 +390,13 @@ function getIndexOfLatestChatSummary(chat) {
async function onChatEvent() { async function onChatEvent() {
// Module not enabled // Module not enabled
if (extension_settings.memory.source === summary_sources.extras) { if (extension_settings.memory.source === summary_sources.extras && !modules.includes('summarize')) {
if (!modules.includes('summarize')) { return;
return;
}
} }
if (extension_settings.memory.source === summary_sources.webllm) { // WebLLM is not supported
if (!isWebLlmSupported()) { if (extension_settings.memory.source === summary_sources.webllm && !isWebLlmSupported()) {
return; return;
}
} }
const context = getContext(); const context = getContext();
@ -510,6 +508,11 @@ async function summarizeCallback(args, text) {
return await callExtrasSummarizeAPI(text); return await callExtrasSummarizeAPI(text);
case summary_sources.main: case summary_sources.main:
return await generateRaw(text, '', false, false, prompt, extension_settings.memory.overrideResponseLength); return await generateRaw(text, '', false, false, prompt, extension_settings.memory.overrideResponseLength);
case summary_sources.webllm: {
const messages = [{ role: 'system', content: prompt }, { role: 'user', content: text }].filter(m => m.content);
const params = extension_settings.memory.overrideResponseLength > 0 ? { max_tokens: extension_settings.memory.overrideResponseLength } : {};
return await generateWebLlmChatPrompt(messages, params);
}
default: default:
toastr.warning('Invalid summarization source specified'); toastr.warning('Invalid summarization source specified');
return ''; return '';
@ -1045,7 +1048,7 @@ jQuery(async function () {
name: 'summarize', name: 'summarize',
callback: summarizeCallback, callback: summarizeCallback,
namedArgumentList: [ namedArgumentList: [
new SlashCommandNamedArgument('source', 'API to use for summarization', [ARGUMENT_TYPE.STRING], false, false, '', ['main', 'extras']), new SlashCommandNamedArgument('source', 'API to use for summarization', [ARGUMENT_TYPE.STRING], false, false, '', Object.values(summary_sources)),
SlashCommandNamedArgument.fromProps({ SlashCommandNamedArgument.fromProps({
name: 'prompt', name: 'prompt',
description: 'prompt to use for summarization', description: 'prompt to use for summarization',