Expressions: Add WebLLM extension classification

f1bc217e79f7854c094090cbb378cd633d9cf218

Cohee <18619528+Cohee1207@users.noreply.github.com>

2 files changed, +50 -15Showing whitespace changes
public/scripts/extensions/expressions/index.js+48 -14
@@ -15,6 +15,7 @@ import { SlashCommandEnumValue, enumTypes } from '../../slash-commands/SlashComm
1515import { commonEnumProviders } from '../../slash-commands/SlashCommandCommonEnumsProvider.js';
1616import { slashCommandReturnHelper } from '../../slash-commands/SlashCommandReturnHelper.js';
1717import { SlashCommandClosure } from '../../slash-commands/SlashCommandClosure.js';
18+import { generateWebLlmChatPrompt, isWebLlmSupported } from '../shared.js';
1819export { MODULE_NAME };
1920
2021const MODULE_NAME = 'expressions';
@@ -59,6 +60,7 @@ const EXPRESSION_API = {
5960 local: 0,
6061 extras: 1,
6162 llm: 2,
63+ webllm: 3,
6264};
6365
6466let expressionsList = null;
@@ -852,7 +854,7 @@ function setTalkingHeadState(newState) {
852854 extension_settings.expressions.talkinghead = newState; // Store setting
853855 saveSettingsDebounced();
854856
855857 if (extension_settings.expressions.api == [EXPRESSION_API.local, ||EXPRESSION_API.llm, EXPRESSION_API.webllm].includes(extension_settings.expressions.api == EXPRESSION_API.llm)) {
856858 return;
857859 }
858860
@@ -1057,21 +1059,25 @@ function parseLlmResponse(emotionResponse, labels) {
10571059 console.debug(`fuzzy search found: ${result[0].item} as closest for the LLM response:`, emotionResponse);
10581060 return result[0].item;
10591061 }
1062+ const lowerCaseResponse = String(emotionResponse || '').toLowerCase();
1063+ for (const label of labels) {
1064+ if (lowerCaseResponse.includes(label.toLowerCase())) {
1065+ console.debug(`Found label ${label} in the LLM response:`, emotionResponse);
1066+ return label;
1067+ }
1068+ }
10601069 }
10611070
10621071 throw new Error('Could not parse emotion response ' + emotionResponse);
10631072}
10641073
1065-function onTextGenSettingsReady(args) {
1074+/**
1066- // Only call if inside an API call
1075+ * Gets the JSON schema for the LLM API.
1067- if (inApiCall && extension_settings.expressions.api === EXPRESSION_API.llm && isJsonSchemaSupported()) {
1076+ * @param {string[]} emotions A list of emotions to search for.
1068- const emotions = DEFAULT_EXPRESSIONS.filter((e) => e != 'talkinghead');
1077+ * @returns {object} The JSON schema for the LLM API.
1069- Object.assign(args, {
1078+ */
1070- top_k: 1,
1079+function getJsonSchema(emotions) {
1071- stop: [],
1080+ return {
1072- stopping_strings: [],
1073- custom_token_bans: [],
1074- json_schema: {
10751081 $schema: 'http://json-schema.org/draft-04/schema#',
10761082 type: 'object',
10771083 properties: {
@@ -1083,7 +1089,19 @@ function onTextGenSettingsReady(args) {
10831089 required: [
10841090 'emotion',
10851091 ],
1086- },
1092+ };
1093+}
1094+
1095+function onTextGenSettingsReady(args) {
1096+ // Only call if inside an API call
1097+ if (inApiCall && extension_settings.expressions.api === EXPRESSION_API.llm && isJsonSchemaSupported()) {
1098+ const emotions = DEFAULT_EXPRESSIONS.filter((e) => e != 'talkinghead');
1099+ Object.assign(args, {
1100+ top_k: 1,
1101+ stop: [],
1102+ stopping_strings: [],
1103+ custom_token_bans: [],
1104+ json_schema: getJsonSchema(emotions),
10871105 });
10881106 }
10891107}
@@ -1139,6 +1157,22 @@ export async function getExpressionLabel(text, expressionsApi = extension_settin
11391157 const emotionResponse = await generateRaw(text, main_api, false, false, prompt);
11401158 return parseLlmResponse(emotionResponse, expressionsList);
11411159 }
1160+ // Using WebLLM
1161+ case EXPRESSION_API.webllm: {
1162+ if (!isWebLlmSupported()) {
1163+ console.warn('WebLLM is not supported. Using fallback expression');
1164+ return getFallbackExpression();
1165+ }
1166+
1167+ const expressionsList = await getExpressionsList();
1168+ const prompt = substituteParamsExtended(customPrompt, { labels: expressionsList }) || await getLlmPrompt(expressionsList);
1169+ const messages = [
1170+ { role: 'user', content: text + '\n\n' + prompt },
1171+ ];
1172+
1173+ const emotionResponse = await generateWebLlmChatPrompt(messages);
1174+ return parseLlmResponse(emotionResponse, expressionsList);
1175+ }
11421176 // Extras
11431177 default: {
11441178 const url = new URL(getApiUrl());
@@ -1603,7 +1637,7 @@ function onExpressionApiChanged() {
16031637 const tempApi = this.value;
16041638 if (tempApi) {
16051639 extension_settings.expressions.api = Number(tempApi);
16061640 $('.expression_llm_prompt_block').toggle([EXPRESSION_API.llm, EXPRESSION_API.webllm].includes(extension_settings.expressions.api === EXPRESSION_API.llm));
16071641 expressionsList = null;
16081642 spriteCache = {};
16091643 moduleWorker();
@@ -1940,7 +1974,7 @@ function migrateSettings() {
19401974
19411975 await renderAdditionalExpressionSettings();
19421976 $('#expression_api').val(extension_settings.expressions.api ?? EXPRESSION_API.extras);
19431977 $('.expression_llm_prompt_block').toggle([EXPRESSION_API.llm, EXPRESSION_API.webllm].includes(extension_settings.expressions.api === EXPRESSION_API.llm));
19441978 $('#expression_llm_prompt').val(extension_settings.expressions.llmPrompt ?? '');
19451979 $('#expression_llm_prompt').on('input', function () {
19461980 extension_settings.expressions.llmPrompt = $(this).val();
public/scripts/extensions/expressions/settings.html+2 -1
@@ -24,7 +24,8 @@
2424 <select id="expression_api" class="flex1 margin0">
2525 <option value="0" data-i18n="Local">Local</option>
2626 <option value="1" data-i18n="Extras">Extras</option>
2727 <option value="2" data-i18n="LLMMain API">LLMMain API</option>
28+ <option value="3" data-i18n="WebLLM Extension">WebLLM Extension</option>
2829 </select>
2930 </div>
3031 <div class="expression_llm_prompt_block m-b-1 m-t-1">