From 8dab4ecb06ab2e5c1d16a479db4bbca4445db6d2 Mon Sep 17 00:00:00 2001
From: Cohee <18619528+Cohee1207@users.noreply.github.com>
Date: Tue, 25 Jun 2024 21:53:10 +0300
Subject: [PATCH] Add simple custom macros registration

---
 public/script.js         |  3 ++-
 public/scripts/macros.js | 58 ++++++++++++++++++++++++++++++++++++++--
 2 files changed, 58 insertions(+), 3 deletions(-)

diff --git a/public/script.js b/public/script.js
index 93b1bbc01..70eb2726e 100644
--- a/public/script.js
+++ b/public/script.js
@@ -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_TYPE, Popup, callGenericPopup, fixToastrForDialogs } from './scripts/popup.js';
 import { renderTemplate, renderTemplateAsync } from './scripts/templates.js';
@@ -7763,6 +7763,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.
diff --git a/public/scripts/macros.js b/public/scripts/macros.js
index 2b4ea17a7..23722aaa2 100644
--- a/public/scripts/macros.js
+++ b/public/scripts/macros.js
@@ -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 } 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,56 @@ Handlebars.registerHelper('helperMissing', function () {
     return substituteParams(`{{${macroName}}}`);
 });
 
+export class MacrosParser {
+    /**
+     * A map of registered macros.
+     * @type {Map<string, string|(() => string)>}
+     */
+    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|(() => string)} 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');
+        }
+
+        if (this.#macros.has(key)) {
+            console.warn(`Macro ${key} is already registered`);
+        }
+
+        if (typeof value !== 'string' && typeof value !== 'function') {
+            throw new Error('Macro value must be a string or a function that returns a string');
+        }
+
+        this.#macros.set(key, value);
+    }
+
+    /**
+     * Populate the env object with macro values from the current context.
+     * @param {Object<string, *>} 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;
+        }
+    }
+}
+
 /**
  * Gets a hashed id of the current chat from the metadata.
  * If no metadata exists, creates a new hash and saves it.
@@ -315,12 +365,16 @@ 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
+    MacrosParser.populateEnv(env);
+
     // Substitute passed-in variables
     for (const varName in env) {
         if (!Object.hasOwn(env, varName)) continue;
 
         const param = env[varName];
-        content = content.replace(new RegExp(`{{${varName}}}`, 'gi'), param);
+        const value = typeof param === 'function' ? param() : param;
+        content = content.replace(new RegExp(`{{${escapeRegex(varName)}}}`, 'gi'), value);
     }
 
     content = content.replace(/{{maxPrompt}}/gi, () => String(getMaxContextSize()));