@aksolab/recall
Version:
A memory management package for AI SDK memory functionality
229 lines (228 loc) • 8.2 kB
JavaScript
import { summarizeMessages } from "./ai/summarizer";
import { encoding_for_model } from "tiktoken";
import { AGENT_PROMPT } from "./ai/prompts";
export 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
constructor(provider, archiveProvider, openaiApiKey, memoryKey, threadId, maxContextSize, coreBlockTokenLimit) {
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;
}
}
async initialize(previousState) {
// Load core and archive memory
const state = await this.provider.initializeMemoryState(this.memoryKey, this.threadId, previousState);
console.log({ state });
this.coreMemory = state.coreMemory;
if (state?.chatHistory && state.chatHistory.length > 0) {
this.chatHistory = state.chatHistory;
}
else {
console.log("initializing new chat history", state?.chatHistory);
// 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();
}
}
async saveChatHistory() {
await this.provider.updateChatHistory({
memoryKey: this.memoryKey,
threadId: this.threadId,
messages: this.chatHistory
});
}
async saveCoreMemory() {
await this.provider.updateCoreMemory(this.memoryKey, this.coreMemory);
}
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 `${AGENT_PROMPT}\n\nCore Memory:\n${coreMemoryEntries}\n\n`;
}
async getChatHistory() {
return this.chatHistory;
}
async addUserMessage(message) {
this.chatHistory.push(message);
await this.checkChatHistorySize();
await this.saveChatHistory();
}
async getCoreMemory() {
return this.provider.getCoreMemory(this.memoryKey) || {};
}
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();
}
async searchArchiveMemory(query) {
const result = await this.archiveProvider.searchBySimilarity(query);
return result.map(r => r.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;
}
async updateArchiveMemory(id, payload) {
const updatedEntry = {
...payload,
timestamp: Date.now(),
};
const entry = await this.archiveProvider.updateEntry(id, updatedEntry);
return entry;
}
async removeArchivalMemory(id) {
const entry = await this.archiveProvider.getEntry(id);
if (!entry) {
return null;
}
await this.archiveProvider.deleteEntry(id);
return entry;
}
async addAIMessage(message) {
this.chatHistory.push(message);
await this.checkChatHistorySize();
await this.saveChatHistory();
}
async addAIMessages(messages) {
this.chatHistory.push(...messages);
await this.checkChatHistorySize();
await this.saveChatHistory();
}
/**
* Get the current context size in tokens
*/
get contextSize() {
return this.totalTokenCount();
}
/**
* Get the maximum allowed context size in tokens
*/
get maxContextSize() {
return this._maxContextSize;
}
/**
* Set the maximum allowed context size in tokens
*/
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);
});
}
getEncoder() {
if (!this.encoder) {
// Using gpt-4o tokenizer as it's compatible with most OpenAI models
this.encoder = encoding_for_model("gpt-4o");
}
return this.encoder;
}
countTokens(text) {
const encoder = this.getEncoder();
return encoder.encode(text).length;
}
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);
}
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 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();
}
}
// Clean up encoder when the instance is no longer needed
dispose() {
if (this.encoder) {
this.encoder.free();
this.encoder = null;
}
}
/**
* Get the core block token limit
*/
get coreMemoryBlockLimit() {
return this.coreBlockTokenLimit;
}
/**
* Set the core block token limit
*/
set coreMemoryBlockLimit(limit) {
this.coreBlockTokenLimit = limit;
}
}