diff --git a/api/server/controllers/agents/client.js b/api/server/controllers/agents/client.js index abef2061a7..0e3b7f91bb 100644 --- a/api/server/controllers/agents/client.js +++ b/api/server/controllers/agents/client.js @@ -35,6 +35,7 @@ const { injectSkillPrimes, isSkillPrimeMessage, collectFileIds, + processTextWithTokenLimit, buildAgentScopedContext, buildSkillPrimeContentParts, buildInitialToolSessions, @@ -58,6 +59,7 @@ const { isAgentsEndpoint, isEphemeralAgentId, removeNullishValues, + DEFAULT_MEMORY_MAX_INPUT_TOKENS, } = require('librechat-data-provider'); const { filterFilesByAgentAccess } = require('~/server/services/Files/permissions'); const { encodeAndFormat } = require('~/server/services/Files/images/encode'); @@ -70,6 +72,8 @@ const db = require('~/models'); const loadAgent = (params) => loadAgentFn(params, { getAgent: db.getAgent, getMCPServerTools }); +const MEMORY_INPUT_CHARS_PER_TOKEN = 8; + class AgentClient extends BaseClient { constructor(options = {}) { super(null, options); @@ -760,7 +764,44 @@ class AgentClient extends BaseClient { const filteredMessages = messagesToProcess.map((msg) => this.filterImageUrls(msg)); const bufferString = getBufferString(filteredMessages); - const bufferMessage = new HumanMessage(`# Current Chat:\n\n${bufferString}`); + const configuredMaxInputTokens = Number.isFinite(memoryConfig?.maxInputTokens) + ? Math.floor(memoryConfig.maxInputTokens) + : undefined; + const maxInputTokens = + configuredMaxInputTokens != null && configuredMaxInputTokens > 0 + ? configuredMaxInputTokens + : DEFAULT_MEMORY_MAX_INPUT_TOKENS; + const maxInputChars = maxInputTokens * MEMORY_INPUT_CHARS_PER_TOKEN; + const isCharTruncated = bufferString.length > maxInputChars; + const memoryInput = `# Current Chat:\n\n${ + isCharTruncated + ? `[Earlier chat content omitted due to memory input limit]\n\n${bufferString.slice( + -maxInputChars, + )}` + : bufferString + }`; + const { + text: limitedMemoryInput, + tokenCount, + wasTruncated, + } = await processTextWithTokenLimit({ + text: memoryInput, + tokenLimit: maxInputTokens, + tokenCountFn: (text) => countTokens(text), + preserve: 'end', + }); + if (isCharTruncated || wasTruncated) { + logger.warn('[MemoryAgent] Memory input truncated before processing', { + tokenCount, + messageId: this.responseMessageId, + conversationId: this.conversationId, + maxInputTokens, + wasTruncated, + maxInputChars, + originalLength: bufferString.length, + }); + } + const bufferMessage = new HumanMessage(limitedMemoryInput); return await this.processMemory([bufferMessage]); } catch (error) { logger.error('Memory Agent failed to process memory', error); diff --git a/api/server/controllers/agents/client.test.js b/api/server/controllers/agents/client.test.js index 193c8d8c06..9fe4623c9c 100644 --- a/api/server/controllers/agents/client.test.js +++ b/api/server/controllers/agents/client.test.js @@ -2010,6 +2010,25 @@ describe('AgentClient - titleConvo', () => { expect(processedMessage.content).not.toContain('Response 1'); }); + it('should cap memory input tokens and preserve recent content', async () => { + const { HumanMessage, AIMessage } = require('@librechat/agents/langchain/messages'); + mockReq.config.memory.maxInputTokens = 12; + const messages = [ + new HumanMessage(`OLDER_CONTENT ${'a'.repeat(600)}`), + new AIMessage('Intermediate response'), + new HumanMessage('Please remember LATEST_MEMORY_MARKER'), + ]; + + await client.runMemory(messages); + + expect(mockProcessMemory).toHaveBeenCalledTimes(1); + const processedMessage = mockProcessMemory.mock.calls[0][0][0]; + + expect(processedMessage.content).toContain('LATEST_MEMORY_MARKER'); + expect(processedMessage.content).not.toContain('OLDER_CONTENT'); + expect(Math.ceil(processedMessage.content.length / 4)).toBeLessThanOrEqual(12); + }); + it('should return early if processMemory is not set', async () => { const { HumanMessage } = require('@librechat/agents/langchain/messages'); client.processMemory = null; diff --git a/librechat.example.yaml b/librechat.example.yaml index 7fb2081e9e..198cf304e4 100644 --- a/librechat.example.yaml +++ b/librechat.example.yaml @@ -758,8 +758,10 @@ endpoints: # disabled: false # # (optional) Restrict memory keys to specific values to limit memory storage and improve consistency # validKeys: ["preferences", "work_info", "personal_info", "skills", "interests", "context"] -# # (optional) Maximum token limit for memory storage (not yet implemented for token counting) +# # (optional) Maximum token limit for stored memory values # tokenLimit: 10000 +# # (optional) Maximum tokens from recent chat sent to the memory agent before truncation +# maxInputTokens: 12000 # # (optional) Enable personalization features (defaults to true if memory is configured) # # When false, users will not see the Personalization tab in settings # personalize: true diff --git a/packages/api/src/agents/__tests__/memory.test.ts b/packages/api/src/agents/__tests__/memory.test.ts index d23a07aa7d..2555102057 100644 --- a/packages/api/src/agents/__tests__/memory.test.ts +++ b/packages/api/src/agents/__tests__/memory.test.ts @@ -410,6 +410,7 @@ describe('processMemory - GPT-5+ handling', () => { llmConfig: expect.objectContaining({ model: 'gpt-4.1-mini', temperature: 0.4, // Default temperature should remain + maxRetries: 0, }), }), }), diff --git a/packages/api/src/agents/memory.ts b/packages/api/src/agents/memory.ts index d47ab2bf6e..4006ea1184 100644 --- a/packages/api/src/agents/memory.ts +++ b/packages/api/src/agents/memory.ts @@ -355,6 +355,7 @@ ${memory ?? 'No existing memories'}`; const finalLLMConfig = { ...defaultLLMConfig, ...normalizeMemoryLLMConfig(llmConfig), + maxRetries: 0, /** * Ensure streaming is always disabled for memory processing */ diff --git a/packages/api/src/memory/config.spec.ts b/packages/api/src/memory/config.spec.ts index 4a836cefa0..3cadc07ed6 100644 --- a/packages/api/src/memory/config.spec.ts +++ b/packages/api/src/memory/config.spec.ts @@ -1,4 +1,5 @@ import { logger } from '@librechat/data-schemas'; +import { DEFAULT_MEMORY_MAX_INPUT_TOKENS } from 'librechat-data-provider'; import type { TCustomConfig } from 'librechat-data-provider'; @@ -14,6 +15,7 @@ describe('memory config', () => { expect(isMemoryEnabled(loaded)).toBe(true); expect(isMemoryAgentEnabled(loaded)).toBe(false); + expect(loaded?.maxInputTokens).toBe(DEFAULT_MEMORY_MAX_INPUT_TOKENS); }); it('requires explicit memory agent enablement before enabling the automatic agent flow', () => { diff --git a/packages/api/src/utils/text.spec.ts b/packages/api/src/utils/text.spec.ts index 30185f9da7..cbafb25af7 100644 --- a/packages/api/src/utils/text.spec.ts +++ b/packages/api/src/utils/text.spec.ts @@ -222,6 +222,24 @@ describe('processTextWithTokenLimit', () => { expect(syncResult.wasTruncated).toBe(asyncResult.wasTruncated); expect(syncResult.text.length).toBe(asyncResult.text.length); }); + + it('should preserve the end of text when requested', async () => { + const { tokenCountFn } = createMockTokenCounter(); + const text = `${'a'.repeat(200)}LATEST_MEMORY_REQUEST`; + const tokenLimit = 10; + + const result = await processTextWithTokenLimit({ + text, + tokenLimit, + tokenCountFn, + preserve: 'end', + }); + + expect(result.wasTruncated).toBe(true); + expect(result.tokenCount).toBeLessThanOrEqual(tokenLimit); + expect(result.text).toContain('LATEST_MEMORY_REQUEST'); + expect(result.text).not.toContain('aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa'); + }); }); describe('when text is under the token limit', () => { diff --git a/packages/api/src/utils/text.ts b/packages/api/src/utils/text.ts index 3099c2bbc4..5273670e10 100644 --- a/packages/api/src/utils/text.ts +++ b/packages/api/src/utils/text.ts @@ -39,10 +39,12 @@ export async function processTextWithTokenLimit({ text, tokenLimit, tokenCountFn, + preserve = 'start', }: { text: string; tokenLimit: number; tokenCountFn: TokenCountFn; + preserve?: 'start' | 'end'; }): Promise<{ text: string; tokenCount: number; wasTruncated: boolean }> { const originalTokenCount = await tokenCountFn(text); @@ -61,7 +63,14 @@ export async function processTextWithTokenLimit({ const ratio = tokenLimit / originalTokenCount; let charPosition = Math.floor(text.length * ratio * TRUNCATION_SAFETY_BUFFER); - let truncatedText = text.substring(0, charPosition); + const sliceText = (position: number) => { + if (position <= 0) { + return ''; + } + return preserve === 'end' ? text.slice(-position) : text.substring(0, position); + }; + + let truncatedText = sliceText(charPosition); let tokenCount = await tokenCountFn(truncatedText); const maxIterations = 5; @@ -70,7 +79,7 @@ export async function processTextWithTokenLimit({ while (tokenCount > tokenLimit && iterations < maxIterations && charPosition > 0) { const overageRatio = tokenLimit / tokenCount; charPosition = Math.floor(charPosition * overageRatio * TRUNCATION_SAFETY_BUFFER); - truncatedText = text.substring(0, charPosition); + truncatedText = sliceText(charPosition); tokenCount = await tokenCountFn(truncatedText); iterations++; } diff --git a/packages/data-provider/src/config.ts b/packages/data-provider/src/config.ts index 69647a1201..34fa03c642 100644 --- a/packages/data-provider/src/config.ts +++ b/packages/data-provider/src/config.ts @@ -1354,11 +1354,14 @@ export const transactionsSchema = z.object({ enabled: z.boolean().optional().default(true), }); +export const DEFAULT_MEMORY_MAX_INPUT_TOKENS = 12000; + export const memorySchema = z.object({ disabled: z.boolean().optional(), validKeys: z.array(z.string()).optional(), tokenLimit: z.number().optional(), charLimit: z.number().optional().default(10000), + maxInputTokens: z.number().int().positive().optional().default(DEFAULT_MEMORY_MAX_INPUT_TOKENS), personalize: z.boolean().default(true), messageWindowSize: z.number().optional().default(5), agent: z diff --git a/packages/data-schemas/src/app/memory.ts b/packages/data-schemas/src/app/memory.ts index a3b89176db..ee298db939 100644 --- a/packages/data-schemas/src/app/memory.ts +++ b/packages/data-schemas/src/app/memory.ts @@ -1,4 +1,4 @@ -import { memorySchema } from 'librechat-data-provider'; +import { DEFAULT_MEMORY_MAX_INPUT_TOKENS, memorySchema } from 'librechat-data-provider'; import type { TCustomConfig, TMemoryConfig } from 'librechat-data-provider'; @@ -23,8 +23,11 @@ export function loadMemoryConfig(config: TCustomConfig['memory']): TMemoryConfig } const charLimit = memorySchema.shape.charLimit.safeParse(config.charLimit).data ?? 10000; + const maxInputTokens = + memorySchema.shape.maxInputTokens.safeParse(config.maxInputTokens).data ?? + DEFAULT_MEMORY_MAX_INPUT_TOKENS; - return { ...config, charLimit }; + return { ...config, charLimit, maxInputTokens }; } export function isMemoryEnabled(config: TMemoryConfig | undefined): boolean {