diff --git a/api/server/controllers/agents/client.js b/api/server/controllers/agents/client.js index c18e0b8b1e..abef2061a7 100644 --- a/api/server/controllers/agents/client.js +++ b/api/server/controllers/agents/client.js @@ -30,6 +30,7 @@ const { createMultiAgentMapper, filterMalformedContentParts, countFormattedMessageTokens, + prependFileContext, hydrateMissingIndexTokenCounts, injectSkillPrimes, isSkillPrimeMessage, @@ -135,6 +136,8 @@ class AgentClient extends BaseClient { this.usage; /** @type {Record} */ this.indexTokenCountMap = {}; + /** @type {Array> | null} */ + this.memoryPayload = null; /** @type {(messages: BaseMessage[]) => Promise} */ this.processMemory; } @@ -315,39 +318,51 @@ class AgentClient extends BaseClient { } /** @type {Record} */ - const canonicalTokenCountMap = {}; + const indexTokenCountMap = {}; /** @type {Record} */ const tokenCountMap = {}; + const memoryPayload = []; + let hasFileContext = false; let promptTokenTotal = 0; + const encoding = this.getEncoding(); const formattedMessages = orderedMessages.map((message, i) => { const formattedMessage = formatMessage({ message, userName: this.options?.name, assistantName: this.options?.modelLabel, }); + const memoryFormattedMessage = formatMessage({ + message, + userName: this.options?.name, + assistantName: this.options?.modelLabel, + }); - /** For non-latest messages, prepend file context directly to message content */ - if (message.fileContext && i !== orderedMessages.length - 1) { - if (typeof formattedMessage.content === 'string') { - formattedMessage.content = message.fileContext + '\n' + formattedMessage.content; - } else { - const textPart = formattedMessage.content.find((part) => part.type === 'text'); - textPart - ? (textPart.text = message.fileContext + '\n' + textPart.text) - : formattedMessage.content.unshift({ type: 'text', text: message.fileContext }); - } + /** + * Bind file context to the message it belongs to. Historical attachments + * are resent inline, so the current turn's text attachment must be inline + * too instead of living only in the dynamic system tail. + */ + if (message.fileContext) { + hasFileContext = true; + prependFileContext(formattedMessage, message.fileContext); } - const dbTokenCount = orderedMessages[i].tokenCount; - const needsTokenCount = !dbTokenCount || message.fileContext; + memoryPayload.push(memoryFormattedMessage); - if (needsTokenCount || (this.isVisionModel && (message.image_urls || message.files))) { - orderedMessages[i].tokenCount = countFormattedMessageTokens( - formattedMessage, - this.getEncoding(), - ); + const dbTokenCount = Number(orderedMessages[i].tokenCount); + const hasDbTokenCount = Number.isFinite(dbTokenCount) && dbTokenCount > 0; + const needsCanonicalTokenCount = + !hasDbTokenCount || (this.isVisionModel && (message.image_urls || message.files)); + + let canonicalTokenCount = hasDbTokenCount ? dbTokenCount : 0; + if (needsCanonicalTokenCount) { + canonicalTokenCount = countFormattedMessageTokens(memoryFormattedMessage, encoding); } + const promptMessageTokenCount = message.fileContext + ? countFormattedMessageTokens(formattedMessage, encoding) + : canonicalTokenCount; + /* If message has files, calculate image token cost */ if (this.message_file_map && this.message_file_map[message.messageId]) { const attachments = this.message_file_map[message.messageId]; @@ -362,13 +377,19 @@ class AgentClient extends BaseClient { } } - const tokenCount = Number(orderedMessages[i].tokenCount); - const normalizedTokenCount = Number.isFinite(tokenCount) && tokenCount > 0 ? tokenCount : 0; - canonicalTokenCountMap[i] = normalizedTokenCount; - promptTokenTotal += normalizedTokenCount; + const normalizedCanonicalTokenCount = + Number.isFinite(canonicalTokenCount) && canonicalTokenCount > 0 ? canonicalTokenCount : 0; + const normalizedPromptTokenCount = + Number.isFinite(promptMessageTokenCount) && promptMessageTokenCount > 0 + ? promptMessageTokenCount + : 0; + + orderedMessages[i].tokenCount = normalizedCanonicalTokenCount; + indexTokenCountMap[i] = normalizedPromptTokenCount; + promptTokenTotal += normalizedPromptTokenCount; if (message.messageId) { - tokenCountMap[message.messageId] = normalizedTokenCount; + tokenCountMap[message.messageId] = normalizedCanonicalTokenCount; } if (isEnabled(process.env.AGENT_DEBUG_LOGGING)) { @@ -377,9 +398,10 @@ class AgentClient extends BaseClient { Array.isArray(message.content) && message.content.some((p) => p && p.type === 'summary'); const suffix = hasSummary ? '[S]' : ''; const id = (message.messageId ?? message.id ?? '').slice(-8); - const recalced = needsTokenCount ? orderedMessages[i].tokenCount : null; + const recalced = needsCanonicalTokenCount ? normalizedCanonicalTokenCount : null; + const promptRecalced = message.fileContext ? normalizedPromptTokenCount : null; logger.debug( - `[AgentClient] msg[${i}] ${role}${suffix} id=…${id} db=${dbTokenCount} needsRecount=${needsTokenCount} recalced=${recalced} tokens=${normalizedTokenCount}`, + `[AgentClient] msg[${i}] ${role}${suffix} id=…${id} db=${dbTokenCount} needsRecount=${needsCanonicalTokenCount} recalced=${recalced} promptRecalced=${promptRecalced} tokens=${normalizedPromptTokenCount}`, ); } @@ -387,22 +409,18 @@ class AgentClient extends BaseClient { }); payload = formattedMessages; + this.memoryPayload = hasFileContext ? memoryPayload : null; messages = orderedMessages; promptTokens = promptTokenTotal; /** * Build shared run context - applies to ALL agents in the run. - * This includes file context from the latest message and augmented prompt (RAG). + * Request attachment file context is already bound inline to the latest + * user message above; only side-channel context belongs here. * Memory context is handled separately and applied per-agent based on config. */ const sharedRunContextParts = []; - /** File context from the latest message (attachments) */ - const latestMessage = orderedMessages[orderedMessages.length - 1]; - if (latestMessage?.fileContext) { - sharedRunContextParts.push(latestMessage.fileContext); - } - /** Augmented prompt from RAG/context handlers */ if (this.contextHandlers) { this.augmentedPrompt = await this.contextHandlers.createContext(); @@ -428,8 +446,8 @@ class AgentClient extends BaseClient { tokenCountFn: (text) => countTokens(text), }); - /** Preserve canonical pre-format token counts for all history entering graph formatting */ - this.indexTokenCountMap = canonicalTokenCountMap; + /** Preserve prompt token counts for graph formatting and pruning. */ + this.indexTokenCountMap = indexTokenCountMap; /** Extract contextMeta from the parent response (second-to-last in ordered chain; * last is the current user message). Seeds the pruner's calibration EMA for this run. */ @@ -970,6 +988,17 @@ class AgentClient extends BaseClient { tokenCounter, }); + const memoryMessages = + this.processMemory && this.memoryPayload + ? formatAgentMessages( + this.memoryPayload, + undefined, + toolSet, + skillPrimeResult?.skills, + formatOptions, + ).messages + : initialMessages; + /** * @param {BaseMessage[]} messages */ @@ -1010,7 +1039,7 @@ class AgentClient extends BaseClient { // } if (this.processMemory) { - memoryPromise = this.runMemory(messages); + memoryPromise = this.runMemory(memoryMessages); } /** Seed calibration state from previous run if encoding matches */ diff --git a/api/server/controllers/agents/client.test.js b/api/server/controllers/agents/client.test.js index 1c656c125c..296c4a1266 100644 --- a/api/server/controllers/agents/client.test.js +++ b/api/server/controllers/agents/client.test.js @@ -1,5 +1,5 @@ const { Providers } = require('@librechat/agents'); -const { Constants, EModelEndpoint } = require('librechat-data-provider'); +const { Constants, ContentTypes, EModelEndpoint } = require('librechat-data-provider'); const AgentClient = require('./client'); jest.mock('@librechat/agents', () => ({ @@ -1500,6 +1500,7 @@ describe('AgentClient - titleConvo', () => { beforeEach(() => { jest.clearAllMocks(); mockFormatInstructions.mockResolvedValue(''); + require('@librechat/api').countFormattedMessageTokens.mockImplementation(() => 42); mockAgent = { id: 'primary-agent', @@ -1544,7 +1545,7 @@ describe('AgentClient - titleConvo', () => { client.useMemory = jest.fn().mockResolvedValue(undefined); }); - it("applies shared request context plus each agent's own context docs only", async () => { + it('places request context inline and applies each agent context doc only once', async () => { const requestFile = makeTextFile('request-file', 'request.txt', 'Shared request context'); const primaryContext = makeTextFile( 'primary-context', @@ -1574,7 +1575,7 @@ describe('AgentClient - titleConvo', () => { ]); client.agentConfigs = new Map([['handoff-agent', handoffAgent]]); - await client.buildMessages( + const result = await client.buildMessages( [ { messageId: 'msg-1', @@ -1588,15 +1589,99 @@ describe('AgentClient - titleConvo', () => { {}, ); - expect(mockAgent.additional_instructions).toContain('Shared request context'); + expect(result.prompt[0].content).toContain('Shared request context'); + expect(mockAgent.additional_instructions).toContain('Primary private context'); + expect(mockAgent.additional_instructions).not.toContain('Shared request context'); expect(mockAgent.additional_instructions).not.toContain('Handoff private context'); - expect(handoffAgent.additional_instructions).toContain('Shared request context'); expect(handoffAgent.additional_instructions).toContain('Handoff private context'); + expect(handoffAgent.additional_instructions).not.toContain('Shared request context'); expect(handoffAgent.additional_instructions).not.toContain('Primary private context'); }); + it('places current request file context on the latest user message', async () => { + const currentFile = makeTextFile('current-file', 'current.txt', 'Current turn file body'); + const previousFileContext = + 'Attached document(s):\n```md\n# "previous.txt"\nPrevious turn file body\n```'; + + client.options.attachments = [currentFile]; + + const result = await client.buildMessages( + [ + { + messageId: 'msg-1', + parentMessageId: null, + sender: 'User', + text: 'What is written here?', + isCreatedByUser: true, + fileContext: previousFileContext, + }, + { + messageId: 'msg-2', + parentMessageId: 'msg-1', + sender: 'Assistant', + text: 'It describes the previous file.', + isCreatedByUser: false, + }, + { + messageId: 'msg-3', + parentMessageId: 'msg-2', + sender: 'User', + text: 'What is written here?', + isCreatedByUser: true, + }, + ], + 'msg-3', + {}, + ); + + expect(result.prompt[0].content).toContain('Previous turn file body'); + expect(result.prompt[2].content).toContain('Current turn file body'); + expect(result.prompt[2].content).toContain('What is written here?'); + expect(result.prompt[2].content).not.toContain('Previous turn file body'); + expect(client.memoryPayload[2].content).toContain('What is written here?'); + expect(client.memoryPayload[2].content).not.toContain('Current turn file body'); + expect(mockAgent.additional_instructions ?? '').not.toContain('Current turn file body'); + expect(result.prompt[2].content.indexOf('Current turn file body')).toBeLessThan( + result.prompt[2].content.indexOf('What is written here?'), + ); + }); + + it('persists canonical token counts while counting request file context for the prompt', async () => { + const { countFormattedMessageTokens } = require('@librechat/api'); + const currentFile = makeTextFile('current-file', 'current.txt', 'Current turn file body'); + + countFormattedMessageTokens.mockImplementation(({ content }) => { + const text = Array.isArray(content) + ? content.map((part) => part.text ?? part[ContentTypes.TEXT] ?? '').join('\n') + : String(content ?? ''); + return text.includes('Current turn file body') ? 200 : 20; + }); + + client.options.attachments = [currentFile]; + + const result = await client.buildMessages( + [ + { + messageId: 'msg-1', + parentMessageId: null, + sender: 'User', + text: 'What is written here?', + isCreatedByUser: true, + }, + ], + 'msg-1', + {}, + ); + + expect(result.prompt[0].content).toContain('Current turn file body'); + expect(result.tokenCountMap['msg-1']).toBe(20); + expect(result.promptTokens).toBe(200); + expect(client.indexTokenCountMap[0]).toBe(200); + expect(client.memoryPayload[0].content).toBe('What is written here?'); + }); + it('does not duplicate a file that is both request context and scoped context', async () => { const sharedFile = makeTextFile('shared-file', 'shared.txt', 'Shared duplicate context'); @@ -1604,7 +1689,7 @@ describe('AgentClient - titleConvo', () => { client.options.agentContextAttachmentsByAgentId = new Map([['primary-agent', [sharedFile]]]); client.agentConfigs = new Map(); - await client.buildMessages( + const result = await client.buildMessages( [ { messageId: 'msg-1', @@ -1618,10 +1703,10 @@ describe('AgentClient - titleConvo', () => { {}, ); - const occurrences = ( - mockAgent.additional_instructions.match(/Shared duplicate context/g) ?? [] - ).length; - expect(occurrences).toBe(1); + const inlineOccurrences = (result.prompt[0].content.match(/Shared duplicate context/g) ?? []) + .length; + expect(inlineOccurrences).toBe(1); + expect(mockAgent.additional_instructions ?? '').not.toContain('Shared duplicate context'); }); it('keeps direct chats with context-doc agents working without request attachments', async () => { diff --git a/packages/api/src/agents/client.spec.ts b/packages/api/src/agents/client.spec.ts new file mode 100644 index 0000000000..ab2c32ac24 --- /dev/null +++ b/packages/api/src/agents/client.spec.ts @@ -0,0 +1,54 @@ +import { ContentTypes } from 'librechat-data-provider'; +import { prependFileContext, type FormattedMessageWithContent } from './client'; + +describe('prependFileContext', () => { + it('prepends file context to string content', () => { + const message: FormattedMessageWithContent = { content: 'Answer this question.' }; + + prependFileContext(message, 'Attached file text'); + + expect(message.content).toBe('Attached file text\nAnswer this question.'); + }); + + it('prepends file context to the first text content part', () => { + const message: FormattedMessageWithContent = { + content: [ + { type: ContentTypes.IMAGE_URL, image_url: { url: 'data:image/png;base64,abc' } }, + { type: ContentTypes.TEXT, text: 'Answer this question.' }, + ], + }; + + prependFileContext(message, 'Attached file text'); + + expect(Array.isArray(message.content)).toBe(true); + if (!Array.isArray(message.content)) { + throw new Error('Expected array content'); + } + expect(message.content[1].text).toBe('Attached file text\nAnswer this question.'); + expect(message.content[0]).toEqual({ + type: ContentTypes.IMAGE_URL, + image_url: { url: 'data:image/png;base64,abc' }, + }); + }); + + it('adds a text content part when an array has no text part', () => { + const message: FormattedMessageWithContent = { + content: [{ type: ContentTypes.IMAGE_URL, image_url: { url: 'data:image/png;base64,abc' } }], + }; + + prependFileContext(message, 'Attached file text'); + + expect(message.content).toEqual([ + { type: ContentTypes.TEXT, text: 'Attached file text' }, + { type: ContentTypes.IMAGE_URL, image_url: { url: 'data:image/png;base64,abc' } }, + ]); + }); + + it('leaves content unchanged when file context is empty', () => { + const message: FormattedMessageWithContent = { content: 'Answer this question.' }; + + prependFileContext(message, ''); + + expect(message.content).toBe('Answer this question.'); + }); +}); diff --git a/packages/api/src/agents/client.ts b/packages/api/src/agents/client.ts index e4969bf7fa..5b6513499c 100644 --- a/packages/api/src/agents/client.ts +++ b/packages/api/src/agents/client.ts @@ -56,6 +56,42 @@ type ContentBlock = { tool_call?: { name?: string; args?: string; output?: string }; }; +export type FormattedMessageContentPart = { + type?: string; + text?: string; + [key: string]: unknown; +}; + +export type FormattedMessageWithContent = { + content?: string | FormattedMessageContentPart[]; +}; + +export function prependFileContext( + formattedMessage: FormattedMessageWithContent, + fileContext?: string | null, +): void { + if (!fileContext) { + return; + } + + if (typeof formattedMessage.content === 'string') { + formattedMessage.content = `${fileContext}\n${formattedMessage.content}`; + return; + } + + if (!Array.isArray(formattedMessage.content)) { + return; + } + + const textPart = formattedMessage.content.find((part) => part.type === ContentTypes.TEXT); + if (textPart != null && typeof textPart.text === 'string') { + textPart.text = `${fileContext}\n${textPart.text}`; + return; + } + + formattedMessage.content.unshift({ type: ContentTypes.TEXT, text: fileContext }); +} + function estimateImageDataTokens(data: string, isClaude: boolean): number { const dims = extractImageDimensions(data); if (dims == null) {