refactor: support array prompts in placeholder services

src/services/pick.ts

@@ -1,10 +1,13 @@
-import type { generationData } from '@/types/SillyTavern'
-
 const pickPlaceholderRegex = /%%pick::([\w\W]*?)%%/g
 
 export const setupPickPlaceholders = () => {
   const { eventTypes, eventSource } = SillyTavern.getContext()
-  eventSource.on(eventTypes.GENERATE_AFTER_DATA, (chat: generationData) => {
+  eventSource.on(
+    eventTypes.GENERATE_AFTER_DATA,
+    (
+      chat: { prompt: string } | { prompt: { role: string; content: string }[] }
+    ) => {
+      if (typeof chat.prompt === 'string') {
         chat.prompt = chat.prompt.replaceAll(
           pickPlaceholderRegex,
           (_, optionsString: string) => {
@@ -15,5 +18,21 @@ export const setupPickPlaceholders = () => {
             return options[Math.floor(Math.random() * options.length)] ?? ''
           }
         )
+      } else if (Array.isArray(chat.prompt)) {
+        chat.prompt = chat.prompt.map((prompt) => {
+          prompt.content = prompt.content.replaceAll(
+            pickPlaceholderRegex,
+            (_, optionsString: string) => {
+              const options = optionsString.split('::')
+              if (options.length === 0) {
+                return ''
+              }
+              return options[Math.floor(Math.random() * options.length)] ?? ''
+            }
+          )
+          return prompt
         })
+      }
+    }
+  )
 }

src/services/randomWord.ts

@@ -1,51 +1,109 @@
 import { useStore } from '@/store'
-import type { generationData } from '@/types/SillyTavern'
 
-const registerRandomWordMacros = (dictName: string) => {
+const registerRandomWordMacros = () => {
   const { macros } = SillyTavern.getContext()
-  const words = useStore.getState().wordLists[dictName]
+  const wordsLists = new Map(Object.entries(useStore.getState().wordLists))
-  if (words === undefined || words.length === 0) {
+
-    // If the word list is empty or undefined, don't register the macro
+  if (wordsLists.size === 0) {
+    console.info(
+      `[st-randomness-helpers] No word lists found, skipping random word macros registration`
+    )
     return
   }
 
-  macros.register(`randomWord::${dictName}`, {
+  macros.register(`randomWord`, {
     category: 'random',
-    handler: () => {
+    description: `Get a random word from the specified word list`,
-      const words = useStore.getState().wordLists[dictName] as string[]
+    returns:
+      'A random word from the specified word list or an empty string if the word list is empty or the dictionary name is not found',
+    exampleUsage: wordsLists
+      .keys()
+      .toArray()
+      .map((wordList) => `{{randomWord::${wordList}}}`),
+    returnType: 'string',
+    unnamedArgs: [
+      {
+        name: 'dictionary name',
+        description: `The name of the word list to use. Available dictionaries: ${wordsLists
+          .keys()
+          .toArray()
+          .map((wordList) => `'${wordList}'`)
+          .join(`, `)}`,
+        type: 'string',
+        optional: false,
+        sampleValue: wordsLists.keys().next().value ?? 'wordlist'
+      }
+    ],
+    handler: ({ unnamedArgs: [dictName] }) => {
+      if (!dictName) {
+        console.error(
+          `[st-randomness-helpers] randomWord: No dictionary name provided, returning empty string`
+        )
+        return ``
+      }
+      if (!wordsLists.has(dictName)) {
+        console.error(
+          `[st-randomness-helpers] randomWord: Dictionary ${dictName} not found, returning empty string`
+        )
+        return ``
+      }
+      const words = wordsLists.get(dictName) as string[] // We know it exists
       return words[Math.floor(Math.random() * words.length)] as string
     }
   })
-  macros.register(`placeholder::randomWord::${dictName}`, {
+  macros.register(`randomWordPlaceholder`, {
     category: 'random',
-    handler: () => {
+    description: `Get a placeholder for a random word from the specified word list. The placeholder is in the form of %%randomWord::dictionary name%%`,
+    returns:
+      'A placeholder for a random word from the specified word list or an empty string if the word list is empty or the dictionary name is not found',
+    exampleUsage: wordsLists
+      .keys()
+      .toArray()
+      .map((wordList) => `{{randomWordPlaceholder::${wordList}}}`),
+    returnType: 'string',
+    unnamedArgs: [
+      {
+        name: 'dictionary name',
+        description: `The name of the word list to use. Available dictionaries: ${wordsLists
+          .keys()
+          .toArray()
+          .map((wordList) => `'${wordList}'`)
+          .join(`, `)}`,
+        type: 'string',
+        optional: false,
+        sampleValue: wordsLists.keys().next().value ?? 'wordlist'
+      }
+    ],
+    handler: ({ unnamedArgs: [dictName] }) => {
+      if (!dictName) {
+        console.error(
+          `[st-randomness-helpers] randomWordPlaceholder: No dictionary name provided, returning empty string`
+        )
+        return ``
+      }
+      if (!wordsLists.has(dictName)) {
+        console.error(
+          `[st-randomness-helpers] randomWordPlaceholder: Dictionary ${dictName} not found, returning empty string`
+        )
+        return ``
+      }
       return `%%randomWord::${dictName}%%`
     }
   })
 }
 
 export const setupRandomWordMacros = () => {
-  for (const wordList of Object.keys(useStore.getState().wordLists)) {
+  registerRandomWordMacros()
-    registerRandomWordMacros(wordList)
-  }
 
   useStore.subscribe(
     (state) => Object.keys(state.wordLists),
-    (wordLists, oldWordLists) => {
+    () => {
-      const wordListsSet = new Set(wordLists)
-      const oldWordListsSet = new Set(oldWordLists)
-      const deletedWordLists = oldWordListsSet.difference(wordListsSet)
-      const newWordLists = wordListsSet.difference(oldWordListsSet)
-
       const { macros } = SillyTavern.getContext()
-      for (const wordList of deletedWordLists) {
-        macros.registry.unregisterMacro(`randomWord::${wordList}`)
-        macros.registry.unregisterMacro(`placeholder::randomWord::${wordList}`)
-      }
 
-      for (const wordList of newWordLists) {
+      macros.registry.unregisterMacro(`randomWord`)
-        registerRandomWordMacros(wordList)
+      macros.registry.unregisterMacro(`randomWordPlaceholder`)
-      }
+
+      registerRandomWordMacros()
     }
   )
 }
@@ -54,16 +112,54 @@ const randomWordPlaceholderRegex = /%%randomWord::([\w\W]*?)%%/g
 
 export const setupRandomWordPlaceholders = () => {
   const { eventTypes, eventSource } = SillyTavern.getContext()
-  eventSource.on(eventTypes.GENERATE_AFTER_DATA, (chat: generationData) => {
+  eventSource.on(
+    eventTypes.GENERATE_AFTER_DATA,
+    (
+      chat: { prompt: string } | { prompt: { role: string; content: string }[] }
+    ) => {
+      if (typeof chat.prompt === 'string') {
         chat.prompt = chat.prompt.replaceAll(
           randomWordPlaceholderRegex,
           (_, wordList: string) => {
-        const words = useStore.getState().wordLists[wordList]
+            const words = new Map(
+              Object.entries(useStore.getState().wordLists)
+            ).get(wordList)
             if (words === undefined || words.length === 0) {
+              console.error(
+                `[st-randomness-helpers] randomWordPlaceholders: No words found for ${wordList}`
+              )
               return ''
             }
-        return words[Math.floor(Math.random() * words.length)] ?? ''
+            const word = words[Math.floor(Math.random() * words.length)]
+            return word ?? ''
+          }
+        )
+      } else if (Array.isArray(chat.prompt)) {
+        const newPrompt = chat.prompt.map(
+          ({ role, content }: { role: string; content: string }) => {
+            return {
+              role,
+              content: content.replaceAll(
+                randomWordPlaceholderRegex,
+                (_, wordList: string) => {
+                  const words = new Map(
+                    Object.entries(useStore.getState().wordLists)
+                  ).get(wordList)
+                  if (words === undefined || words.length === 0) {
+                    console.error(
+                      `[st-randomness-helpers] randomWordPlaceholders: No words found for ${wordList}`
+                    )
+                    return ''
+                  }
+                  const word = words[Math.floor(Math.random() * words.length)]
+                  return word ?? ''
+                }
+              )
+            }
+          }
+        )
+        chat.prompt = newPrompt
+      }
     }
   )
-  })
 }

src/services/shuffle.ts

@@ -1,5 +1,3 @@
-import type { generationData } from '@/types/SillyTavern'
-
 const fishersYatesShuffle = <Type>(arr: Type[]): Type[] => {
   const shuffled = [...arr]
   for (let index = shuffled.length - 1; index > 0; index--) {
@@ -17,7 +15,12 @@ const shufflePlaceholderRegex = /%%shuffle::([\w\W]*?)(;;[\w\W]*?)?%%/g
 
 export const setupShufflePlaceholders = () => {
   const { eventTypes, eventSource } = SillyTavern.getContext()
-  eventSource.on(eventTypes.GENERATE_AFTER_DATA, (chat: generationData) => {
+  eventSource.on(
+    eventTypes.GENERATE_AFTER_DATA,
+    (
+      chat: { prompt: string } | { prompt: { role: string; content: string }[] }
+    ) => {
+      if (typeof chat.prompt === 'string') {
         chat.prompt = chat.prompt.replaceAll(
           shufflePlaceholderRegex,
           (_, toBeshuffleed: string, separator?: string) => {
@@ -27,5 +30,20 @@ export const setupShufflePlaceholders = () => {
             )
           }
         )
+      } else if (Array.isArray(chat.prompt)) {
+        chat.prompt = chat.prompt.map((prompt) => {
+          prompt.content = prompt.content.replaceAll(
+            shufflePlaceholderRegex,
+            (_, toBeshuffleed: string, separator?: string) => {
+              const options = toBeshuffleed.split('::')
+              return fishersYatesShuffle(options).join(
+                separator !== undefined ? separator.substring(2) : ''
+              )
+            }
+          )
+          return prompt
         })
+      }
+    }
+  )
 }