@aksolab/recall
Version:
A memory management package for AI SDK memory functionality
98 lines (97 loc) • 4.1 kB
JavaScript
export class InMemoryProvider {
// Separate storage maps for each component
storage = new Map();
chatHistoryStorage = new Map();
coreMemoryStorage = new Map();
archiveMemoryStorage = new Map();
async flush() {
this.storage.clear();
this.chatHistoryStorage.clear();
this.coreMemoryStorage.clear();
this.archiveMemoryStorage.clear();
}
async export(memoryKey, threadId = 'default') {
const coreMemory = this.coreMemoryStorage.get(this.getCoreMemoryKey(memoryKey)) || null;
const archiveMemory = this.archiveMemoryStorage.get(this.getArchiveMemoryKey(memoryKey)) || [];
const chatHistory = this.chatHistoryStorage.get(this.getChatHistoryKey(memoryKey, threadId)) || [];
return {
chatHistory,
coreMemory,
//archiveMemory
};
}
// Get complete memory state for a key
async getMemoryState(memoryKey) {
const chatHistory = this.chatHistoryStorage.get(this.getChatHistoryKey(memoryKey)) || [];
const coreMemory = this.coreMemoryStorage.get(this.getCoreMemoryKey(memoryKey)) || null;
//const archiveMemory = this.archiveMemoryStorage.get(this.getArchiveMemoryKey(memoryKey)) || [];
return {
chatHistory,
coreMemory,
//archiveMemory
};
}
// Initialize memory state if it doesn't exist
async initializeMemoryState(memoryKey, threadId, previousState) {
const existingState = await this.getMemoryState(memoryKey);
if (!existingState?.coreMemory) {
const initialState = {
chatHistory: previousState?.chatHistory || [],
coreMemory: previousState?.coreMemory || null,
//archiveMemory: previousState?.archiveMemory || []
};
console.log("initialState", initialState);
// Initialize component-specific storages
await this.updateChatHistory({
memoryKey,
threadId,
messages: initialState.chatHistory
});
await this.updateCoreMemory(memoryKey, initialState.coreMemory);
//await this.updateArchiveMemory(memoryKey, initialState.archiveMemory);
return initialState;
}
return existingState;
}
// Update chat history for a specific thread
async updateChatHistory({ memoryKey, threadId, messages }) {
this.chatHistoryStorage.set(this.getChatHistoryKey(memoryKey, threadId), messages);
}
async addChatHistoryMessage({ memoryKey, message, threadId }) {
const key = this.getChatHistoryKey(memoryKey, threadId);
const messages = this.chatHistoryStorage.get(key) || [];
messages.push(message);
this.chatHistoryStorage.set(key, messages);
}
async getCoreMemory(memoryKey) {
return this.coreMemoryStorage.get(this.getCoreMemoryKey(memoryKey)) || null;
}
// Update core memory
async updateCoreMemory(key, memory) {
this.coreMemoryStorage.set(this.getCoreMemoryKey(key), memory);
}
// Update archive memory
async updateArchiveMemory(key, entries) {
this.archiveMemoryStorage.set(this.getArchiveMemoryKey(key), entries);
}
// Delete memory state for a key
async deleteMemoryState(memoryKey) {
this.storage.delete(memoryKey);
this.chatHistoryStorage.delete(this.getChatHistoryKey(memoryKey));
this.coreMemoryStorage.delete(this.getCoreMemoryKey(memoryKey));
this.archiveMemoryStorage.delete(this.getArchiveMemoryKey(memoryKey));
}
// Get chat history key
getChatHistoryKey(memoryKey, threadId = 'default') {
return `${memoryKey}::chat-history::thread::${threadId}`;
}
// Get core memory key
getCoreMemoryKey(memoryKey) {
// Remove any thread-specific part before adding core memory suffix
return `${memoryKey}::core-memory`;
}
// Get archive memory key
getArchiveMemoryKey(memoryKey) {
return `${memoryKey}::archive-memory`;
}
}