From 98f444c11f9a6b07e08bc4fbf9968555f79bf027 Mon Sep 17 00:00:00 2001 From: Beatrix Date: Thu, 5 Dec 2024 11:36:16 -0800 Subject: [PATCH 1/2] feat(cody-chat-memory): Enhance chat memory management This PR introduces the following improvements to the CodyChatMemory module: - Transition to a static utility class with a Map-based storage approach - Implement timestamp-based memory organization and automatic trimming to maintain the 10 most recent items - Enhance the retrieve() method to return a formatted ContextItem with the memory content - Add support for REMOVE command to clear the chat memory - Introduce comprehensive unit tests to cover various chat session workflows ## Test plan - Run the unit tests for the CodyChatMemory module - Verify the chat memory functionality in the Cody extension, including adding, retrieving, and clearing memories ## Changelog - Improved chat memory management with timestamp-based organization and automatic trimming - Added support for REMOVE command to clear the chat memory - Enhanced unit test coverage for chat memory workflows --- .../src/chat/agentic/CodyChatMemory.test.ts | 157 ++++++++++++++++++ vscode/src/chat/agentic/CodyChatMemory.ts | 78 ++++++--- vscode/src/chat/agentic/CodyTool.ts | 6 +- 3 files changed, 213 insertions(+), 28 deletions(-) create mode 100644 vscode/src/chat/agentic/CodyChatMemory.test.ts diff --git a/vscode/src/chat/agentic/CodyChatMemory.test.ts b/vscode/src/chat/agentic/CodyChatMemory.test.ts new file mode 100644 index 000000000000..55b55b83fed0 --- /dev/null +++ b/vscode/src/chat/agentic/CodyChatMemory.test.ts @@ -0,0 +1,157 @@ +import { ContextItemSource } from '@sourcegraph/cody-shared' +import { afterAll, beforeEach, describe, expect, it, vi } from 'vitest' +import { URI } from 'vscode-uri' +import { localStorage } from '../../services/LocalStorageProvider' +import { CodyChatMemory } from './CodyChatMemory' + +// Mock localStorage +vi.mock('../../services/LocalStorageProvider', () => ({ + localStorage: { + getChatMemory: vi.fn(), + setChatMemory: vi.fn(), + }, +})) + +describe('CodyChatMemory Workflows', () => { + beforeEach(() => { + vi.clearAllMocks() + vi.useFakeTimers() + }) + + afterAll(() => { + vi.useRealTimers() + }) + + describe('Chat Session Workflows', () => { + interface TestScenario { + name: string + actions: Array<{ + type: 'initialize' | 'load' | 'retrieve' | 'unload' + input?: string + expectedContent?: string | null + expectedStorageCall?: boolean + }> + } + + const scenarios: TestScenario[] = [ + { + name: 'New user first chat session', + actions: [ + { + type: 'initialize', + expectedContent: null, + expectedStorageCall: true, + }, + { + type: 'load', + input: 'User prefers TypeScript', + // Update to match new timestamp + content format + expectedContent: + '## \\d{4}-\\d{2}-\\d{2}T\\d{2}:\\d{2}:\\d{2}.\\d{3}Z\\nUser prefers TypeScript', + expectedStorageCall: true, + }, + { + type: 'retrieve', + expectedContent: 'User prefers TypeScript', + }, + ], + }, + { + name: 'Multiple chat interactions in one session', + actions: [ + { + type: 'load', + input: 'User likes unit testing', + expectedContent: 'User likes unit testing', + }, + { + type: 'load', + input: 'User works on VS Code extensions', + expectedContent: 'User works on VS Code extensions', + }, + { + type: 'retrieve', + // Update regex to match new Map-based format with timestamps + expectedContent: + '## \\d{4}.*User likes unit testing.*## \\d{4}.*User works on VS Code extensions', + }, + ], + }, + { + name: 'Memory capacity management with timestamps', + actions: [ + ...Array.from({ length: 10 }, (_, i) => ({ + type: 'load' as const, + input: `Memory item ${i}`, + // Verify only last 8 items are kept + expectedContent: i >= 2 ? `Memory item ${i}` : null, + })), + { + type: 'retrieve', + // Verify chronological order is maintained + expectedContent: 'Memory item 2.*Memory item 9', + }, + ], + }, + // Add new test scenario for timestamp ordering + { + name: 'Timestamp ordering verification', + actions: [ + { + type: 'load', + input: 'First message', + }, + { + type: 'load', + input: 'Second message', + }, + { + type: 'retrieve', + // Verify messages appear in chronological order with timestamps + expectedContent: '.*First message.*Second message', + }, + ], + }, + ] + + for (const scenario of scenarios) { + it(scenario.name, () => { + for (const action of scenario.actions) { + switch (action.type) { + case 'load': + // Advance by 1 second to ensure unique timestamps + vi.advanceTimersByTime(1000) + CodyChatMemory.load(action.input!) + if (action.expectedStorageCall) { + expect(localStorage.setChatMemory).toHaveBeenCalled() + } + break + + case 'retrieve': { + const retrieved = CodyChatMemory.retrieve() + if (action.expectedContent === null) { + expect(retrieved).toBeUndefined() + } else { + if (action.expectedContent) { + expect(retrieved?.content).toMatch( + new RegExp(action.expectedContent, 's') + ) + } + expect(retrieved?.source).toBe(ContextItemSource.Agentic) + expect(retrieved?.uri).toEqual(URI.file('Cody Memory')) + } + break + } + case 'unload': { + const lastState = CodyChatMemory.reset() + if (action.expectedContent) { + expect(lastState?.content).toContain(action.expectedContent) + } + break + } + } + } + }) + } + }) +}) diff --git a/vscode/src/chat/agentic/CodyChatMemory.ts b/vscode/src/chat/agentic/CodyChatMemory.ts index 499780918f17..47648a0c807c 100644 --- a/vscode/src/chat/agentic/CodyChatMemory.ts +++ b/vscode/src/chat/agentic/CodyChatMemory.ts @@ -3,51 +3,63 @@ import { URI } from 'vscode-uri' import { localStorage } from '../../services/LocalStorageProvider' /** - * CodyChatMemory is a singleton class that manages short-term memory storage for chat conversations. - * It maintains a maximum of 8 most recent memory items in a static Store. - * We store the memory items in local storage to persist them across sessions. - * NOTE: The memory items set to a maximum of 8 to avoid overloading the local storage. + * CodyChatMemory is a static utility class that manages persistent chat memory storage. + * It maintains the 10 most recent memory items using a static Map and localStorage for persistence. * - * @remarks - * This class should never be instantiated directly. All operations should be performed - * through static methods. The only instance creation happens internally during initialization. + * The class handles: + * - Memory persistence across sessions via localStorage + * - Automatic trimming to maintain the 10 most recent items + * - Timestamp-based memory organization + * - Context retrieval for chat interactions * - * Key features: - * - Maintains a static Set of up to 8 chat memory items - * - Persists memory items to local storage - * - Provides memory retrieval as ContextItem for chat context + * Key Features: + * - Static interface - all operations performed through static methods + * - Automatic initialization on first use + * - Memory items formatted with timestamps + * - Integration with chat context system via ContextItem format * * Usage: - * - Call CodyChatMemory.initialize() once at startup - * - Use static methods load(), retrieve(), and unload() for memory operations + * - CodyChatMemory.initialize() - Called at startup to load existing memories + * - CodyChatMemory.load(memory) - Add new memory + * - CodyChatMemory.retrieve() - Get memories as chat context + * - CodyChatMemory.reset() - Clear and return last state */ export class CodyChatMemory { - private static readonly MAX_MEMORY_ITEMS = 8 - private static Store = new Set([]) + private static readonly MAX_MEMORY_ITEMS = 10 + private static Store = new Map() public static initialize(): void { if (CodyChatMemory.Store.size === 0) { const newMemory = new CodyChatMemory() - CodyChatMemory.Store = new Set(newMemory.getChatMemory()) + const memories = newMemory.getChatMemory() + CodyChatMemory.Store = new Map( + memories.map(memory => { + const [timestamp, ...content] = memory.split('\n') + return [timestamp.replace('## ', ''), content.join('\n')] + }) + ) } } public static load(memory: string): void { - CodyChatMemory.Store.add(memory) - // If store exceeds the max, remove oldest items - if (CodyChatMemory.Store.size > CodyChatMemory.MAX_MEMORY_ITEMS) { - const storeArray = Array.from(CodyChatMemory.Store) - CodyChatMemory.Store = new Set(storeArray.slice(-5)) - } - // TODO - persist to local file system - localStorage?.setChatMemory(Array.from(CodyChatMemory.Store)) + const timestamp = new Date().toISOString() + CodyChatMemory.Store.set(timestamp, memory) + // Convert existing entries to array for manipulation + const entries = Array.from(CodyChatMemory.Store.entries()) + // Keep only the most recent MAX_MEMORY_ITEMS entries & + // update stores with trimmed entries + const trimmedEntries = entries.slice(-CodyChatMemory.MAX_MEMORY_ITEMS) + CodyChatMemory.Store = new Map(trimmedEntries) + localStorage?.setChatMemory( + Array.from(trimmedEntries.entries()).map(([ts, mem]) => `## ${ts}\n${mem}`) + ) } public static retrieve(): ContextItem | undefined { return CodyChatMemory.Store.size > 0 ? { type: 'file', - content: '# Chat Memory\n' + Array.from(CodyChatMemory.Store).reverse().join('\n- '), + content: populateMemoryContent(CodyChatMemory.Store), uri: URI.file('Cody Memory'), source: ContextItemSource.Agentic, title: 'Cody Chat Memory', @@ -55,9 +67,9 @@ export class CodyChatMemory { : undefined } - public static unload(): ContextItem | undefined { + public static reset(): ContextItem | undefined { const stored = CodyChatMemory.retrieve() - CodyChatMemory.Store = new Set() + CodyChatMemory.Store.clear() return stored } @@ -65,3 +77,15 @@ export class CodyChatMemory { return localStorage?.getChatMemory() || [] } } + +export const CHAT_MEMORY_CONTEXT_TEMPLATE = `# Chat Memory +Here are the notes you made about the user (me) from previous chat: +{memoryItems}` + +function populateMemoryContent(memoryMap: Map): string { + const memories = Array.from(memoryMap.entries()) + .map(([timestamp, content]) => `\n## ${timestamp}\n${content}`) + .join('') + + return CHAT_MEMORY_CONTEXT_TEMPLATE.replace('{memoryItems}', memories) +} diff --git a/vscode/src/chat/agentic/CodyTool.ts b/vscode/src/chat/agentic/CodyTool.ts index ce3118bc3633..10034ec70441 100644 --- a/vscode/src/chat/agentic/CodyTool.ts +++ b/vscode/src/chat/agentic/CodyTool.ts @@ -283,12 +283,16 @@ class MemoryTool extends CodyTool { const newMemories = this.parse() for (const memory of newMemories) { if (memory === 'FORGET') { - CodyChatMemory.unload() + CodyChatMemory.reset() return } if (memory === 'GET') { return } + if (memory.startsWith('REMOVE')) { + CodyChatMemory.reset() + return + } CodyChatMemory.load(memory) logDebug('Cody Memory', 'added', { verbose: memory }) } From ae5023596497170752b9bd26844e68549d1fe8e4 Mon Sep 17 00:00:00 2001 From: Beatrix Date: Thu, 5 Dec 2024 12:36:29 -0800 Subject: [PATCH 2/2] left over --- vscode/src/chat/agentic/CodyChatMemory.ts | 2 +- vscode/src/chat/agentic/CodyTool.ts | 4 ---- 2 files changed, 1 insertion(+), 5 deletions(-) diff --git a/vscode/src/chat/agentic/CodyChatMemory.ts b/vscode/src/chat/agentic/CodyChatMemory.ts index 47648a0c807c..033ac4a556a3 100644 --- a/vscode/src/chat/agentic/CodyChatMemory.ts +++ b/vscode/src/chat/agentic/CodyChatMemory.ts @@ -79,7 +79,7 @@ export class CodyChatMemory { } export const CHAT_MEMORY_CONTEXT_TEMPLATE = `# Chat Memory -Here are the notes you made about the user (me) from previous chat: +Here are the notes you made about me (user) from previous chat: {memoryItems}` function populateMemoryContent(memoryMap: Map): string { diff --git a/vscode/src/chat/agentic/CodyTool.ts b/vscode/src/chat/agentic/CodyTool.ts index 10034ec70441..c7b1868325fe 100644 --- a/vscode/src/chat/agentic/CodyTool.ts +++ b/vscode/src/chat/agentic/CodyTool.ts @@ -289,10 +289,6 @@ class MemoryTool extends CodyTool { if (memory === 'GET') { return } - if (memory.startsWith('REMOVE')) { - CodyChatMemory.reset() - return - } CodyChatMemory.load(memory) logDebug('Cody Memory', 'added', { verbose: memory }) }