@aksolab/recall
Version:
A memory management package for AI SDK memory functionality
360 lines (359 loc) • 12.7 kB
JavaScript
"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
exports.MemoryManager = void 0;
const summarizer_1 = require("./ai/summarizer");
const tiktoken_1 = require("tiktoken");
const prompts_1 = require("./ai/prompts");
class MemoryManager {
provider;
archiveProvider;
openaiApiKey;
memoryKey;
threadId;
chatHistory = [];
coreMemory = null;
encoder = null;
chatTokenLimit = 10000; // default token limit
_maxContextSize = 20000; // default max context size
coreBlockTokenLimit = 2000; // default core block token limit
/**
* Creates a new instance of the MemoryManager.
* @param provider - Storage provider for chat history and core memory
* @param archiveProvider - Provider for archival memory and RAG capabilities
* @param openaiApiKey - OpenAI API key for embeddings and summarization
* @param memoryKey - Unique identifier for this memory instance
* @param threadId - Conversation thread identifier
* @param maxContextSize - Optional maximum context size in tokens
* @param coreBlockTokenLimit - Optional token limit for core memory blocks
*/
constructor(provider, archiveProvider, openaiApiKey, memoryKey, threadId, maxContextSize, coreBlockTokenLimit, chatTokenLimit) {
this.provider = provider;
this.archiveProvider = archiveProvider;
this.openaiApiKey = openaiApiKey;
this.memoryKey = memoryKey;
this.threadId = threadId;
if (maxContextSize) {
this._maxContextSize = maxContextSize;
}
if (coreBlockTokenLimit) {
this.coreBlockTokenLimit = coreBlockTokenLimit;
}
if (chatTokenLimit) {
this.chatTokenLimit = chatTokenLimit;
}
}
/**
* Initializes the memory manager with optional previous state.
* @param previousState - Optional previous memory state to restore
*/
async initialize(previousState) {
// Load core and archive memory
const state = await this.provider.initializeMemoryState(this.memoryKey, this.threadId, previousState);
this.coreMemory = state.coreMemory;
if (state?.chatHistory && state.chatHistory.length > 0) {
this.chatHistory = state.chatHistory;
}
else {
// Initialize new chat history with system message
this.chatHistory = [{ role: 'system', content: this.coreMemoryToString() }];
await this.saveChatHistory();
}
// Ensure core memory changes are persisted
if (this.coreMemory) {
await this.saveCoreMemory();
}
}
/**
* Saves the current chat history to the storage provider.
* @private
*/
async saveChatHistory() {
await this.provider.updateChatHistory({
memoryKey: this.memoryKey,
threadId: this.threadId,
messages: this.chatHistory
});
}
/**
* Saves the current core memory to the storage provider.
* @private
*/
async saveCoreMemory() {
await this.provider.updateCoreMemory(this.memoryKey, this.coreMemory);
}
/**
* Converts core memory blocks to a string representation.
* @private
* @returns Formatted string of core memory blocks
*/
coreMemoryToString() {
const coreMemoryEntries = this.coreMemory
? Object.entries(this.coreMemory)
.map(([key, entry]) => {
return `Name: ${key}\nDescription: ${entry.description}\nContent: ${entry.content}`;
})
.join("\n---\n")
: 'No core memory available';
return `${prompts_1.AGENT_PROMPT}\n\nCore Memory:\n${coreMemoryEntries}\n\n`;
}
/**
* Retrieves the complete chat history.
* @returns Promise resolving to an array of chat messages
*/
async getChatHistory() {
return this.chatHistory;
}
/**
* Adds one or more messages to the chat history.
* @param messages - A single message or array of messages to add
* @example
* // Add a single message
* await memoryManager.addMessages({ role: 'user', content: 'Hello!' });
*
* // Add multiple messages
* await memoryManager.addMessages([
* { role: 'assistant', content: 'Hi!' },
* { role: 'user', content: 'How are you?' }
* ]);
*/
async addMessages(messages) {
const messageArray = Array.isArray(messages) ? messages : [messages];
this.chatHistory.push(...messageArray);
await this.checkChatHistorySize();
await this.saveChatHistory();
}
/**
* @deprecated Use {@link addMessages} instead. This method will be removed in the next major version.
* @example
* // Instead of:
* await memoryManager.addUserMessage(message);
*
* // Use:
* await memoryManager.addMessages(message);
*/
async addUserMessage(message) {
return this.addMessages(message);
}
/**
* @deprecated Use {@link addMessages} instead. This method will be removed in the next major version.
* @example
* // Instead of:
* await memoryManager.addAIMessage(message);
*
* // Use:
* await memoryManager.addMessages(message);
*/
async addAIMessage(message) {
return this.addMessages(message);
}
/**
* @deprecated Use {@link addMessages} instead. This method will be removed in the next major version.
* @example
* // Instead of:
* await memoryManager.addAIMessages(messages);
*
* // Use:
* await memoryManager.addMessages(messages);
*/
async addAIMessages(messages) {
return this.addMessages(messages);
}
/**
* Retrieves the current state of core memory.
* @returns Promise resolving to a record of core memory blocks and their entries
*/
async getCoreMemory() {
return this.provider.getCoreMemory(this.memoryKey) || {};
}
/**
* Updates a core memory block with new content.
* @param block - The core memory block to update
* @param content - New content for the block
* @param description - Optional description for the block
* @throws {Error} If content exceeds token limit
*/
async updateCoreMemory(block, content, description) {
if (!this.coreMemory) {
this.coreMemory = await this.getCoreMemory() || {};
}
// Check token count for the content
const contentTokens = this.countTokens(content);
if (contentTokens > this.coreBlockTokenLimit) {
throw new Error(`Core memory block content exceeds token limit of ${this.coreBlockTokenLimit} tokens. Current: ${contentTokens} tokens.`);
}
this.coreMemory[block] = { content, description: description || this.coreMemory[block]?.description || '' };
if (this.chatHistory[0]?.role === 'system') {
this.chatHistory[0].content = this.coreMemoryToString();
await this.saveChatHistory();
}
await this.saveCoreMemory();
}
/**
* Searches the archive memory using semantic similarity.
* @param query - Search query text
* @returns Promise resolving to array of matching archive entries
*/
async searchArchiveMemory(query) {
const result = await this.archiveProvider.searchBySimilarity(query);
return result.map(r => r.entry);
}
/**
* Adds a new entry to the archive memory.
* @param payload - Archive entry data to add
* @returns Promise resolving to the created archive entry
*/
async addToArchiveMemory(payload) {
const timestamp = Date.now();
const id = payload.id || `archival_memory_${timestamp}`;
const newEntry = {
id,
content: payload.content,
name: payload.name,
timestamp,
};
const entry = await this.archiveProvider.addEntry(newEntry);
return entry;
}
/**
* Updates an existing archive memory entry.
* @param id - ID of the entry to update
* @param payload - Updated archive entry data
* @returns Promise resolving to the updated archive entry
*/
async updateArchiveMemory(id, payload) {
const updatedEntry = {
...payload,
timestamp: Date.now(),
};
const entry = await this.archiveProvider.updateEntry(id, updatedEntry);
return entry;
}
/**
* Removes an entry from archive memory.
* @param id - ID of the entry to remove
* @returns Promise resolving to the removed entry, or null if not found
*/
async removeArchivalMemory(id) {
const entry = await this.archiveProvider.getEntry(id);
if (!entry) {
return null;
}
await this.archiveProvider.deleteEntry(id);
return entry;
}
/**
* Gets the current context size in tokens.
*/
get contextSize() {
return this.totalTokenCount();
}
/**
* Gets the maximum allowed context size in tokens.
*/
get maxContextSize() {
return this._maxContextSize;
}
/**
* Sets the maximum allowed context size in tokens.
* Triggers chat history summarization if needed.
*/
set maxContextSize(size) {
this._maxContextSize = size;
// Check if we need to summarize due to new limit
this.checkChatHistorySize().catch(error => {
console.error('Error checking chat history size after maxContextSize update:', error);
});
}
/**
* Gets or initializes the token encoder.
* @private
* @returns Token encoder instance
*/
getEncoder() {
if (!this.encoder) {
// Using gpt-4o tokenizer as it's compatible with most OpenAI models
this.encoder = (0, tiktoken_1.encoding_for_model)("gpt-4o");
}
return this.encoder;
}
/**
* Counts the number of tokens in a text string.
* @private
* @param text - Text to count tokens for
* @returns Number of tokens
*/
countTokens(text) {
const encoder = this.getEncoder();
return encoder.encode(text).length;
}
/**
* Calculates the total token count of the chat history.
* @private
* @returns Total number of tokens
*/
totalTokenCount() {
return this.chatHistory.reduce((total, message) => {
if (typeof message.content === 'string') {
return total + this.countTokens(message.content);
}
if (Array.isArray(message.content)) {
return total + message.content.reduce((acc, item) => {
if (item.type === 'text') {
acc += this.countTokens(item.text);
}
else if (['tool-call', 'tool-result'].includes(item.type)) {
acc += this.countTokens(JSON.stringify(item));
}
return acc;
}, 0);
}
return total;
}, 0);
}
/**
* Checks and manages chat history size, summarizing if needed.
* @private
*/
async checkChatHistorySize() {
while (this.totalTokenCount() > this.chatTokenLimit) {
// Keep the system message (index 0) and last message
const messagesToSummarize = this.chatHistory.slice(1, -1);
if (messagesToSummarize.length === 0)
break;
const summary = await (0, summarizer_1.summarizeMessages)(messagesToSummarize, this.openaiApiKey);
const firstMessage = this.chatHistory[0];
const lastMessage = this.chatHistory[this.chatHistory.length - 1];
if (!firstMessage || !lastMessage) {
return;
}
this.chatHistory = [
firstMessage,
{ role: 'system', content: `Previous conversation summary: ${summary}` },
lastMessage
];
await this.saveChatHistory();
}
}
/**
* Cleans up resources when the instance is no longer needed.
*/
dispose() {
if (this.encoder) {
this.encoder.free();
this.encoder = null;
}
}
/**
* Gets the token limit for core memory blocks.
*/
get coreMemoryBlockLimit() {
return this.coreBlockTokenLimit;
}
/**
* Sets the token limit for core memory blocks.
*/
set coreMemoryBlockLimit(limit) {
this.coreBlockTokenLimit = limit;
}
}
exports.MemoryManager = MemoryManager;