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,5 +1,4 @@
 import { useStore } from '@/store'
-import type { generationData } from '@/types/SillyTavern'
 
 const registerRandomWordMacros = () => {
   const { macros } = SillyTavern.getContext()
@@ -113,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
         })
       }
+    }
+  )
+}