refactor: support array prompts in placeholder services

src/services/pick.ts

@@ -1,19 +1,38 @@
-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) => {
-    chat.prompt = chat.prompt.replaceAll(
-      pickPlaceholderRegex,
-      (_, optionsString: string) => {
-        const options = optionsString.split('::')
-        if (options.length === 0) {
-          return ''
-        }
-        return options[Math.floor(Math.random() * options.length)] ?? ''
+  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) => {
+            const options = optionsString.split('::')
+            if (options.length === 0) {
+              return ''
+            }
+            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]
-  if (words === undefined || words.length === 0) {
-    // If the word list is empty or undefined, don't register the macro
+  const wordsLists = new Map(Object.entries(useStore.getState().wordLists))
+
+  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: () => {
-      const words = useStore.getState().wordLists[dictName] as string[]
+    description: `Get a random word from the specified word list`,
+    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(wordList)
-  }
+  registerRandomWordMacros()
 
   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) {
-        registerRandomWordMacros(wordList)
-      }
+      macros.registry.unregisterMacro(`randomWord`)
+      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) => {
-    chat.prompt = chat.prompt.replaceAll(
-      randomWordPlaceholderRegex,
-      (_, wordList: string) => {
-        const words = useStore.getState().wordLists[wordList]
-        if (words === undefined || words.length === 0) {
-          return ''
-        }
-        return words[Math.floor(Math.random() * words.length)] ?? ''
+  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 = 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 ?? ''
+          }
+        )
+      } 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,15 +15,35 @@ const shufflePlaceholderRegex = /%%shuffle::([\w\W]*?)(;;[\w\W]*?)?%%/g
 
 export const setupShufflePlaceholders = () => {
   const { eventTypes, eventSource } = SillyTavern.getContext()
-  eventSource.on(eventTypes.GENERATE_AFTER_DATA, (chat: generationData) => {
-    chat.prompt = chat.prompt.replaceAll(
-      shufflePlaceholderRegex,
-      (_, toBeshuffleed: string, separator?: string) => {
-        const options = toBeshuffleed.split('::')
-        return fishersYatesShuffle(options).join(
-          separator !== undefined ? separator.substring(2) : ''
+  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) => {
+            const options = toBeshuffleed.split('::')
+            return fishersYatesShuffle(options).join(
+              separator !== undefined ? separator.substring(2) : ''
+            )
+          }
         )
+      } 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
+        })
       }
-    )
-  })
+    }
+  )
 }