From 3d5fc5fee199fb47e0994cfc565af17c336ba893 Mon Sep 17 00:00:00 2001 From: tomt610 Date: Fri, 9 Jan 2026 20:10:21 +0000 Subject: [PATCH] Refactor historical context injection logic to support dynamic message indexing based on injection position --- src/systems/generation/injector.js | 104 ++++++++++++++++-------- src/systems/generation/promptBuilder.js | 93 +++++++++++++++++---- 2 files changed, 148 insertions(+), 49 deletions(-) diff --git a/src/systems/generation/injector.js b/src/systems/generation/injector.js index 0b82d66..2801a56 100644 --- a/src/systems/generation/injector.js +++ b/src/systems/generation/injector.js @@ -44,8 +44,9 @@ let originalMessageContent = new Map(); /** * Builds a map of historical context data from ST chat messages with rpg_companion_swipes data. * Returns a map keyed by message index with formatted context strings. + * The index stored depends on the injection position setting. * - * @returns {Map} Map of message index to context data + * @returns {Map} Map of target message index to formatted context string */ function buildHistoricalContextMap() { const historyPersistence = extensionSettings.historyPersistence; @@ -61,19 +62,41 @@ function buildHistoricalContextMap() { const trackerConfig = extensionSettings.trackerConfig; const userName = context.name1; + const position = historyPersistence.injectionPosition || 'assistant_message_end'; const contextMap = new Map(); // Determine how many messages to include (0 = all available) const messageCount = historyPersistence.messageCount || 0; const maxMessages = messageCount === 0 ? chat.length : Math.min(messageCount, chat.length); - // Start from the second-to-last message (skip the most recent one as it gets current context) - // and work backwards - let processedCount = 0; + // Find the last assistant message - this is the one that gets current context via setExtensionPrompt + // We should NOT add historical context to it + let lastAssistantIndex = -1; + for (let i = chat.length - 1; i >= 0; i--) { + if (!chat[i].is_user && !chat[i].is_system) { + lastAssistantIndex = i; + break; + } + } - for (let i = chat.length - 2; i >= 0 && (messageCount === 0 || processedCount < maxMessages); i--) { + // Iterate through messages to find those with tracker data + // Start from before the last assistant message + let processedCount = 0; + const startIndex = lastAssistantIndex > 0 ? lastAssistantIndex - 1 : chat.length - 2; + + for (let i = startIndex; i >= 0 && (messageCount === 0 || processedCount < maxMessages); i--) { const message = chat[i]; + // Skip system messages + if (message.is_system) { + continue; + } + + // Only assistant messages have rpg_companion_swipes data + if (message.is_user) { + continue; + } + // Get the rpg_companion_swipes data for current swipe const swipeData = message.extra?.rpg_companion_swipes; if (!swipeData) { @@ -96,11 +119,32 @@ function buildHistoricalContextMap() { const preamble = historyPersistence.contextPreamble || '[Context at this point:]'; const wrappedContext = `\n${preamble}\n${formattedContext}`; - // Store with message index and whether it's a user message - contextMap.set(i, { - context: wrappedContext, - isUserMessage: message.is_user - }); + // Determine which message index to store based on injection position + let targetIndex = i; // Default: the assistant message itself + + if (position === 'user_message_end') { + // Find the next user message after this assistant message + for (let j = i + 1; j < chat.length; j++) { + if (chat[j].is_user && !chat[j].is_system) { + targetIndex = j; + break; + } + } + // If no user message found after, skip this one + if (targetIndex === i) { + continue; + } + } + // For assistant_message_end, extra_user_message, extra_assistant_message: + // We inject into the assistant message itself (for now - extra messages handled differently) + + // Store the context keyed by target index + // If multiple assistant messages map to the same user message, append + if (contextMap.has(targetIndex)) { + contextMap.set(targetIndex, contextMap.get(targetIndex) + wrappedContext); + } else { + contextMap.set(targetIndex, wrappedContext); + } processedCount++; } @@ -141,41 +185,23 @@ function injectHistoricalContextIntoChat() { console.log(`[RPG Companion] Injecting historical context into ${contextMap.size} messages`); - const position = historyPersistence.injectionPosition || 'assistant_message_end'; - // Clear any previous stored content originalMessageContent.clear(); let injectedCount = 0; - for (const [msgIdx, data] of contextMap) { + for (const [msgIdx, ctxContent] of contextMap) { const message = chat[msgIdx]; if (!message || typeof message.mes !== 'string') { continue; } - const { context: ctxContent, isUserMessage } = data; + // Store original content for restoration + originalMessageContent.set(msgIdx, message.mes); - // Determine if we should inject based on position and message type - let shouldInject = false; - - if (position === 'user_message_end' && isUserMessage) { - shouldInject = true; - } else if (position === 'assistant_message_end' && !isUserMessage) { - shouldInject = true; - } else if (position === 'extra_user_message' || position === 'extra_assistant_message') { - // For these positions, inject regardless of message type - shouldInject = true; - } - - if (shouldInject) { - // Store original content for restoration - originalMessageContent.set(msgIdx, message.mes); - - // Modify the message in-place - message.mes = message.mes + ctxContent; - injectedCount++; - console.log(`[RPG Companion] Injected context into message ${msgIdx}`); - } + // Modify the message in-place + message.mes = message.mes + ctxContent; + injectedCount++; + console.log(`[RPG Companion] Injected context into message ${msgIdx}`); } console.log(`[RPG Companion] Successfully injected historical context into ${injectedCount} messages`); @@ -533,6 +559,14 @@ Ensure these details naturally reflect and influence the narrative. Character be // Inject historical context directly into chat messages // This temporarily modifies messages and will be restored after generation injectHistoricalContextIntoChat(); + + // Register a one-time listener to restore messages after prompt is built + // This ensures messages are restored even if generation is cancelled or for dry runs + const restoreAfterPrompt = () => { + restoreOriginalMessageContent(); + eventSource.off(event_types.GENERATE_AFTER_COMBINE_PROMPTS, restoreAfterPrompt); + }; + eventSource.once(event_types.GENERATE_AFTER_COMBINE_PROMPTS, restoreAfterPrompt); } /** diff --git a/src/systems/generation/promptBuilder.js b/src/systems/generation/promptBuilder.js index b861e4d..74b8080 100644 --- a/src/systems/generation/promptBuilder.js +++ b/src/systems/generation/promptBuilder.js @@ -1123,26 +1123,91 @@ export async function generateSeparateUpdatePrompt() { // Add chat history as separate user/assistant messages with per-message historical context const recentMessages = chat.slice(-depth); const startIndex = chat.length - depth; + const position = historyPersistence?.injectionPosition || 'assistant_message_end'; - for (let i = 0; i < recentMessages.length; i++) { - const message = recentMessages[i]; - const chatIndex = startIndex + i; - let content = message.mes; + // Build a map of which messages should get context based on position setting + // Key: message index in recentMessages, Value: context string + const contextInjectionMap = new Map(); + + if (historyPersistence?.enabled) { + // Find the last assistant message index (in recentMessages) + let lastAssistantIdx = -1; + for (let i = recentMessages.length - 1; i >= 0; i--) { + if (!recentMessages[i].is_user && !recentMessages[i].is_system) { + lastAssistantIdx = i; + break; + } + } + + // Iterate through assistant messages to find tracker data + for (let i = 0; i < recentMessages.length; i++) { + const message = recentMessages[i]; + + // Skip user and system messages - only assistant messages have tracker data + if (message.is_user || message.is_system) { + continue; + } + + // Skip the last assistant message - it gets current context elsewhere + if (i === lastAssistantIdx) { + continue; + } - // Append historical tracker context to this message if enabled and available - if (historyPersistence?.enabled && chatIndex < chat.length - 1) { const swipeData = message.extra?.rpg_companion_swipes; - if (swipeData) { - const currentSwipeId = message.swipe_id || 0; - const trackerData = swipeData[currentSwipeId]; - if (trackerData) { - const formattedContext = formatHistoricalTrackerData(trackerData, trackerConfig, userName); - if (formattedContext) { - const preamble = historyPersistence.contextPreamble || '[Context at this point:]'; - content += `\n${preamble}\n${formattedContext}`; + if (!swipeData) { + continue; + } + + const currentSwipeId = message.swipe_id || 0; + const trackerData = swipeData[currentSwipeId]; + if (!trackerData) { + continue; + } + + const formattedContext = formatHistoricalTrackerData(trackerData, trackerConfig, userName); + if (!formattedContext) { + continue; + } + + const preamble = historyPersistence.contextPreamble || '[Context at this point:]'; + const wrappedContext = `\n${preamble}\n${formattedContext}`; + + // Determine target message based on position + let targetIdx = i; + + if (position === 'user_message_end') { + // Find next user message after this assistant message + for (let j = i + 1; j < recentMessages.length; j++) { + if (recentMessages[j].is_user && !recentMessages[j].is_system) { + targetIdx = j; + break; } } + // If no user message found, skip + if (targetIdx === i) { + continue; + } } + // For assistant_message_end, extra_user_message, extra_assistant_message: + // Inject into the assistant message itself + + // Append to existing or create new entry + if (contextInjectionMap.has(targetIdx)) { + contextInjectionMap.set(targetIdx, contextInjectionMap.get(targetIdx) + wrappedContext); + } else { + contextInjectionMap.set(targetIdx, wrappedContext); + } + } + } + + // Now build the messages array with injected context + for (let i = 0; i < recentMessages.length; i++) { + const message = recentMessages[i]; + let content = message.mes; + + // Add historical context if this message is a target + if (contextInjectionMap.has(i)) { + content += contextInjectionMap.get(i); } messages.push({