Refactor historical context injection logic to support dynamic message indexing based on injection position

This commit is contained in:
tomt610
2026-01-09 20:10:21 +00:00
parent 98ef751a9f
commit 3d5fc5fee1
2 changed files with 148 additions and 49 deletions
+69 -35
View File
@@ -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<number, {context: string, isUserMessage: boolean}>} Map of message index to context data
* @returns {Map<number, string>} 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);
}
/**
+79 -14
View File
@@ -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({