Merge pull request #2421 from SillyTavern/macro-register

Add simple custom macros registration
This commit is contained in:
Cohee 2024-06-26 22:01:24 +03:00 committed by GitHub
commit aceca89080
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 115 additions and 5 deletions

View File

@ -226,7 +226,7 @@ import { BulkEditOverlay, CharacterContextMenu } from './scripts/BulkEditOverlay
import { loadMancerModels, loadOllamaModels, loadTogetherAIModels, loadInfermaticAIModels, loadOpenRouterModels, loadVllmModels, loadAphroditeModels, loadDreamGenModels } from './scripts/textgen-models.js';
import { appendFileContent, hasPendingFileAttachment, populateFileAttachment, decodeStyleTags, encodeStyleTags, isExternalMediaAllowed, getCurrentEntityId } from './scripts/chats.js';
import { initPresetManager } from './scripts/preset-manager.js';
import { evaluateMacros } from './scripts/macros.js';
import { MacrosParser, evaluateMacros } from './scripts/macros.js';
import { currentUser, setUserControls } from './scripts/user.js';
import { POPUP_RESULT, POPUP_TYPE, Popup, callGenericPopup, fixToastrForDialogs } from './scripts/popup.js';
import { renderTemplate, renderTemplateAsync } from './scripts/templates.js';
@ -7790,6 +7790,7 @@ window['SillyTavern'].getContext = function () {
* @deprecated Handlebars for extensions are no longer supported.
*/
registerHelper: () => { },
registerMacro: MacrosParser.registerMacro.bind(MacrosParser),
registedDebugFunction: registerDebugFunction,
/**
* @deprecated Use renderExtensionTemplateAsync instead.

View File

@ -1,5 +1,5 @@
import { chat, chat_metadata, main_api, getMaxContextSize, getCurrentChatId, substituteParams } from '../script.js';
import { timestampToMoment, isDigitsOnly, getStringHash } from './utils.js';
import { timestampToMoment, isDigitsOnly, getStringHash, escapeRegex, uuidv4 } from './utils.js';
import { textgenerationwebui_banned_in_macros } from './textgen-settings.js';
import { replaceInstructMacros } from './instruct-mode.js';
import { replaceVariableMacros } from './variables.js';
@ -13,6 +13,108 @@ Handlebars.registerHelper('helperMissing', function () {
return substituteParams(`{{${macroName}}}`);
});
/**
* @typedef {Object<string, *>} EnvObject
* @typedef {(nonce: string) => string} MacroFunction
*/
export class MacrosParser {
/**
* A map of registered macros.
* @type {Map<string, string|MacroFunction>}
*/
static #macros = new Map();
/**
* Registers a global macro that can be used anywhere where substitution is allowed.
* @param {string} key Macro name (key)
* @param {string|MacroFunction} value A string or a function that returns a string
*/
static registerMacro(key, value) {
if (typeof key !== 'string') {
throw new Error('Macro key must be a string');
}
// Allowing surrounding whitespace would just create more confusion...
key = key.trim();
if (!key) {
throw new Error('Macro key must not be empty or whitespace only');
}
if (key.startsWith('{{') || key.endsWith('}}')) {
throw new Error('Macro key must not include the surrounding braces');
}
if (typeof value !== 'string' && typeof value !== 'function') {
console.warn(`Macro value for "${key}" will be converted to a string`);
value = this.sanitizeMacroValue(value);
}
if (this.#macros.has(key)) {
console.warn(`Macro ${key} is already registered`);
}
this.#macros.set(key, value);
}
/**
* Populate the env object with macro values from the current context.
* @param {EnvObject} env Env object for the current evaluation context
* @returns {void}
*/
static populateEnv(env) {
if (!env || typeof env !== 'object') {
console.warn('Env object is not provided');
return;
}
// No macros are registered
if (this.#macros.size === 0) {
return;
}
for (const [key, value] of this.#macros) {
env[key] = value;
}
}
/**
* Performs a type-check on the macro value and returns a sanitized version of it.
* @param {any} value Value returned by a macro
* @returns {string} Sanitized value
*/
static sanitizeMacroValue(value) {
if (typeof value === 'string') {
return value;
}
if (value === null || value === undefined) {
return '';
}
if (value instanceof Promise) {
console.warn('Promises are not supported as macro values');
return '';
}
if (typeof value === 'function') {
console.warn('Functions are not supported as macro values');
return '';
}
if (value instanceof Date) {
return value.toISOString();
}
if (typeof value === 'object') {
return JSON.stringify(value);
}
return String(value);
}
}
/**
* Gets a hashed id of the current chat from the metadata.
* If no metadata exists, creates a new hash and saves it.
@ -284,7 +386,7 @@ function timeDiffReplace(input) {
/**
* Substitutes {{macro}} parameters in a string.
* @param {string} content - The string to substitute parameters in.
* @param {Object<string, *>} env - Map of macro names to the values they'll be substituted with. If the param
* @param {EnvObject} env - Map of macro names to the values they'll be substituted with. If the param
* values are functions, those functions will be called and their return values are used.
* @returns {string} The string with substituted parameters.
*/
@ -315,12 +417,19 @@ export function evaluateMacros(content, env) {
content = content.replace(/{{noop}}/gi, '');
content = content.replace(/{{input}}/gi, () => String($('#send_textarea').val()));
// Add all registered macros to the env object
const nonce = uuidv4();
MacrosParser.populateEnv(env);
// Substitute passed-in variables
for (const varName in env) {
if (!Object.hasOwn(env, varName)) continue;
content = content.replace(new RegExp(`{{${escapeRegex(varName)}}}`, 'gi'), () => {
const param = env[varName];
content = content.replace(new RegExp(`{{${varName}}}`, 'gi'), param);
const value = MacrosParser.sanitizeMacroValue(typeof param === 'function' ? param(nonce) : param);
return value;
});
}
content = content.replace(/{{maxPrompt}}/gi, () => String(getMaxContextSize()));