Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix(deep-cody): update chat memory management #6264

Draft
wants to merge 2 commits into
base: main
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
157 changes: 157 additions & 0 deletions vscode/src/chat/agentic/CodyChatMemory.test.ts
Original file line number Diff line number Diff line change
@@ -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
}
}
}
})
}
})
})
78 changes: 51 additions & 27 deletions vscode/src/chat/agentic/CodyChatMemory.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,65 +3,89 @@ 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<string>([])
private static readonly MAX_MEMORY_ITEMS = 10
private static Store = new Map<string, string>()

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',
}
: undefined
}

public static unload(): ContextItem | undefined {
public static reset(): ContextItem | undefined {
const stored = CodyChatMemory.retrieve()
CodyChatMemory.Store = new Set<string>()
CodyChatMemory.Store.clear()
return stored
}

private getChatMemory(): string[] {
return localStorage?.getChatMemory() || []
}
}

export const CHAT_MEMORY_CONTEXT_TEMPLATE = `# Chat Memory
Copy link
Collaborator

@PriNova PriNova Dec 5, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

What do you think about keeping the prompt more generally? Similar to the following:
'Here are important memos/notes from past conversations' (Up to you about the phrasing)

This would also make it possible for users to store notes that are not directly related to them, but correspond to specifications/requirements and may not confuse Cody.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@PriNova Are you referring to the about me (user) part of the prompt?

Something like:

Here are the notes from our past conversations

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@PriNova Are you referring to the about me (user) part of the prompt?

Something like:

Here are the notes from our past conversations

Yes. I missed the correct line comment. It is meant for the prompt.

Here are the notes you made about me (user) from previous chat:
{memoryItems}`

function populateMemoryContent(memoryMap: Map<string, string>): string {
const memories = Array.from(memoryMap.entries())
.map(([timestamp, content]) => `\n## ${timestamp}\n${content}`)
.join('')

return CHAT_MEMORY_CONTEXT_TEMPLATE.replace('{memoryItems}', memories)
}
2 changes: 1 addition & 1 deletion vscode/src/chat/agentic/CodyTool.ts
Original file line number Diff line number Diff line change
Expand Up @@ -283,7 +283,7 @@ class MemoryTool extends CodyTool {
const newMemories = this.parse()
for (const memory of newMemories) {
if (memory === 'FORGET') {
CodyChatMemory.unload()
CodyChatMemory.reset()
return
}
if (memory === 'GET') {
Expand Down
Loading