diff --git a/src/LLMProviders/chainManager.ts b/src/LLMProviders/chainManager.ts index 46c2edca..44b055a2 100644 --- a/src/LLMProviders/chainManager.ts +++ b/src/LLMProviders/chainManager.ts @@ -374,13 +374,12 @@ export default class ChainManager { async updateMemoryWithLoadedMessages(messages: ChatMessage[]) { await this.memoryManager.clearChatMemory(); + // Use memoryManager.saveContext to apply compaction for any old uncompacted messages for (let i = 0; i < messages.length; i += 2) { const userMsg = messages[i]; const aiMsg = messages[i + 1]; if (userMsg && aiMsg && userMsg.sender === USER_SENDER) { - await this.memoryManager - .getMemory() - .saveContext({ input: userMsg.message }, { output: aiMsg.message }); + await this.memoryManager.saveContext({ input: userMsg.message }, { output: aiMsg.message }); } } } diff --git a/src/LLMProviders/chainRunner/BaseChainRunner.ts b/src/LLMProviders/chainRunner/BaseChainRunner.ts index 10099ccb..1deb8461 100644 --- a/src/LLMProviders/chainRunner/BaseChainRunner.ts +++ b/src/LLMProviders/chainRunner/BaseChainRunner.ts @@ -70,14 +70,16 @@ export abstract class BaseChainRunner implements ChainRunner { !(abortController.signal.aborted && abortController.signal.reason === ABORT_REASON.NEW_CHAT); if (shouldAddMessage) { - // Use saveContext for atomic operation and proper memory management - // Note: LangChain's memory expects text content, not multimodal arrays, so multimodal content is not saved + // Use MemoryManager.saveContext for atomic operation and proper memory management + // This applies chat history compaction to reduce memory bloat from tool results + // Note: LangChain's memory expects text content, not multimodal arrays // For truncated empty responses, save a placeholder message const outputForMemory = llmFormattedOutput || fullAIResponse || "[Response truncated - no content generated]"; - await this.chainManager.memoryManager - .getMemory() - .saveContext({ input: userMessage.message }, { output: outputForMemory }); + await this.chainManager.memoryManager.saveContext( + { input: userMessage.message }, + { output: outputForMemory } + ); // For empty truncated responses, show a helpful message const displayMessage = diff --git a/src/LLMProviders/chainRunner/utils/chatHistoryUtils.ts b/src/LLMProviders/chainRunner/utils/chatHistoryUtils.ts index 90696a37..544e61ef 100644 --- a/src/LLMProviders/chainRunner/utils/chatHistoryUtils.ts +++ b/src/LLMProviders/chainRunner/utils/chatHistoryUtils.ts @@ -219,17 +219,8 @@ export function extractConversationTurns(processedHistory: ProcessedMessage[]): * Load chat history from memory and add to messages array. * This is the single entry point for all chain runners to use. * - * NOTE: Chat history compaction was intentionally removed for simplicity. - * A previous implementation would summarize older conversation turns when - * total context exceeded a threshold. This was removed because: - * 1. It added complexity (LLM calls for summarization) - * 2. Context compaction in ContextManager already handles large context - * 3. BufferWindowMemory already limits conversation history length - * - * If chat history compaction is needed in the future, consider: - * - Extracting conversation turns with extractConversationTurns() - * - Summarizing older turns while keeping recent ones intact - * - Using the same autoCompactThreshold setting for consistency + * Note: Chat history is already compacted at save time (in MemoryManager.saveContext) + * so tool results (localSearch, readNote, etc.) are stored as compact summaries. * * @param memory - LangChain memory instance * @param messages - Target messages array (system message should already be added) @@ -248,7 +239,7 @@ export async function loadAndAddChatHistory( const processedHistory = processRawChatHistory(rawHistory); - // Add history messages directly + // Add history messages directly (already compacted at save time) for (const msg of processedHistory) { messages.push({ role: msg.role, content: msg.content }); } diff --git a/src/LLMProviders/memoryManager.ts b/src/LLMProviders/memoryManager.ts index bac41e69..85d34014 100644 --- a/src/LLMProviders/memoryManager.ts +++ b/src/LLMProviders/memoryManager.ts @@ -1,3 +1,4 @@ +import { compactAssistantOutput } from "@/context/ChatHistoryCompactor"; import { getSettings, subscribeToSettingsChange } from "@/settings/model"; import { BaseChatMemory, BufferWindowMemory } from "@langchain/classic/memory"; import { BaseChatMessageHistory } from "@langchain/core/chat_history"; @@ -52,8 +53,21 @@ export default class MemoryManager { return variables; } + /** + * Save a conversation turn to memory. + * The output (assistant response) is compacted to reduce memory bloat from + * accumulated tool results (localSearch, readNote, etc.). + */ async saveContext(input: any, output: any): Promise { - if (this.debug) console.log("Saving to memory - Input:", input, "Output:", output); - await this.memory.saveContext(input, output); + // Compact the output to prevent memory bloat from tool results + const compactedOutput = + typeof output === "string" + ? compactAssistantOutput(output) + : { ...output, output: compactAssistantOutput(output.output) }; + + if (this.debug) { + console.log("Saving to memory - Input:", input, "Output (compacted):", compactedOutput); + } + await this.memory.saveContext(input, compactedOutput); } } diff --git a/src/chatUtils.toolMarkers.test.ts b/src/chatUtils.toolMarkers.test.ts index e18c3f42..961996e3 100644 --- a/src/chatUtils.toolMarkers.test.ts +++ b/src/chatUtils.toolMarkers.test.ts @@ -17,6 +17,10 @@ class MockMemoryManager { getMemory() { return this.memory; } + async saveContext(input: { input: string }, output: { output: string }) { + // Mock the MemoryManager.saveContext behavior (compaction is tested separately) + await this.memory.saveContext(input, output); + } } describe("updateChatMemory with tool call markers", () => { diff --git a/src/chatUtils.ts b/src/chatUtils.ts index 34bdd176..85c895d3 100644 --- a/src/chatUtils.ts +++ b/src/chatUtils.ts @@ -10,15 +10,14 @@ export async function updateChatMemory( await memoryManager.clearChatMemory(); // Process each message in sequence + // Use memoryManager.saveContext to apply compaction for any old uncompacted messages for (let i = 0; i < messages.length - 1; i++) { const msg = messages[i]; if (msg.sender === USER_SENDER) { const nextMsg = messages[i + 1]; if (nextMsg?.sender === AI_SENDER) { - await memoryManager - .getMemory() - .saveContext({ input: msg.message }, { output: nextMsg.message }); + await memoryManager.saveContext({ input: msg.message }, { output: nextMsg.message }); } } } diff --git a/src/context/ChatHistoryCompactor.test.ts b/src/context/ChatHistoryCompactor.test.ts new file mode 100644 index 00000000..e490a033 --- /dev/null +++ b/src/context/ChatHistoryCompactor.test.ts @@ -0,0 +1,219 @@ +import { compactAssistantOutput } from "./ChatHistoryCompactor"; + +describe("ChatHistoryCompactor", () => { + describe("compactAssistantOutput", () => { + it("should return small content unchanged", () => { + const output = "This is a short response without tool results."; + expect(compactAssistantOutput(output)).toBe(output); + }); + + it("should compact large localSearch blocks", () => { + const largeContent = "A".repeat(10000); + const output = `I searched for that. + + + +Note 1 +notes/note1.md +${largeContent} + + + +Based on my search, here's what I found.`; + + const result = compactAssistantOutput(output, { verbatimThreshold: 1000 }); + expect(typeof result).toBe("string"); + expect((result as string).length).toBeLessThan(output.length); + expect(result).toContain("I searched for that"); + expect(result).toContain("here's what I found"); + expect(result).toContain("prior_context"); + }); + + it("should compact large readNote JSON results", () => { + const largeContent = `## Introduction +${"This is intro. ".repeat(200)} + +## Methods +${"This is methods. ".repeat(200)}`; + + const readNoteResult = JSON.stringify({ + notePath: "research/paper.md", + noteTitle: "Research Paper", + content: largeContent, + chunkIndex: 0, + totalChunks: 1, + }); + + const output = `Let me read that note. + +Tool 'readNote' result: ${readNoteResult} + +Here's a summary.`; + + const result = compactAssistantOutput(output, { verbatimThreshold: 1000 }); + expect(typeof result).toBe("string"); + expect((result as string).length).toBeLessThan(output.length); + expect(result).toContain("COMPACTED"); + expect(result).toContain("## Introduction"); + }); + + it("should handle multimodal content arrays", () => { + const largeContent = "B".repeat(10000); + const textItem = { type: "text", text: `${largeContent}` }; + const imageItem = { type: "image_url", url: "data:image/png;base64,..." }; + const output = [textItem, imageItem]; + + const result = compactAssistantOutput(output, { verbatimThreshold: 1000 }); + expect(Array.isArray(result)).toBe(true); + const resultArray = result as any[]; + expect(resultArray[0].text.length).toBeLessThan(textItem.text.length); + expect(resultArray[1]).toEqual(imageItem); // Image unchanged + }); + + it("should keep small tool results verbatim", () => { + const smallResult = JSON.stringify({ + notePath: "notes/short.md", + content: "Brief content.", + }); + const output = `Tool 'readNote' result: ${smallResult}`; + expect(compactAssistantOutput(output)).toBe(output); + }); + + it("should never compact selected_text blocks", () => { + const largeContent = "X".repeat(20000); + const output = ` +${largeContent} +`; + + // selected_text is not in the TOOL_RESULT_PATTERNS, so it won't be processed + // But if it were, it should remain unchanged due to isRecoverable check + const result = compactAssistantOutput(output, { verbatimThreshold: 1000 }); + expect(result).toBe(output); + }); + + it("should handle multiple tool results in one message", () => { + const large1 = "A".repeat(8000); + const large2 = "B".repeat(8000); + const output = `First: +${large1} + +Second: +${large2} + +Done.`; + + const result = compactAssistantOutput(output, { verbatimThreshold: 1000 }); + expect(typeof result).toBe("string"); + expect((result as string).length).toBeLessThan(output.length * 0.3); + }); + + it("should return non-string/non-array content unchanged", () => { + const output = { some: "object" }; + expect(compactAssistantOutput(output as any)).toBe(output); + }); + + it("should handle readNote JSON with nested braces in content", () => { + // Content contains code with nested braces - this used to break the regex + const codeContent = `## Code Example +${"function test() { if (true) { return { value: 1 }; } } ".repeat(100)} + +## Another Section +${"More content here. ".repeat(100)}`; + + const readNoteResult = JSON.stringify({ + notePath: "code/example.md", + noteTitle: "Code Example", + content: codeContent, + chunkIndex: 0, + totalChunks: 1, + }); + + const output = `Here's the code: + +Tool 'readNote' result: ${readNoteResult} + +That's the implementation.`; + + const result = compactAssistantOutput(output, { verbatimThreshold: 1000 }); + expect(typeof result).toBe("string"); + expect((result as string).length).toBeLessThan(output.length); + expect(result).toContain("COMPACTED"); + expect(result).toContain("## Code Example"); + expect(result).toContain("That's the implementation"); + }); + + it("should handle multiple readNote results in sequence", () => { + const content1 = "A".repeat(3000); + const content2 = "B".repeat(3000); + + const result1 = JSON.stringify({ notePath: "a.md", content: content1 }); + const result2 = JSON.stringify({ notePath: "b.md", content: content2 }); + + const output = `First note: +Tool 'readNote' result: ${result1} + +Second note: +Tool 'readNote' result: ${result2} + +Done.`; + + const result = compactAssistantOutput(output, { verbatimThreshold: 1000 }); + expect(typeof result).toBe("string"); + expect((result as string).length).toBeLessThan(output.length); + expect(result).toContain("First note:"); + expect(result).toContain("Second note:"); + expect(result).toContain("Done."); + // Both should be compacted + const compactedCount = ((result as string).match(/COMPACTED/g) || []).length; + expect(compactedCount).toBe(2); + }); + + it("should preserve all documents in localSearch results", () => { + const largeContent1 = "First document content. ".repeat(200); + const largeContent2 = "Second document content. ".repeat(200); + const largeContent3 = "Third document content. ".repeat(200); + + const output = `I found these notes: + + + +First Note +folder/first.md +${largeContent1} + + +Second Note +folder/second.md +${largeContent2} + + +Third Note +third.md +${largeContent3} + + + +Based on my search, here's what I found.`; + + const result = compactAssistantOutput(output, { verbatimThreshold: 1000 }); + expect(typeof result).toBe("string"); + const resultStr = result as string; + + // Should be compacted (smaller than original) + expect(resultStr.length).toBeLessThan(output.length); + + // All three documents should be preserved + expect(resultStr).toContain("First Note"); + expect(resultStr).toContain("Second Note"); + expect(resultStr).toContain("Third Note"); + expect(resultStr).toContain("folder/first.md"); + expect(resultStr).toContain("folder/second.md"); + expect(resultStr).toContain("third.md"); + expect(resultStr).toContain("3 search results"); + + // Surrounding text should be preserved + expect(resultStr).toContain("I found these notes:"); + expect(resultStr).toContain("here's what I found"); + }); + }); +}); diff --git a/src/context/ChatHistoryCompactor.ts b/src/context/ChatHistoryCompactor.ts new file mode 100644 index 00000000..777b4913 --- /dev/null +++ b/src/context/ChatHistoryCompactor.ts @@ -0,0 +1,332 @@ +/** + * ChatHistoryCompactor - Compacts tool results in chat history. + * + * When assistant responses contain large tool results (localSearch, readNote, etc.), + * this compactor replaces them with compact summaries to prevent memory bloat. + * + * This is applied at save time in MemoryManager.saveContext() so that + * both LangChain memory and LLM payloads contain compacted content. + */ + +import { + extractContentFromBlock, + extractSourceFromBlock, + getSourceType, + isRecoverable, +} from "./contextBlockRegistry"; +import { + CompactionConfig, + DEFAULT_COMPACTION_CONFIG, + compactBySection, + escapeXmlAttr, +} from "./compactionUtils"; + +/** + * Build regex patterns for all registered block types that appear in tool results. + * These are XML blocks embedded in assistant responses. + */ +function buildToolResultPatterns(): Array<{ pattern: RegExp; tag: string }> { + // Tags that commonly appear as tool results in chat history + const toolResultTags = [ + "localSearch", + "note_context", + "active_note", + "retrieved_document", + "url_content", + "youtube_video_context", + ]; + + return toolResultTags.map((tag) => ({ + pattern: new RegExp(`<${tag}[^>]*>[\\s\\S]*?`, "g"), + tag, + })); +} + +const TOOL_RESULT_PATTERNS = buildToolResultPatterns(); + +/** + * Pattern to find readNote tool result prefix in chat history. + * The JSON is extracted separately using brace-balanced parsing. + */ +const READ_NOTE_PREFIX = "Tool 'readNote' result: "; + +/** + * Extract a brace-balanced JSON object starting at position. + * Returns the JSON string and end position, or null if not valid. + */ +function extractBalancedJson( + content: string, + startPos: number +): { json: string; endPos: number } | null { + if (content[startPos] !== "{") return null; + + let depth = 0; + let inString = false; + let escape = false; + + for (let i = startPos; i < content.length; i++) { + const char = content[i]; + + if (escape) { + escape = false; + continue; + } + + if (char === "\\" && inString) { + escape = true; + continue; + } + + if (char === '"') { + inString = !inString; + continue; + } + + if (inString) continue; + + if (char === "{") depth++; + else if (char === "}") { + depth--; + if (depth === 0) { + return { json: content.slice(startPos, i + 1), endPos: i + 1 }; + } + } + } + + return null; // Unbalanced braces +} + +/** + * Compact an assistant's output before saving to memory. + * + * This is the main entry point - a simple function that takes the assistant's + * response and returns a compacted version with large tool results summarized. + * + * @param output - The assistant's response (string or multimodal array) + * @param config - Optional configuration overrides + * @returns Compacted output + */ +export function compactAssistantOutput( + output: string | any[], + config: Partial = {} +): string | any[] { + if (Array.isArray(output)) { + // Handle multimodal content - compact text parts + return output.map((item) => { + if (item.type === "text" && typeof item.text === "string") { + return { ...item, text: compactOutputString(item.text, config) }; + } + return item; + }); + } + + if (typeof output === "string") { + return compactOutputString(output, config); + } + + return output; +} + +/** + * Compact a string that may contain tool results. + */ +function compactOutputString(content: string, config: Partial = {}): string { + const threshold = config.verbatimThreshold ?? DEFAULT_COMPACTION_CONFIG.verbatimThreshold; + let result = content; + + // Compact XML-based tool results + for (const { pattern, tag } of TOOL_RESULT_PATTERNS) { + // Reset regex state for each use + pattern.lastIndex = 0; + + result = result.replace(pattern, (match) => { + // Only compact if the block is large + if (match.length < threshold) { + return match; + } + return compactToolResultBlock(match, tag, config); + }); + } + + // Compact readNote JSON results using brace-balanced extraction + result = compactReadNoteResults(result, threshold, config); + + return result; +} + +/** + * Find and compact all readNote JSON results in content using brace-balanced extraction. + */ +function compactReadNoteResults( + content: string, + threshold: number, + config: Partial +): string { + let result = ""; + let searchPos = 0; + + while (searchPos < content.length) { + const prefixPos = content.indexOf(READ_NOTE_PREFIX, searchPos); + if (prefixPos === -1) { + result += content.slice(searchPos); + break; + } + + // Add content before the prefix + result += content.slice(searchPos, prefixPos); + + const jsonStart = prefixPos + READ_NOTE_PREFIX.length; + const extracted = extractBalancedJson(content, jsonStart); + + if (!extracted) { + // No valid JSON found, keep original prefix and continue + result += READ_NOTE_PREFIX; + searchPos = jsonStart; + continue; + } + + try { + const parsed = JSON.parse(extracted.json); + if (parsed.content && parsed.content.length > threshold) { + const compactedResult = compactReadNoteResult(parsed, config); + result += `${READ_NOTE_PREFIX}${JSON.stringify(compactedResult)}`; + } else { + // Keep original if under threshold + result += READ_NOTE_PREFIX + extracted.json; + } + } catch { + // JSON parse failed (shouldn't happen with balanced extraction), keep original + result += READ_NOTE_PREFIX + extracted.json; + } + + searchPos = extracted.endPos; + } + + return result; +} + +/** + * Compact an XML tool result block. + */ +function compactToolResultBlock( + xmlBlock: string, + tag: string, + config: Partial = {} +): string { + // Never compact non-recoverable content + if (!isRecoverable(tag)) { + return xmlBlock; + } + + // Special handling for localSearch with multiple documents + if (tag === "localSearch") { + return compactLocalSearchBlock(xmlBlock, config); + } + + const source = extractSourceFromBlock(xmlBlock, tag); + const content = extractContentFromBlock(xmlBlock); + const sourceType = getSourceType(tag); + + const previewChars = + config.previewCharsPerSection ?? DEFAULT_COMPACTION_CONFIG.previewCharsPerSection; + const maxSections = config.maxSections ?? DEFAULT_COMPACTION_CONFIG.maxSections; + const compactedContent = compactBySection(content, previewChars, maxSections); + + return ` +${compactedContent} +`; +} + +/** + * Compact a localSearch block, preserving all documents with their paths. + * Each document gets its content compacted but title/path are preserved. + */ +function compactLocalSearchBlock(xmlBlock: string, config: Partial = {}): string { + const previewChars = + config.previewCharsPerSection ?? DEFAULT_COMPACTION_CONFIG.previewCharsPerSection; + + // Extract all document blocks + const documentRegex = /([\s\S]*?)<\/document>/g; + const documents: Array<{ path: string; title: string; preview: string }> = []; + + let match; + while ((match = documentRegex.exec(xmlBlock)) !== null) { + const docContent = match[1]; + + // Extract path and title from the document + const pathMatch = /([^<]+)<\/path>/.exec(docContent); + const titleMatch = /([^<]+)<\/title>/.exec(docContent); + const contentMatch = /<content>([\s\S]*?)<\/content>/.exec(docContent); + + const path = pathMatch?.[1] ?? ""; + const title = titleMatch?.[1] ?? path.split("/").pop() ?? "Untitled"; + const content = contentMatch?.[1] ?? ""; + + // Truncate content to preview length + const preview = + content.length > previewChars ? content.slice(0, previewChars) + "..." : content; + + documents.push({ path, title, preview: preview.trim() }); + } + + if (documents.length === 0) { + // Fallback: no documents found, use generic compaction + const content = extractContentFromBlock(xmlBlock); + const compactedContent = compactBySection( + content, + previewChars, + config.maxSections ?? DEFAULT_COMPACTION_CONFIG.maxSections + ); + return `<prior_context source="localSearch" type="note"> +${compactedContent} +</prior_context>`; + } + + // Build compacted output with all documents listed + const docList = documents + .map((doc, i) => `${i + 1}. [[${doc.title}]] (${doc.path})\n ${doc.preview}`) + .join("\n\n"); + + return `<prior_context source="localSearch" type="note"> +[${documents.length} search results - use localSearch to re-query] + +${docList} +</prior_context>`; +} + +/** + * Compact a readNote tool result by replacing full content with structure + preview. + */ +function compactReadNoteResult( + result: { + notePath?: string; + noteTitle?: string; + content?: string; + chunkIndex?: number; + totalChunks?: number; + [key: string]: unknown; + }, + config: Partial<CompactionConfig> = {} +): Record<string, unknown> { + const { content, notePath, noteTitle, ...rest } = result; + + if (!content || typeof content !== "string") { + return result; + } + + const previewChars = + config.previewCharsPerSection ?? DEFAULT_COMPACTION_CONFIG.previewCharsPerSection; + const maxSections = config.maxSections ?? DEFAULT_COMPACTION_CONFIG.maxSections; + const compactedContent = compactBySection(content, previewChars, maxSections); + + return { + ...rest, + notePath, + noteTitle, + content: `[COMPACTED - use readNote to get full content]\n\n${compactedContent}`, + _wasCompacted: true, + }; +} + +// Re-export for backwards compatibility during transition +export { compactAssistantOutput as compactChatHistoryContent }; diff --git a/src/context/L2ContextCompactor.test.ts b/src/context/L2ContextCompactor.test.ts new file mode 100644 index 00000000..f744ad78 --- /dev/null +++ b/src/context/L2ContextCompactor.test.ts @@ -0,0 +1,701 @@ +import { + compactL3ForL2, + compactBySection, + truncateWithEllipsis, + detectSourceType, + extractSource, + extractContent, + compactXmlBlock, + compactChatHistoryContent, + getL2RefetchInstruction, +} from "./L2ContextCompactor"; + +describe("L2ContextCompactor", () => { + describe("truncateWithEllipsis", () => { + it("should return text unchanged if under maxLength", () => { + const text = "Short text."; + expect(truncateWithEllipsis(text, 100)).toBe(text); + }); + + it("should truncate at sentence boundary when possible", () => { + const text = + "First sentence. Second sentence. Third sentence that goes on and on to make it longer."; + const result = truncateWithEllipsis(text, 50); + // Should break after "Second sentence." since it's the last sentence ending > 50% mark + expect(result).toContain("First sentence."); + expect(result).toContain("Second sentence."); + expect(result).toContain("..."); + expect(result.length).toBeLessThan(text.length); + }); + + it("should truncate at paragraph boundary when no good sentence break", () => { + const text = "First paragraph with no sentence breaks\n\nSecond paragraph here"; + const result = truncateWithEllipsis(text, 50); + expect(result).toBe("First paragraph with no sentence breaks\n\n..."); + }); + + it("should truncate at word boundary as fallback", () => { + const text = "word1 word2 word3 word4 word5 word6 word7 word8"; + const result = truncateWithEllipsis(text, 25); + expect(result).toBe("word1 word2 word3 word4 ..."); + }); + + it("should add ellipsis when hard truncating", () => { + const text = "verylongwordwithoutanyspaces"; + const result = truncateWithEllipsis(text, 10); + expect(result).toBe("verylongwo..."); + }); + }); + + describe("compactBySection", () => { + it("should truncate content without headings", () => { + const content = "A".repeat(3000); + const result = compactBySection(content, 500, 20); + expect(result.length).toBeLessThan(2500); + expect(result).toContain("..."); + }); + + it("should preserve headings and extract previews from each section", () => { + const content = `## Introduction +This is the introduction section with some content that explains what the document is about. It goes on for a while with more details. + +## Methodology +This describes the methodology used in the research. It includes various approaches and techniques. + +## Results +Here are the results of the study. They show interesting findings.`; + + const result = compactBySection(content, 50, 20); + + expect(result).toContain("## Introduction"); + expect(result).toContain("## Methodology"); + expect(result).toContain("## Results"); + // Should truncate long sections + expect(result.length).toBeLessThan(content.length); + }); + + it("should keep small sections verbatim", () => { + const content = `## Section 1 +Short content. + +## Section 2 +Also short.`; + + const result = compactBySection(content, 500, 20); + expect(result).toBe(content); + }); + + it("should limit number of sections", () => { + const sections = Array.from( + { length: 30 }, + (_, i) => `## Section ${i + 1}\nContent ${i + 1}` + ); + const content = sections.join("\n\n"); + + const result = compactBySection(content, 500, 10); + + expect(result).toContain("## Section 1"); + expect(result).toContain("## Section 10"); + expect(result).not.toContain("## Section 11"); + expect(result).toContain("20 more sections omitted"); + }); + + it("should handle nested headings", () => { + const content = `# Main Title +Introduction text. + +## Section 1 +Content for section 1. + +### Subsection 1.1 +Details here. + +## Section 2 +More content.`; + + const result = compactBySection(content, 500, 20); + + expect(result).toContain("# Main Title"); + expect(result).toContain("## Section 1"); + expect(result).toContain("### Subsection 1.1"); + expect(result).toContain("## Section 2"); + }); + + it("should handle content with code blocks", () => { + const content = `## Code Example +Here is some code: + +\`\`\`javascript +function example() { + return "hello"; +} +\`\`\` + +## Another Section +More text here.`; + + const result = compactBySection(content, 500, 20); + + expect(result).toContain("## Code Example"); + expect(result).toContain("```javascript"); + expect(result).toContain("## Another Section"); + }); + }); + + describe("detectSourceType", () => { + it("should detect note types", () => { + expect(detectSourceType("note_context")).toBe("note"); + expect(detectSourceType("active_note")).toBe("note"); + expect(detectSourceType("embedded_note")).toBe("note"); + expect(detectSourceType("vault_note")).toBe("note"); + expect(detectSourceType("retrieved_document")).toBe("note"); + }); + + it("should detect URL types", () => { + expect(detectSourceType("url_content")).toBe("url"); + expect(detectSourceType("web_tab_context")).toBe("url"); + expect(detectSourceType("active_web_tab")).toBe("url"); + }); + + it("should detect YouTube type", () => { + expect(detectSourceType("youtube_video_context")).toBe("youtube"); + }); + + it("should detect PDF type", () => { + expect(detectSourceType("embedded_pdf")).toBe("pdf"); + }); + + it("should detect selected text types", () => { + expect(detectSourceType("selected_text")).toBe("selected_text"); + expect(detectSourceType("web_selected_text")).toBe("selected_text"); + }); + + it("should return unknown for unrecognized types", () => { + expect(detectSourceType("random_type")).toBe("unknown"); + expect(detectSourceType("")).toBe("unknown"); + }); + }); + + describe("extractSource", () => { + it("should extract path from note context", () => { + const xml = `<note_context> +<title>My Note +folder/my-note.md +Content here +`; + expect(extractSource(xml)).toBe("folder/my-note.md"); + }); + + it("should extract URL from web content", () => { + const xml = ` +Web Page +https://example.com/page +Content here +`; + expect(extractSource(xml)).toBe("https://example.com/page"); + }); + + it("should extract name from PDF", () => { + const xml = ` +document.pdf +PDF content +`; + expect(extractSource(xml)).toBe("document.pdf"); + }); + + it("should return empty string if no source found", () => { + const xml = ` +Just some text +`; + expect(extractSource(xml)).toBe(""); + }); + + it("should prefer path over url if both present", () => { + const xml = ` +local/file.md +https://example.com +`; + expect(extractSource(xml)).toBe("local/file.md"); + }); + }); + + describe("extractContent", () => { + it("should extract content from content tags", () => { + const xml = ` +Title +path.md +This is the actual content +with multiple lines. +`; + expect(extractContent(xml)).toBe("This is the actual content\nwith multiple lines."); + }); + + it("should return whole block if no content tags", () => { + const xml = "Just plain text"; + expect(extractContent(xml)).toBe(xml); + }); + + it("should handle content with special characters", () => { + const xml = ` +Content with and "quotes" & ampersands +`; + expect(extractContent(xml)).toBe('Content with and "quotes" & ampersands'); + }); + }); + + describe("compactL3ForL2", () => { + it("should keep small content verbatim", () => { + const content = "Small content"; + const result = compactL3ForL2(content, "notes/test.md", "note"); + expect(result).toBe(content); + }); + + it("should compact large content with proper structure", () => { + const content = `## Section 1 +${"A".repeat(2000)} + +## Section 2 +${"B".repeat(2000)}`; + + const result = compactL3ForL2(content, "notes/research.md", "note", { + verbatimThreshold: 1000, + previewCharsPerSection: 100, + }); + + expect(result).toContain(''); + expect(result).toContain("## Section 1"); + expect(result).toContain("## Section 2"); + expect(result).toContain(""); + expect(result.length).toBeLessThan(content.length); + }); + + it("should set correct type for URL content", () => { + const content = "A".repeat(10000); + const result = compactL3ForL2(content, "https://example.com", "url", { + verbatimThreshold: 1000, + }); + + expect(result).toContain('type="url"'); + expect(result).toContain('source="https://example.com"'); + }); + + it("should set correct type for YouTube content", () => { + const content = "A".repeat(10000); + const result = compactL3ForL2(content, "https://youtube.com/watch?v=123", "youtube", { + verbatimThreshold: 1000, + }); + + expect(result).toContain('type="youtube"'); + expect(result).toContain('source="https://youtube.com/watch?v=123"'); + }); + + it("should escape special characters in source attribute", () => { + const content = "A".repeat(10000); + const result = compactL3ForL2(content, 'path/with"quotes&special', "note", { + verbatimThreshold: 1000, + }); + + expect(result).toContain("""); + expect(result).toContain("&"); + expect(result).toContain("<"); + }); + }); + + describe("compactXmlBlock", () => { + it("should keep small XML blocks verbatim", () => { + const xml = ` +Short Note +notes/short.md +Brief content. +`; + + const result = compactXmlBlock(xml, "note_context"); + expect(result).toBe(xml); + }); + + it("should compact large XML blocks", () => { + const largeContent = `## Introduction +${"X".repeat(3000)} + +## Details +${"Y".repeat(3000)}`; + + const xml = ` +Research Paper +research/paper.md +${largeContent} +`; + + const result = compactXmlBlock(xml, "note_context", { + verbatimThreshold: 1000, + previewCharsPerSection: 200, + }); + + expect(result).toContain(''); + expect(result).toContain("## Introduction"); + expect(result).toContain("## Details"); + expect(result.length).toBeLessThan(xml.length); + }); + + it("should handle URL content blocks", () => { + const largeContent = "A".repeat(10000); + const xml = ` +Web Page +https://example.com/article +${largeContent} +`; + + const result = compactXmlBlock(xml, "url_content", { + verbatimThreshold: 1000, + }); + + expect(result).toContain('source="https://example.com/article"'); + expect(result).toContain('type="url"'); + }); + + it("should handle YouTube blocks", () => { + const transcript = "Speaker: " + "transcript content ".repeat(500); + const xml = ` +Tutorial Video +https://www.youtube.com/watch?v=abc123 +${transcript} +`; + + const result = compactXmlBlock(xml, "youtube_video_context", { + verbatimThreshold: 1000, + }); + + expect(result).toContain('type="youtube"'); + }); + + it("should handle PDF blocks", () => { + const pdfContent = "Page 1: " + "content ".repeat(1000); + const xml = ` +documents/report.pdf +${pdfContent} +`; + + const result = compactXmlBlock(xml, "embedded_pdf", { + verbatimThreshold: 1000, + }); + + expect(result).toContain('source="documents/report.pdf"'); + expect(result).toContain('type="pdf"'); + }); + + it("should NEVER compact selected_text blocks regardless of size", () => { + const largeContent = "A".repeat(20000); + const xml = ` +User Selection +notes/document.md +${largeContent} +`; + + const result = compactXmlBlock(xml, "selected_text", { + verbatimThreshold: 1000, + }); + + // Should return verbatim, not compacted + expect(result).toBe(xml); + expect(result).not.toContain("prior_context"); + }); + + it("should NEVER compact web_selected_text blocks regardless of size", () => { + const largeContent = "B".repeat(20000); + const xml = ` +Web Selection +https://example.com/page +${largeContent} +`; + + const result = compactXmlBlock(xml, "web_selected_text", { + verbatimThreshold: 1000, + }); + + // Should return verbatim, not compacted + expect(result).toBe(xml); + expect(result).not.toContain("prior_context"); + }); + }); + + describe("getL2RefetchInstruction", () => { + it("should return a single consolidated instruction", () => { + const instruction = getL2RefetchInstruction(); + + expect(instruction).toContain("prior_context_note"); + expect(instruction).toContain("previews"); + expect(instruction).toContain("[[note title]]"); + }); + }); + + describe("real-world scenarios", () => { + it("should handle a typical research note with large sections", () => { + // Generate a realistic large research note where each section has substantial content + const longParagraph = (topic: string) => + `This section discusses ${topic} in great detail. `.repeat(30); + + const content = `## Abstract +${longParagraph("the research overview")} + +## 1. Introduction +${longParagraph("the background and motivation")} + +### 1.1 Background +${longParagraph("historical context")} + +### 1.2 Motivation +${longParagraph("why this research matters")} + +## 2. Methodology +${longParagraph("experimental design")} + +### 2.1 Data Collection +${longParagraph("data gathering procedures")} + +### 2.2 Model Architectures +${longParagraph("model implementations")} + +## 3. Results +${longParagraph("experimental findings")} + +### 3.1 Classification Tasks +${longParagraph("classification results")} + +### 3.2 Generation Tasks +${longParagraph("generation results")} + +## 4. Discussion +${longParagraph("implications of findings")} + +## 5. Conclusion +${longParagraph("summary and future work")}`; + + const xml = ` +ML Research Paper +research/ml-nlp-study.md +${content} +`; + + const result = compactXmlBlock(xml, "note_context", { + verbatimThreshold: 1000, + previewCharsPerSection: 300, + }); + + // Should preserve all major sections + expect(result).toContain("## Abstract"); + expect(result).toContain("## 1. Introduction"); + expect(result).toContain("### 1.1 Background"); + expect(result).toContain("## 2. Methodology"); + expect(result).toContain("## 3. Results"); + expect(result).toContain("## 4. Discussion"); + expect(result).toContain("## 5. Conclusion"); + + // Should be significantly smaller (each section was ~1500 chars, now ~300) + expect(result.length).toBeLessThan(xml.length * 0.3); + + // Should be wrapped in prior_context + expect(result).toContain("prior_context"); + }); + + it("should handle content with mixed markdown elements", () => { + const content = `## Overview +This document contains various markdown elements. + +## Code Examples +Here's a code block: + +\`\`\`python +def hello(): + print("Hello, World!") + +def complex_function(x, y, z): + # This is a longer function + result = x + y + z + for i in range(100): + result += i + return result +\`\`\` + +## Tables +| Column 1 | Column 2 | Column 3 | +|----------|----------|----------| +| Data 1 | Data 2 | Data 3 | +| More | Data | Here | + +## Lists +- Item 1 +- Item 2 + - Nested item + - Another nested +- Item 3 + +## Blockquotes +> This is a quote +> that spans multiple lines +> and contains important information. + +## Final Section +Some concluding remarks.`; + + const result = compactBySection(content, 400, 20); + + // Should preserve structure + expect(result).toContain("## Overview"); + expect(result).toContain("## Code Examples"); + expect(result).toContain("```python"); + expect(result).toContain("## Tables"); + expect(result).toContain("## Lists"); + expect(result).toContain("## Blockquotes"); + expect(result).toContain("## Final Section"); + }); + }); + + describe("compactChatHistoryContent", () => { + it("should return small content unchanged", () => { + const content = "This is a short message without any tool results."; + const result = compactChatHistoryContent(content); + expect(result).toBe(content); + }); + + it("should compact large localSearch blocks in chat history", () => { + const largeContent = "A".repeat(10000); + const content = `I'll search for that information. + +Tool 'localSearch' result: +Answer the question based only on the following context: + + +${largeContent} + + + +CITATION RULES: +1. START with [^1] and increment sequentially + + + +Based on my search, I found the answer.`; + + const result = compactChatHistoryContent(content, { verbatimThreshold: 1000 }); + + expect(typeof result).toBe("string"); + expect((result as string).length).toBeLessThan(content.length); + expect(result).toContain("I'll search for that information"); + expect(result).toContain("Based on my search, I found the answer"); + expect(result).toContain("prior_context"); + }); + + it("should compact large readNote JSON results", () => { + const largeNoteContent = `## Introduction +${"This is introduction content. ".repeat(200)} + +## Methods +${"This describes the methodology. ".repeat(200)} + +## Results +${"Here are the results. ".repeat(200)}`; + + const readNoteResult = JSON.stringify({ + notePath: "research/paper.md", + noteTitle: "Research Paper", + content: largeNoteContent, + chunkIndex: 0, + totalChunks: 1, + }); + + const content = `I'll read that note for you. + +Tool 'readNote' result: ${readNoteResult} + +Here's what I found in the note.`; + + const result = compactChatHistoryContent(content, { verbatimThreshold: 1000 }); + + expect(typeof result).toBe("string"); + expect((result as string).length).toBeLessThan(content.length); + expect(result).toContain("I'll read that note for you"); + expect(result).toContain("Here's what I found in the note"); + expect(result).toContain("COMPACTED"); + expect(result).toContain("## Introduction"); + expect(result).toContain("## Methods"); + }); + + it("should handle multimodal content arrays", () => { + const largeContent = "A".repeat(10000); + const originalText = `Tool 'localSearch' result: ${largeContent}`; + const content = [ + { type: "text", text: originalText }, + { type: "image_url", url: "data:image/png;base64,..." }, + ]; + + const result = compactChatHistoryContent(content, { verbatimThreshold: 1000 }); + + expect(Array.isArray(result)).toBe(true); + const resultArray = result as any[]; + expect(resultArray[0].type).toBe("text"); + expect(resultArray[0].text.length).toBeLessThan(originalText.length); + expect(resultArray[1]).toEqual(content[1]); // Image unchanged + }); + + it("should keep small tool results verbatim", () => { + const smallContent = "Brief note content."; + const readNoteResult = JSON.stringify({ + notePath: "notes/brief.md", + noteTitle: "Brief Note", + content: smallContent, + chunkIndex: 0, + totalChunks: 1, + }); + + const content = `Tool 'readNote' result: ${readNoteResult}`; + const result = compactChatHistoryContent(content); + + expect(result).toBe(content); // Unchanged + }); + + it("should handle multiple tool results in one message", () => { + const largeContent1 = "A".repeat(8000); + const largeContent2 = "B".repeat(8000); + + const content = `First search: + +Tool 'localSearch' result: ${largeContent1} + +Second search: + +Tool 'localSearch' result: ${largeContent2} + +Summary of findings.`; + + const result = compactChatHistoryContent(content, { verbatimThreshold: 1000 }); + + expect(typeof result).toBe("string"); + expect((result as string).length).toBeLessThan(content.length * 0.3); + expect(result).toContain("First search"); + expect(result).toContain("Second search"); + expect(result).toContain("Summary of findings"); + }); + + it("should compact note_context blocks", () => { + const largeContent = `## Section 1 +${"Content for section 1. ".repeat(200)} + +## Section 2 +${"Content for section 2. ".repeat(200)}`; + + const content = ` +My Note +notes/mynote.md +${largeContent} +`; + + const result = compactChatHistoryContent(content, { verbatimThreshold: 1000 }); + + expect(typeof result).toBe("string"); + expect((result as string).length).toBeLessThan(content.length); + expect(result).toContain("prior_context"); + expect(result).toContain("## Section 1"); + expect(result).toContain("## Section 2"); + }); + }); +}); diff --git a/src/context/L2ContextCompactor.ts b/src/context/L2ContextCompactor.ts new file mode 100644 index 00000000..40f08f44 --- /dev/null +++ b/src/context/L2ContextCompactor.ts @@ -0,0 +1,139 @@ +/** + * L2ContextCompactor - Compacts L3 context for inclusion in L2 (previous turn context). + * + * When context from previous turns (L3) flows into L2, we don't need the full verbatim + * content. Instead, we extract structure and previews. A single instruction at the end + * of L2 tells the LLM how to re-fetch full content if needed. + * + * This is NOT LLM-based summarization - it's deterministic extraction that: + * 1. Preserves document structure (headings) + * 2. Keeps preview content from each section + * 3. Includes source path/URL for re-fetching + * + * Typical compression: 100KB -> 3-5KB (95%+ reduction) + */ + +import { + ContextSourceType, + extractContentFromBlock, + extractSourceFromBlock, + getSourceType, + isRecoverable, +} from "./contextBlockRegistry"; +import { + CompactionConfig, + DEFAULT_COMPACTION_CONFIG, + compactBySection, + escapeXmlAttr, +} from "./compactionUtils"; + +// Re-export types and utilities for backwards compatibility +export type { ContextSourceType } from "./contextBlockRegistry"; +export type { CompactionConfig as L2CompactorConfig } from "./compactionUtils"; +export { compactBySection, truncateWithEllipsis } from "./compactionUtils"; +export { getSourceType as detectSourceType } from "./contextBlockRegistry"; + +// Re-export chat history compaction for backwards compatibility +export { compactAssistantOutput, compactChatHistoryContent } from "./ChatHistoryCompactor"; + +/** + * Extract the source identifier from an XML block. + * @deprecated Use extractSourceFromBlock from contextBlockRegistry instead + */ +export function extractSource(xmlBlock: string): string { + // Try common source extractors in order + const pathMatch = /([^<]+)<\/path>/.exec(xmlBlock); + if (pathMatch) return pathMatch[1]; + + const urlMatch = /([^<]+)<\/url>/.exec(xmlBlock); + if (urlMatch) return urlMatch[1]; + + const nameMatch = /([^<]+)<\/name>/.exec(xmlBlock); + if (nameMatch) return nameMatch[1]; + + return ""; +} + +/** + * Extract the inner content from an XML block. + * @deprecated Use extractContentFromBlock from contextBlockRegistry instead + */ +export function extractContent(xmlBlock: string): string { + return extractContentFromBlock(xmlBlock); +} + +/** + * Compact a single L3 segment's content for inclusion in L2. + * + * @param content - The full content from the L3 segment + * @param source - The source path/URL (for display and re-fetching) + * @param sourceType - The type of source + * @param config - Optional configuration overrides + * @returns Compacted content with structure and previews + */ +export function compactL3ForL2( + content: string, + source: string, + sourceType: ContextSourceType, + config: Partial = {} +): string { + const threshold = config.verbatimThreshold ?? DEFAULT_COMPACTION_CONFIG.verbatimThreshold; + + // Keep small content verbatim + if (content.length <= threshold) { + return content; + } + + const previewChars = + config.previewCharsPerSection ?? DEFAULT_COMPACTION_CONFIG.previewCharsPerSection; + const maxSections = config.maxSections ?? DEFAULT_COMPACTION_CONFIG.maxSections; + const compactedContent = compactBySection(content, previewChars, maxSections); + + return ` +${compactedContent} +`; +} + +/** + * Compact an entire XML context block for L2. + * This is the main entry point for compacting L3 segments. + * + * @param xmlBlock - The full XML block (e.g., ...) + * @param blockType - The XML tag type (e.g., "note_context") + * @param config - Optional configuration overrides + * @returns Compacted block ready for L2 + */ +export function compactXmlBlock( + xmlBlock: string, + blockType: string, + config: Partial = {} +): string { + // Never compact non-recoverable content (e.g., selected_text) + if (!isRecoverable(blockType)) { + return xmlBlock; + } + + const threshold = config.verbatimThreshold ?? DEFAULT_COMPACTION_CONFIG.verbatimThreshold; + + // Keep small blocks verbatim + if (xmlBlock.length <= threshold) { + return xmlBlock; + } + + const source = extractSourceFromBlock(xmlBlock, blockType) || extractSource(xmlBlock); + const content = extractContentFromBlock(xmlBlock); + const sourceType = getSourceType(blockType); + + return compactL3ForL2(content, source, sourceType, config); +} + +/** + * Returns a single instruction to append after all L2 context, explaining how to + * re-fetch compacted content. This is more efficient than per-block hints. + */ +export function getL2RefetchInstruction(): string { + return ` +The above prior_context blocks contain previews of content from earlier turns. +To access full content: use [[note title]] for notes, or ask to read a specific URL/video. +`; +} diff --git a/src/context/compactionUtils.test.ts b/src/context/compactionUtils.test.ts new file mode 100644 index 00000000..b2e446ef --- /dev/null +++ b/src/context/compactionUtils.test.ts @@ -0,0 +1,101 @@ +import { + compactBySection, + truncateWithEllipsis, + escapeXmlAttr, + mergeConfig, + DEFAULT_COMPACTION_CONFIG, +} from "./compactionUtils"; + +describe("compactionUtils", () => { + describe("DEFAULT_COMPACTION_CONFIG", () => { + it("should have expected default values", () => { + expect(DEFAULT_COMPACTION_CONFIG.previewCharsPerSection).toBe(500); + expect(DEFAULT_COMPACTION_CONFIG.maxSections).toBe(20); + expect(DEFAULT_COMPACTION_CONFIG.verbatimThreshold).toBe(5000); + }); + }); + + describe("mergeConfig", () => { + it("should use defaults when no config provided", () => { + const config = mergeConfig(); + expect(config).toEqual(DEFAULT_COMPACTION_CONFIG); + }); + + it("should override specific values", () => { + const config = mergeConfig({ previewCharsPerSection: 100 }); + expect(config.previewCharsPerSection).toBe(100); + expect(config.maxSections).toBe(DEFAULT_COMPACTION_CONFIG.maxSections); + }); + }); + + describe("truncateWithEllipsis", () => { + it("should return text unchanged if under maxLength", () => { + expect(truncateWithEllipsis("Short.", 100)).toBe("Short."); + }); + + it("should truncate at sentence boundary when possible", () => { + const text = "First sentence. Second sentence. Third sentence is much longer."; + const result = truncateWithEllipsis(text, 50); + expect(result).toContain("First sentence."); + expect(result).toContain("..."); + }); + + it("should truncate at word boundary as fallback", () => { + const text = "word1 word2 word3 word4 word5"; + const result = truncateWithEllipsis(text, 20); + expect(result).toContain("..."); + expect(result.length).toBeLessThanOrEqual(25); // 20 + ellipsis + }); + }); + + describe("compactBySection", () => { + it("should preserve all headings", () => { + const content = `## Section 1 +Content 1 + +## Section 2 +Content 2`; + const result = compactBySection(content, 500, 20); + expect(result).toContain("## Section 1"); + expect(result).toContain("## Section 2"); + }); + + it("should truncate long sections", () => { + const content = `## Section 1 +${"A".repeat(1000)}`; + const result = compactBySection(content, 100, 20); + expect(result).toContain("## Section 1"); + expect(result.length).toBeLessThan(content.length); + }); + + it("should limit number of sections", () => { + const sections = Array.from({ length: 30 }, (_, i) => `## Section ${i}\nContent`); + const content = sections.join("\n\n"); + const result = compactBySection(content, 500, 10); + expect(result).toContain("## Section 0"); + expect(result).toContain("## Section 9"); + expect(result).not.toContain("## Section 10"); + expect(result).toContain("more sections omitted"); + }); + + it("should handle content without headings", () => { + const content = "A".repeat(5000); + const result = compactBySection(content, 500, 20); + expect(result.length).toBeLessThan(content.length); + expect(result).toContain("..."); + }); + }); + + describe("escapeXmlAttr", () => { + it("should escape special XML characters", () => { + expect(escapeXmlAttr('test "quoted"')).toBe("test "quoted""); + expect(escapeXmlAttr("test & ampersand")).toBe("test & ampersand"); + expect(escapeXmlAttr("test ")).toBe("test <tag>"); + expect(escapeXmlAttr("test 'apostrophe'")).toBe("test 'apostrophe'"); + }); + + it("should handle empty string", () => { + expect(escapeXmlAttr("")).toBe(""); + }); + }); +}); diff --git a/src/context/compactionUtils.ts b/src/context/compactionUtils.ts new file mode 100644 index 00000000..ebe10b3f --- /dev/null +++ b/src/context/compactionUtils.ts @@ -0,0 +1,135 @@ +/** + * Shared utilities for content compaction. + * + * These are pure functions used by both L2ContextCompactor and ChatHistoryCompactor + * for truncating and extracting structure from content. + */ + +/** + * Configuration for compaction operations + */ +export interface CompactionConfig { + /** Characters to keep per section (default: 500) */ + previewCharsPerSection: number; + /** Max total sections to include (default: 20) */ + maxSections: number; + /** Threshold below which content is kept verbatim (default: 5000) */ + verbatimThreshold: number; +} + +export const DEFAULT_COMPACTION_CONFIG: CompactionConfig = { + previewCharsPerSection: 500, + maxSections: 20, + verbatimThreshold: 5000, +}; + +/** + * Merge partial config with defaults + */ +export function mergeConfig(config: Partial = {}): CompactionConfig { + return { ...DEFAULT_COMPACTION_CONFIG, ...config }; +} + +/** + * Compact content by extracting headings and preview from each section. + * + * @param content - The full content to compact + * @param previewCharsPerSection - Max chars to keep per section + * @param maxSections - Max number of sections to include + * @returns Compacted content with structure preserved + */ +export function compactBySection( + content: string, + previewCharsPerSection = 500, + maxSections = 20 +): string { + // Split content by markdown headings (keep the heading with its section) + const sections = content.split(/(?=^#{1,6}\s+)/m).filter((s) => s.trim()); + + if (sections.length <= 1) { + // No headings - just truncate intelligently + return truncateWithEllipsis(content, previewCharsPerSection * 4); + } + + // Limit number of sections + const limitedSections = sections.slice(0, maxSections); + const hasMoreSections = sections.length > maxSections; + + const compacted = limitedSections + .map((section) => { + const lines = section.trim().split("\n"); + const heading = lines[0]; // The heading line + const body = lines.slice(1).join("\n").trim(); + + if (body.length <= previewCharsPerSection) { + // Keep full section if small enough + return section.trim(); + } + + // Keep heading + truncated body + return `${heading}\n${truncateWithEllipsis(body, previewCharsPerSection)}`; + }) + .join("\n\n"); + + if (hasMoreSections) { + return `${compacted}\n\n[... ${sections.length - maxSections} more sections omitted ...]`; + } + + return compacted; +} + +/** + * Truncate text at a sensible boundary (sentence or paragraph) with ellipsis. + * + * @param text - Text to truncate + * @param maxLength - Maximum length + * @returns Truncated text with ellipsis if truncated + */ +export function truncateWithEllipsis(text: string, maxLength: number): string { + if (text.length <= maxLength) { + return text; + } + + const truncated = text.slice(0, maxLength); + + // Try to break at sentence boundary - find last complete sentence + // Look for sentence-ending punctuation followed by space + const sentenceEndPattern = /[.!?]\s+/g; + let lastSentenceEnd = -1; + let match; + while ((match = sentenceEndPattern.exec(truncated)) !== null) { + // Only count if we're past 50% of the text (avoid breaking too early) + if (match.index > maxLength * 0.5) { + lastSentenceEnd = match.index + 1; // Include the punctuation + } + } + if (lastSentenceEnd > 0) { + return truncated.slice(0, lastSentenceEnd) + " ..."; + } + + // Try to break at paragraph boundary + const lastParagraph = truncated.lastIndexOf("\n\n"); + if (lastParagraph > maxLength * 0.5) { + return truncated.slice(0, lastParagraph) + "\n\n..."; + } + + // Fall back to word boundary + const lastSpace = truncated.lastIndexOf(" "); + if (lastSpace > maxLength * 0.8) { + return truncated.slice(0, lastSpace) + " ..."; + } + + return truncated + "..."; +} + +/** + * Escape a string for use in an XML attribute value. + */ +export function escapeXmlAttr(str: string): string { + return str + .replace(/&/g, "&") + .replace(/"/g, """) + .replace(/'/g, "'") + .replace(//g, ">"); +} diff --git a/src/context/contextBlockRegistry.test.ts b/src/context/contextBlockRegistry.test.ts new file mode 100644 index 00000000..c4f05949 --- /dev/null +++ b/src/context/contextBlockRegistry.test.ts @@ -0,0 +1,143 @@ +import { + getBlockType, + getSourceType, + isRecoverable, + getNeverCompactTags, + extractSourceFromBlock, + extractContentFromBlock, + detectBlockTag, + CONTEXT_BLOCK_TYPES, +} from "./contextBlockRegistry"; + +describe("contextBlockRegistry", () => { + describe("CONTEXT_BLOCK_TYPES", () => { + it("should have all expected block types registered", () => { + const tags = CONTEXT_BLOCK_TYPES.map((bt) => bt.tag); + expect(tags).toContain("note_context"); + expect(tags).toContain("active_note"); + expect(tags).toContain("url_content"); + expect(tags).toContain("youtube_video_context"); + expect(tags).toContain("selected_text"); + expect(tags).toContain("localSearch"); + }); + }); + + describe("getBlockType", () => { + it("should return block type for known tags", () => { + const noteContext = getBlockType("note_context"); + expect(noteContext).toBeDefined(); + expect(noteContext?.sourceType).toBe("note"); + expect(noteContext?.recoverable).toBe(true); + }); + + it("should return undefined for unknown tags", () => { + expect(getBlockType("unknown_tag")).toBeUndefined(); + }); + }); + + describe("getSourceType", () => { + it("should return correct source type for note blocks", () => { + expect(getSourceType("note_context")).toBe("note"); + expect(getSourceType("active_note")).toBe("note"); + expect(getSourceType("embedded_note")).toBe("note"); + }); + + it("should return correct source type for URL blocks", () => { + expect(getSourceType("url_content")).toBe("url"); + expect(getSourceType("web_tab_context")).toBe("url"); + }); + + it("should return correct source type for YouTube blocks", () => { + expect(getSourceType("youtube_video_context")).toBe("youtube"); + }); + + it("should return unknown for unregistered tags", () => { + expect(getSourceType("random_tag")).toBe("unknown"); + }); + }); + + describe("isRecoverable", () => { + it("should return true for recoverable types", () => { + expect(isRecoverable("note_context")).toBe(true); + expect(isRecoverable("url_content")).toBe(true); + expect(isRecoverable("youtube_video_context")).toBe(true); + }); + + it("should return false for non-recoverable types", () => { + expect(isRecoverable("selected_text")).toBe(false); + expect(isRecoverable("web_selected_text")).toBe(false); + }); + + it("should return false for unknown types", () => { + expect(isRecoverable("unknown_type")).toBe(false); + }); + }); + + describe("getNeverCompactTags", () => { + it("should return set of non-recoverable tags", () => { + const tags = getNeverCompactTags(); + expect(tags.has("selected_text")).toBe(true); + expect(tags.has("web_selected_text")).toBe(true); + expect(tags.has("note_context")).toBe(false); + }); + }); + + describe("extractSourceFromBlock", () => { + it("should extract path from note blocks", () => { + const xml = ` +My Note +folder/my-note.md +Content here +`; + expect(extractSourceFromBlock(xml, "note_context")).toBe("folder/my-note.md"); + }); + + it("should extract URL from url_content blocks", () => { + const xml = ` +https://example.com/page +Content +`; + expect(extractSourceFromBlock(xml, "url_content")).toBe("https://example.com/page"); + }); + + it("should extract name from PDF blocks", () => { + const xml = ` +document.pdf +PDF content +`; + expect(extractSourceFromBlock(xml, "embedded_pdf")).toBe("document.pdf"); + }); + + it("should return empty string for blocks without source extractor", () => { + const xml = `Just text`; + expect(extractSourceFromBlock(xml, "selected_text")).toBe(""); + }); + }); + + describe("extractContentFromBlock", () => { + it("should extract content from content tags", () => { + const xml = ` +Title +This is the content +`; + expect(extractContentFromBlock(xml)).toBe("This is the content"); + }); + + it("should return whole block if no content tags", () => { + const xml = "Plain text"; + expect(extractContentFromBlock(xml)).toBe(xml); + }); + }); + + describe("detectBlockTag", () => { + it("should detect tag from block start", () => { + expect(detectBlockTag("...")).toBe("note_context"); + expect(detectBlockTag('...')).toBe("url_content"); + }); + + it("should return null for invalid blocks", () => { + expect(detectBlockTag("not xml")).toBeNull(); + expect(detectBlockTag("")).toBeNull(); + }); + }); +}); diff --git a/src/context/contextBlockRegistry.ts b/src/context/contextBlockRegistry.ts new file mode 100644 index 00000000..05b8e63b --- /dev/null +++ b/src/context/contextBlockRegistry.ts @@ -0,0 +1,144 @@ +/** + * Unified registry for all context block types used in the plugin. + * + * This is the single source of truth for: + * - XML tag names used for context blocks + * - Source types (note, url, youtube, etc.) + * - Whether content is recoverable (can LLM re-fetch it?) + * - How to extract the source identifier from each block type + */ + +/** + * Context source types that determine which tool to use for re-fetching + */ +export type ContextSourceType = "note" | "url" | "youtube" | "pdf" | "selected_text" | "unknown"; + +/** + * How to extract the source identifier from an XML block + */ +export type SourceExtractor = "path" | "url" | "name" | null; + +/** + * Metadata for a context block type + */ +export interface ContextBlockType { + /** XML tag name (e.g., "note_context", "url_content") */ + tag: string; + /** Category of source for re-fetching hints */ + sourceType: ContextSourceType; + /** Whether the LLM can programmatically re-fetch this content */ + recoverable: boolean; + /** Which XML child element contains the source identifier */ + sourceExtractor: SourceExtractor; +} + +/** + * All registered context block types + */ +export const CONTEXT_BLOCK_TYPES: ContextBlockType[] = [ + // Note-based content (recoverable via readNote or wiki-links) + { tag: "note_context", sourceType: "note", recoverable: true, sourceExtractor: "path" }, + { tag: "active_note", sourceType: "note", recoverable: true, sourceExtractor: "path" }, + { tag: "embedded_note", sourceType: "note", recoverable: true, sourceExtractor: "path" }, + { tag: "vault_note", sourceType: "note", recoverable: true, sourceExtractor: "path" }, + { tag: "retrieved_document", sourceType: "note", recoverable: true, sourceExtractor: "path" }, + + // Web content (recoverable via /url command) + { tag: "url_content", sourceType: "url", recoverable: true, sourceExtractor: "url" }, + { tag: "web_tab_context", sourceType: "url", recoverable: true, sourceExtractor: "url" }, + { tag: "active_web_tab", sourceType: "url", recoverable: true, sourceExtractor: "url" }, + + // YouTube (recoverable via /yt command) + { + tag: "youtube_video_context", + sourceType: "youtube", + recoverable: true, + sourceExtractor: "url", + }, + + // PDF (partially recoverable - needs the file in vault) + { tag: "embedded_pdf", sourceType: "pdf", recoverable: true, sourceExtractor: "name" }, + + // Selected text (NOT recoverable - user must manually re-select) + { tag: "selected_text", sourceType: "selected_text", recoverable: false, sourceExtractor: null }, + { + tag: "web_selected_text", + sourceType: "selected_text", + recoverable: false, + sourceExtractor: null, + }, + + // Tool results (recoverable by re-running the tool) + { tag: "localSearch", sourceType: "note", recoverable: true, sourceExtractor: null }, +]; + +// Pre-computed lookup maps for performance +const blockTypeByTag = new Map(); +for (const blockType of CONTEXT_BLOCK_TYPES) { + blockTypeByTag.set(blockType.tag, blockType); +} + +/** + * Get block type metadata by tag name + */ +export function getBlockType(tag: string): ContextBlockType | undefined { + return blockTypeByTag.get(tag); +} + +/** + * Get the source type for a given tag name + */ +export function getSourceType(tag: string): ContextSourceType { + return blockTypeByTag.get(tag)?.sourceType ?? "unknown"; +} + +/** + * Check if a block type is recoverable (can be re-fetched by the LLM) + */ +export function isRecoverable(tag: string): boolean { + return blockTypeByTag.get(tag)?.recoverable ?? false; +} + +/** + * Get all tags that should never be compacted + */ +export function getNeverCompactTags(): Set { + const tags = new Set(); + for (const blockType of CONTEXT_BLOCK_TYPES) { + if (!blockType.recoverable) { + tags.add(blockType.tag); + } + } + return tags; +} + +/** + * Extract the source identifier from an XML block based on the block type + */ +export function extractSourceFromBlock(xmlBlock: string, tag: string): string { + const blockType = blockTypeByTag.get(tag); + if (!blockType?.sourceExtractor) { + return ""; + } + + const extractorTag = blockType.sourceExtractor; + const regex = new RegExp(`<${extractorTag}>([^<]+)`); + const match = regex.exec(xmlBlock); + return match?.[1] ?? ""; +} + +/** + * Extract the content from an XML block (inside tags) + */ +export function extractContentFromBlock(xmlBlock: string): string { + const contentMatch = /([\s\S]*?)<\/content>/.exec(xmlBlock); + return contentMatch ? contentMatch[1] : xmlBlock; +} + +/** + * Detect the XML tag name from a block string + */ +export function detectBlockTag(xmlBlock: string): string | null { + const match = xmlBlock.match(/^<(\w+)[\s>]/); + return match?.[1] ?? null; +} diff --git a/src/core/ContextManager.ts b/src/core/ContextManager.ts index 6ddf4b23..1605645d 100644 --- a/src/core/ContextManager.ts +++ b/src/core/ContextManager.ts @@ -3,6 +3,8 @@ import { ChainType } from "@/chainFactory"; import { processPrompt } from "@/commands/customCommandUtils"; import { LOADING_MESSAGES } from "@/constants"; import { PromptContextEngine } from "@/context/PromptContextEngine"; +import { compactXmlBlock, getL2RefetchInstruction } from "@/context/L2ContextCompactor"; +import { CONTEXT_BLOCK_TYPES, detectBlockTag } from "@/context/contextBlockRegistry"; import { PromptContextEnvelope, PromptLayerId, @@ -405,7 +407,10 @@ export class ContextManager { const segmentContent: string[] = []; for (const segment of l3Layer.segments || []) { if (segment.content) { - segmentContent.push(segment.content); + // Compact large L3 segments for L2 to reduce context size + // This extracts structure + preview with tool hints for re-fetching + const compacted = this.compactSegmentForL2(segment.content); + segmentContent.push(compacted); } } if (segmentContent.length > 0) { @@ -448,7 +453,16 @@ export class ContextManager { // - Lazy rebuild: process context.notes when first accessed after load } - return { l2Context: l2Parts.join("\n"), l2Paths }; + // Build the L2 context string + const l2Content = l2Parts.join("\n"); + + // Append re-fetch instruction if there's any L2 content with prior_context blocks + const hasCompactedContent = l2Content.includes(" and blocks and creates segments with path as ID + * Extracts , , and blocks and creates segments with path as ID */ private parseContextIntoSegments(contextXml: string, stable: boolean): PromptLayerSegment[] { if (!contextXml.trim()) { @@ -616,22 +630,23 @@ export class ContextManager { const segments: PromptLayerSegment[] = []; - // Match both and blocks + // Match , , and blocks + // prior_context has attributes: const noteContextRegex = /<(?:note_context|active_note)>[\s\S]*?<\/(?:note_context|active_note)>/g; + const priorContextRegex = /]*>[\s\S]*?<\/prior_context>/g; + let match: RegExpExecArray | null; + // Parse original note blocks while ((match = noteContextRegex.exec(contextXml)) !== null) { const block = match[0]; - - // Extract path from the block const pathMatch = /([^<]+)<\/path>/.exec(block); if (!pathMatch) continue; const notePath = pathMatch[1]; - segments.push({ - id: notePath, // Use note path as unique ID for comparison + id: notePath, content: block, stable, metadata: { @@ -641,6 +656,22 @@ export class ContextManager { }); } + // Parse compacted prior_context blocks (L2 from previous turns) + while ((match = priorContextRegex.exec(contextXml)) !== null) { + const block = match[0]; + const source = match[1]; // Extracted from source="..." attribute + + segments.push({ + id: source, + content: block, + stable, + metadata: { + source: "previous_turns_compacted", + notePath: source, + }, + }); + } + return segments; } @@ -710,6 +741,48 @@ export class ContextManager { getSelectedTextContexts() { return getSelectedTextContexts(); } + + /** + * Compact an L3 segment's content for inclusion in L2 (previous turn context). + * + * Large content is compacted by extracting structure (headings) and previews + * from each section. Small content and selected text are kept verbatim. + * A single re-fetch instruction is appended at the L2 level (not per-block). + * + * Handles multiple concatenated XML blocks of different types (e.g., web_tab_context + * mixed with youtube_video_context) by finding and compacting each block independently. + * + * @param content - The segment content (one or more XML blocks) + * @returns Compacted content for L2 + */ + private compactSegmentForL2(content: string): string { + // Build a regex that matches any known block type + const blockTags = CONTEXT_BLOCK_TYPES.map((bt) => bt.tag).join("|"); + const allBlocksRegex = new RegExp(`<(${blockTags})[^>]*>[\\s\\S]*?`, "g"); + + const blocks = content.match(allBlocksRegex); + + if (!blocks || blocks.length === 0) { + // No recognized XML blocks, return as-is + return content; + } + + if (blocks.length === 1) { + // Single block - detect type and compact + const blockType = detectBlockTag(blocks[0]); + if (!blockType) return content; + return compactXmlBlock(blocks[0], blockType); + } + + // Multiple blocks found - compact each one independently + const compactedBlocks = blocks.map((block) => { + const blockType = detectBlockTag(block); + if (!blockType) return block; + return compactXmlBlock(block, blockType); + }); + + return compactedBlocks.join("\n\n"); + } } /** diff --git a/src/hooks/use-streaming-chat-session.ts b/src/hooks/use-streaming-chat-session.ts index ce7a1da4..cfef4d8c 100644 --- a/src/hooks/use-streaming-chat-session.ts +++ b/src/hooks/use-streaming-chat-session.ts @@ -16,6 +16,7 @@ import type { BaseChatMemory } from "@langchain/classic/memory"; import type { CustomModel } from "@/aiParams"; import { createChatChain, createChatMemory } from "@/commands/customCommandChatEngine"; +import { compactAssistantOutput } from "@/context/ChatHistoryCompactor"; import { ThinkBlockStreamer } from "@/LLMProviders/chainRunner/utils/ThinkBlockStreamer"; import { ABORT_REASON } from "@/constants"; import { logError } from "@/logger"; @@ -333,9 +334,14 @@ export function useStreamingChatSession( committed = result; // Best-effort persistence - failure doesn't affect UI + // Compact the output to reduce memory bloat from tool results if (memory) { try { - await memory.saveContext({ input: prompt }, { output: result }); + const compactedResult = compactAssistantOutput(result); + await memory.saveContext( + { input: prompt }, + { output: typeof compactedResult === "string" ? compactedResult : result } + ); hasSavedContextOnceRef.current = true; } catch (error) { logError("Error saving chat context:", error); diff --git a/src/integration_tests/composerPrompt.test.ts b/src/integration_tests/composerPrompt.test.ts index b2eebfe1..15525cc2 100644 --- a/src/integration_tests/composerPrompt.test.ts +++ b/src/integration_tests/composerPrompt.test.ts @@ -32,7 +32,9 @@ jest.mock("../encryptionService", () => ({ getDecryptedKey: jest.fn().mockImplementation((key) => Promise.resolve(key)), })); -describe("Composer Instructions - Integration Tests", () => { +// TODO(@logancyang, @wenzhengjiang): Re-enable once composer is revamped. +// Currently skipped due to flaky 503 errors from Gemini API causing CI failures. +describe.skip("Composer Instructions - Integration Tests", () => { let genAI: GoogleGenerativeAI; let model: GenerativeModel;