Refactor history injection to modify prompts instead of chat messages
This prevents any risk of injected context being accidentally saved to the chat. Instead of modifying chat[].mes directly, we now: 1. Build a context map during GENERATION_STARTED 2. Inject into the prompt string (GENERATE_AFTER_COMBINE_PROMPTS) for text completion 3. Inject into the message array (CHAT_COMPLETION_PROMPT_READY) for chat completion The original chat messages are never modified.
This commit is contained in:
tomt610
@@ -37,9 +37,8 @@ let currentSuppressionState = false;
|
||||
// Track last chat length we committed at to prevent duplicate commits from streaming
|
||||
let lastCommittedChatLength = -1;
|
||||
|
||||
// Store original message content for restoration after generation
|
||||
// Map of message index -> original mes content
|
||||
let originalMessageContent = new Map();
|
||||
// Store context map for prompt injection (used by event handlers)
|
||||
let pendingContextMap = new Map();
|
||||
|
||||
/**
|
||||
* Builds a map of historical context data from ST chat messages with rpg_companion_swipes data.
|
||||
@@ -162,81 +161,171 @@ function buildHistoricalContextMap() {
|
||||
}
|
||||
|
||||
/**
|
||||
* Injects historical context into chat messages by modifying them in-place.
|
||||
* Stores original content for restoration after generation.
|
||||
* This approach works for ALL API types (text completion and chat completion).
|
||||
* Prepares historical context for injection into prompts.
|
||||
* This builds the context map and stores it for use by prompt event handlers.
|
||||
* Does NOT modify the original chat messages.
|
||||
*/
|
||||
function injectHistoricalContextIntoChat() {
|
||||
function prepareHistoricalContextInjection() {
|
||||
const historyPersistence = extensionSettings.historyPersistence;
|
||||
if (!historyPersistence || !historyPersistence.enabled) {
|
||||
// console.log('[RPG Companion] History persistence not enabled, skipping injection');
|
||||
pendingContextMap = new Map();
|
||||
return;
|
||||
}
|
||||
|
||||
if (currentSuppressionState || !extensionSettings.enabled) {
|
||||
// console.log('[RPG Companion] Skipping history injection: suppressed or disabled');
|
||||
pendingContextMap = new Map();
|
||||
return;
|
||||
}
|
||||
|
||||
const context = getContext();
|
||||
const chat = context.chat;
|
||||
if (!chat || chat.length < 2) {
|
||||
// console.log('[RPG Companion] Chat too short, skipping history injection');
|
||||
pendingContextMap = new Map();
|
||||
return;
|
||||
}
|
||||
|
||||
// Build the context map
|
||||
const contextMap = buildHistoricalContextMap();
|
||||
if (contextMap.size === 0) {
|
||||
// console.log('[RPG Companion] No historical context to inject');
|
||||
return;
|
||||
// Build and store the context map for use by prompt handlers
|
||||
pendingContextMap = buildHistoricalContextMap();
|
||||
}
|
||||
|
||||
/**
|
||||
* Injects historical context into a text completion prompt string.
|
||||
* Searches for message content in the prompt and appends context after matches.
|
||||
*
|
||||
* @param {string} prompt - The text completion prompt
|
||||
* @returns {string} - The modified prompt with injected context
|
||||
*/
|
||||
function injectContextIntoTextPrompt(prompt) {
|
||||
if (pendingContextMap.size === 0) {
|
||||
return prompt;
|
||||
}
|
||||
|
||||
// console.log(`[RPG Companion] Injecting historical context into ${contextMap.size} messages`);
|
||||
|
||||
// Clear any previous stored content
|
||||
originalMessageContent.clear();
|
||||
|
||||
const context = getContext();
|
||||
const chat = context.chat;
|
||||
let modifiedPrompt = prompt;
|
||||
let injectedCount = 0;
|
||||
for (const [msgIdx, ctxContent] of contextMap) {
|
||||
|
||||
// Process each message that needs context injection
|
||||
for (const [msgIdx, ctxContent] of pendingContextMap) {
|
||||
const message = chat[msgIdx];
|
||||
if (!message || typeof message.mes !== 'string') {
|
||||
continue;
|
||||
}
|
||||
|
||||
// Store original content for restoration
|
||||
originalMessageContent.set(msgIdx, message.mes);
|
||||
// Find the message content in the prompt
|
||||
// Use a portion of the message to find it (last 100 chars should be unique enough)
|
||||
const searchContent = message.mes.length > 100
|
||||
? message.mes.slice(-100)
|
||||
: message.mes;
|
||||
|
||||
const searchIndex = modifiedPrompt.lastIndexOf(searchContent);
|
||||
if (searchIndex === -1) {
|
||||
// Message not found in prompt (might be truncated)
|
||||
continue;
|
||||
}
|
||||
|
||||
// Modify the message in-place
|
||||
message.mes = message.mes + ctxContent;
|
||||
// Find the end of this message content in the prompt
|
||||
const insertPosition = searchIndex + searchContent.length;
|
||||
|
||||
// Insert the context after the message
|
||||
modifiedPrompt = modifiedPrompt.slice(0, insertPosition) + ctxContent + modifiedPrompt.slice(insertPosition);
|
||||
injectedCount++;
|
||||
// console.log(`[RPG Companion] Injected context into message ${msgIdx}`);
|
||||
}
|
||||
|
||||
// console.log(`[RPG Companion] Successfully injected historical context into ${injectedCount} messages`);
|
||||
if (injectedCount > 0) {
|
||||
console.log(`[RPG Companion] Injected historical context into ${injectedCount} positions in text prompt`);
|
||||
}
|
||||
|
||||
return modifiedPrompt;
|
||||
}
|
||||
|
||||
/**
|
||||
* Restores original message content after generation completes.
|
||||
* This ensures the injected context doesn't persist in the actual chat data.
|
||||
* Injects historical context into a chat completion message array.
|
||||
* Modifies the content of messages in the array directly.
|
||||
*
|
||||
* @param {Array} chatMessages - The chat completion message array
|
||||
* @returns {Array} - The modified message array with injected context
|
||||
*/
|
||||
function restoreOriginalMessageContent() {
|
||||
if (originalMessageContent.size === 0) {
|
||||
return;
|
||||
function injectContextIntoChatPrompt(chatMessages) {
|
||||
if (pendingContextMap.size === 0 || !Array.isArray(chatMessages)) {
|
||||
return chatMessages;
|
||||
}
|
||||
|
||||
const context = getContext();
|
||||
const chat = context.chat;
|
||||
let injectedCount = 0;
|
||||
|
||||
// console.log(`[RPG Companion] Restoring ${originalMessageContent.size} messages to original content`);
|
||||
// Process each message that needs context injection
|
||||
for (const [msgIdx, ctxContent] of pendingContextMap) {
|
||||
const originalMessage = chat[msgIdx];
|
||||
if (!originalMessage || typeof originalMessage.mes !== 'string') {
|
||||
continue;
|
||||
}
|
||||
|
||||
for (const [msgIdx, originalContent] of originalMessageContent) {
|
||||
if (chat[msgIdx]) {
|
||||
chat[msgIdx].mes = originalContent;
|
||||
// Find this message in the chat completion array by matching content
|
||||
// Use a portion of the message to find it
|
||||
const searchContent = originalMessage.mes.length > 100
|
||||
? originalMessage.mes.slice(-100)
|
||||
: originalMessage.mes;
|
||||
|
||||
for (const promptMsg of chatMessages) {
|
||||
if (promptMsg.content && typeof promptMsg.content === 'string' &&
|
||||
promptMsg.content.includes(searchContent)) {
|
||||
// Found the message - append context
|
||||
promptMsg.content = promptMsg.content + ctxContent;
|
||||
injectedCount++;
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
originalMessageContent.clear();
|
||||
if (injectedCount > 0) {
|
||||
console.log(`[RPG Companion] Injected historical context into ${injectedCount} messages in chat prompt`);
|
||||
}
|
||||
|
||||
return chatMessages;
|
||||
}
|
||||
|
||||
/**
|
||||
* Event handler for GENERATE_AFTER_COMBINE_PROMPTS (text completion).
|
||||
* Injects historical context into the prompt string.
|
||||
*
|
||||
* @param {Object} eventData - Event data with prompt property
|
||||
*/
|
||||
function onGenerateAfterCombinePrompts(eventData) {
|
||||
if (!eventData || typeof eventData.prompt !== 'string') {
|
||||
return;
|
||||
}
|
||||
|
||||
if (eventData.dryRun) {
|
||||
return;
|
||||
}
|
||||
|
||||
eventData.prompt = injectContextIntoTextPrompt(eventData.prompt);
|
||||
|
||||
// Clear the pending context after injection
|
||||
pendingContextMap = new Map();
|
||||
}
|
||||
|
||||
/**
|
||||
* Event handler for CHAT_COMPLETION_PROMPT_READY.
|
||||
* Injects historical context into the chat message array.
|
||||
*
|
||||
* @param {Object} eventData - Event data with chat property
|
||||
*/
|
||||
function onChatCompletionPromptReady(eventData) {
|
||||
if (!eventData || !Array.isArray(eventData.chat)) {
|
||||
return;
|
||||
}
|
||||
|
||||
if (eventData.dryRun) {
|
||||
return;
|
||||
}
|
||||
|
||||
eventData.chat = injectContextIntoChatPrompt(eventData.chat);
|
||||
|
||||
// Clear the pending context after injection
|
||||
pendingContextMap = new Map();
|
||||
}
|
||||
|
||||
/**
|
||||
@@ -565,22 +654,24 @@ Ensure these details naturally reflect and influence the narrative. Character be
|
||||
// Set suppression state for the historical context injection
|
||||
currentSuppressionState = shouldSuppress;
|
||||
|
||||
// Inject historical context directly into chat messages
|
||||
// This temporarily modifies messages and will be restored after generation
|
||||
injectHistoricalContextIntoChat();
|
||||
// Prepare historical context for injection into prompts
|
||||
// This builds the context map but does NOT modify original chat messages
|
||||
prepareHistoricalContextInjection();
|
||||
|
||||
// Register a one-time listener to restore messages after prompt is built
|
||||
// Using .once() so it auto-removes after firing
|
||||
eventSource.once(event_types.GENERATE_AFTER_COMBINE_PROMPTS, () => {
|
||||
restoreOriginalMessageContent();
|
||||
});
|
||||
// Register one-time listeners to inject context into the actual prompt
|
||||
// These modify only the prompt sent to the API, not the stored chat data
|
||||
if (pendingContextMap.size > 0) {
|
||||
eventSource.once(event_types.GENERATE_AFTER_COMBINE_PROMPTS, onGenerateAfterCombinePrompts);
|
||||
eventSource.once(event_types.CHAT_COMPLETION_PROMPT_READY, onChatCompletionPromptReady);
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Called when generation ends to restore original message content.
|
||||
* Called when generation ends to clean up any pending context.
|
||||
* This should be called from the GENERATION_ENDED event handler.
|
||||
*/
|
||||
export function onGenerationEndedCleanup() {
|
||||
restoreOriginalMessageContent();
|
||||
// Clear any pending context that wasn't used (e.g., if generation was cancelled)
|
||||
pendingContextMap = new Map();
|
||||
}
|
||||
|
||||
|
||||
Reference in New Issue
Block a user