Andreas/LibreChat
1
0
Fork 0
mirror of https://github.com/danny-avila/LibreChat.git synced 2025-12-21 10:50:14 +01:00
Code Issues Projects Releases Packages Wiki Activity Actions

🧠 feat: Enforce Token Limit for Memory Usage (#8401)

Browse source
This commit is contained in:
Samuel Path 2025-07-11 20:46:19 +02:00 • committed by GitHub
parent 2e1874e596
commit 8e869f2274
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 765 additions and 31 deletions
Download patch file Download diff file Expand all files Collapse all files

63
packages/api/src/agents/memory.ts
View file

@ -71,7 +71,7 @@ const getDefaultInstructions = (
/**
* Creates a memory tool instance with user context
*/
const createMemoryTool = ({
export const createMemoryTool = ({
userId,
setMemory,
validKeys,
@ -84,6 +84,9 @@ const createMemoryTool = ({
tokenLimit?: number;
totalTokens?: number;
}) => {
const remainingTokens = tokenLimit ? tokenLimit - totalTokens : Infinity;
const isOverflowing = tokenLimit ? remainingTokens <= 0 : false;
return tool(
async ({ key, value }) => {
try {
@ -93,24 +96,48 @@ const createMemoryTool = ({
', ',
)}`,
);
return `Invalid key "${key}". Must be one of: ${validKeys.join(', ')}`;
return [`Invalid key "${key}". Must be one of: ${validKeys.join(', ')}`, undefined];
}
const tokenCount = Tokenizer.getTokenCount(value, 'o200k_base');
if (tokenLimit && tokenCount > tokenLimit) {
logger.warn(
`Memory Agent failed to set memory: Value exceeds token limit. Value has ${tokenCount} tokens, but limit is ${tokenLimit}`,
);
return `Memory value too large: ${tokenCount} tokens exceeds limit of ${tokenLimit}`;
if (isOverflowing) {
const errorArtifact: Record<Tools.memory, MemoryArtifact> = {
[Tools.memory]: {
key: 'system',
type: 'error',
value: JSON.stringify({
errorType: 'already_exceeded',
tokenCount: Math.abs(remainingTokens),
totalTokens: totalTokens,
tokenLimit: tokenLimit!,
}),
tokenCount: totalTokens,
},
};
return [`Memory storage exceeded. Cannot save new memories.`, errorArtifact];
}
if (tokenLimit && totalTokens + tokenCount > tokenLimit) {
const remainingCapacity = tokenLimit - totalTokens;
logger.warn(
`Memory Agent failed to set memory: Would exceed total token limit. Current usage: ${totalTokens}, new memory: ${tokenCount} tokens, limit: ${tokenLimit}`,
);
return `Cannot add memory: would exceed token limit. Current usage: ${totalTokens}/${tokenLimit} tokens. This memory requires ${tokenCount} tokens, but only ${remainingCapacity} tokens available.`;
if (tokenLimit) {
const newTotalTokens = totalTokens + tokenCount;
const newRemainingTokens = tokenLimit - newTotalTokens;
if (newRemainingTokens < 0) {
const errorArtifact: Record<Tools.memory, MemoryArtifact> = {
[Tools.memory]: {
key: 'system',
type: 'error',
value: JSON.stringify({
errorType: 'would_exceed',
tokenCount: Math.abs(newRemainingTokens),
totalTokens: newTotalTokens,
tokenLimit,
}),
tokenCount: totalTokens,
},
};
return [`Memory storage would exceed limit. Cannot save this memory.`, errorArtifact];
}
}
const artifact: Record<Tools.memory, MemoryArtifact> = {
@ -177,7 +204,7 @@ const createDeleteMemoryTool = ({
', ',
)}`,
);
return `Invalid key "${key}". Must be one of: ${validKeys.join(', ')}`;
return [`Invalid key "${key}". Must be one of: ${validKeys.join(', ')}`, undefined];
}
const artifact: Record<Tools.memory, MemoryArtifact> = {
@ -269,7 +296,13 @@ export async function processMemory({
llmConfig?: Partial<LLMConfig>;
}): Promise<(TAttachment | null)[] | undefined> {
try {
const memoryTool = createMemoryTool({ userId, tokenLimit, setMemory, validKeys, totalTokens });
const memoryTool = createMemoryTool({
userId,
tokenLimit,
setMemory,
validKeys,
totalTokens,
});
const deleteMemoryTool = createDeleteMemoryTool({
userId,
validKeys,