From 68eac104adbab4f7d62fb64da7e2bb14d35e3d0b Mon Sep 17 00:00:00 2001 From: Danny Avila Date: Mon, 18 May 2026 15:36:22 -0400 Subject: [PATCH] =?UTF-8?q?=F0=9F=97=82=EF=B8=8F=20fix:=20Scope=20Handoff?= =?UTF-8?q?=20Agent=20Context=20Docs=20(#13167)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * fix: Scope agent context docs to handoff agents * fix: Deduplicate scoped request context * refactor: Extract agent attachment helpers --- api/server/controllers/agents/client.js | 30 ++- api/server/controllers/agents/client.test.js | 179 ++++++++++++++++++ .../services/Endpoints/agents/initialize.js | 9 +- .../Endpoints/agents/initialize.spec.js | 21 +- .../src/agents/__tests__/initialize.test.ts | 39 ++++ packages/api/src/agents/attachments.test.ts | 81 ++++++++ packages/api/src/agents/attachments.ts | 112 +++++++++++ packages/api/src/agents/index.ts | 1 + packages/api/src/agents/initialize.ts | 30 ++- packages/api/src/agents/resources.test.ts | 8 + packages/api/src/agents/resources.ts | 33 +++- packages/api/src/files/context.ts | 3 +- 12 files changed, 529 insertions(+), 17 deletions(-) create mode 100644 packages/api/src/agents/attachments.test.ts create mode 100644 packages/api/src/agents/attachments.ts diff --git a/api/server/controllers/agents/client.js b/api/server/controllers/agents/client.js index 0338918412..d621b830f5 100644 --- a/api/server/controllers/agents/client.js +++ b/api/server/controllers/agents/client.js @@ -12,6 +12,7 @@ const { resolveHeaders, createSafeUser, initializeAgent, + countTokens, getBalanceConfig, omitTitleOptions, getProviderConfig, @@ -31,6 +32,8 @@ const { hydrateMissingIndexTokenCounts, injectSkillPrimes, isSkillPrimeMessage, + collectFileIds, + buildAgentScopedContext, buildSkillPrimeContentParts, buildInitialToolSessions, } = require('@librechat/api'); @@ -270,10 +273,15 @@ class AgentClient extends BaseClient { })) : []), ]; + const sharedRunAttachmentIds = new Set(); if (this.options.attachments) { const attachments = await this.options.attachments; const latestMessage = orderedMessages[orderedMessages.length - 1]; + for (const fileId of collectFileIds(attachments)) { + sharedRunAttachmentIds.add(fileId); + } + if (this.message_file_map) { this.message_file_map[latestMessage.messageId] = attachments; } else { @@ -402,6 +410,14 @@ class AgentClient extends BaseClient { const sharedRunContext = sharedRunContextParts.join('\n\n'); const memoryAgentEnabled = isMemoryAgentEnabled(this.options.req.config?.memory); + const agentScopedContext = await buildAgentScopedContext({ + agentIds: allAgents.map(({ agentId }) => agentId), + attachmentsByAgentId: this.options.agentContextAttachmentsByAgentId, + sharedRunAttachmentIds, + req: this.options.req, + tokenCountFn: (text) => countTokens(text), + }); + /** Preserve canonical pre-format token counts for all history entering graph formatting */ this.indexTokenCountMap = canonicalTokenCountMap; @@ -439,10 +455,14 @@ class AgentClient extends BaseClient { await Promise.all( allAgents.map(({ agent, agentId }) => { - const agentRunContext = - memoryContext && (agentId === this.options.agent.id || memoryAgentEnabled) - ? [sharedRunContext, memoryContext].filter(Boolean).join('\n\n') - : sharedRunContext; + const agentRunContextParts = [sharedRunContext]; + if (memoryContext && (agentId === this.options.agent.id || memoryAgentEnabled)) { + agentRunContextParts.push(memoryContext); + } + const scopedContext = agentScopedContext.get(agentId); + if (scopedContext) { + agentRunContextParts.push(scopedContext); + } return applyContextToAgent({ agent, @@ -450,7 +470,7 @@ class AgentClient extends BaseClient { logger, mcpManager, configServers, - sharedRunContext: agentRunContext, + sharedRunContext: agentRunContextParts.filter(Boolean).join('\n\n'), ephemeralAgent: agentId === this.options.agent.id ? ephemeralAgent : undefined, }); }), diff --git a/api/server/controllers/agents/client.test.js b/api/server/controllers/agents/client.test.js index 31bd5227d5..873cba9c58 100644 --- a/api/server/controllers/agents/client.test.js +++ b/api/server/controllers/agents/client.test.js @@ -13,6 +13,8 @@ jest.mock('@librechat/agents', () => ({ jest.mock('@librechat/api', () => ({ ...jest.requireActual('@librechat/api'), checkAccess: jest.fn(), + countFormattedMessageTokens: jest.fn(() => 42), + countTokens: jest.fn((text) => Math.ceil(String(text ?? '').length / 4)), initializeAgent: jest.fn(), createMemoryProcessor: jest.fn(), isMemoryAgentEnabled: jest.fn((config) => { @@ -1429,6 +1431,183 @@ describe('AgentClient - titleConvo', () => { }); }); + describe('buildMessages with request and agent-scoped context attachments', () => { + let client; + let mockReq; + let mockRes; + let mockAgent; + + const makeTextFile = (file_id, filename, text) => ({ + user: 'user-123', + file_id, + filename, + filepath: `/uploads/${filename}`, + object: 'file', + type: 'text/plain', + bytes: text.length, + embedded: false, + usage: 0, + source: 'text', + text, + }); + + beforeEach(() => { + jest.clearAllMocks(); + mockFormatInstructions.mockResolvedValue(''); + + mockAgent = { + id: 'primary-agent', + endpoint: EModelEndpoint.openAI, + provider: EModelEndpoint.openAI, + instructions: 'Primary instructions', + model_parameters: { + model: 'gpt-4', + }, + tools: [], + }; + + mockReq = { + user: { + id: 'user-123', + personalization: { + memories: true, + }, + }, + body: { + endpoint: EModelEndpoint.openAI, + fileTokenLimit: 1000, + }, + config: { + memory: { + disabled: true, + }, + }, + }; + mockRes = {}; + + client = new AgentClient({ + req: mockReq, + res: mockRes, + agent: mockAgent, + endpoint: EModelEndpoint.agents, + }); + client.conversationId = 'convo-123'; + client.responseMessageId = 'response-123'; + client.shouldSummarize = false; + client.maxContextTokens = 4096; + client.useMemory = jest.fn().mockResolvedValue(undefined); + }); + + it("applies shared request context plus each agent's own context docs only", async () => { + const requestFile = makeTextFile('request-file', 'request.txt', 'Shared request context'); + const primaryContext = makeTextFile( + 'primary-context', + 'primary.txt', + 'Primary private context', + ); + const handoffContext = makeTextFile( + 'handoff-context', + 'handoff.txt', + 'Handoff private context', + ); + const handoffAgent = { + id: 'handoff-agent', + endpoint: EModelEndpoint.openAI, + provider: EModelEndpoint.openAI, + instructions: 'Handoff instructions', + model_parameters: { + model: 'gpt-4', + }, + tools: [], + }; + + client.options.attachments = [requestFile]; + client.options.agentContextAttachmentsByAgentId = new Map([ + ['primary-agent', [primaryContext]], + ['handoff-agent', [handoffContext]], + ]); + client.agentConfigs = new Map([['handoff-agent', handoffAgent]]); + + await client.buildMessages( + [ + { + messageId: 'msg-1', + parentMessageId: null, + sender: 'User', + text: 'Use the available context.', + isCreatedByUser: true, + }, + ], + 'msg-1', + {}, + ); + + expect(mockAgent.additional_instructions).toContain('Shared request context'); + expect(mockAgent.additional_instructions).toContain('Primary private 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('Primary private context'); + }); + + 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'); + + client.options.attachments = [sharedFile]; + client.options.agentContextAttachmentsByAgentId = new Map([['primary-agent', [sharedFile]]]); + client.agentConfigs = new Map(); + + await client.buildMessages( + [ + { + messageId: 'msg-1', + parentMessageId: null, + sender: 'User', + text: 'Use the available context.', + isCreatedByUser: true, + }, + ], + 'msg-1', + {}, + ); + + const occurrences = ( + mockAgent.additional_instructions.match(/Shared duplicate context/g) ?? [] + ).length; + expect(occurrences).toBe(1); + }); + + it('keeps direct chats with context-doc agents working without request attachments', async () => { + const primaryContext = makeTextFile( + 'primary-context', + 'primary.txt', + 'Direct primary context', + ); + + client.options.agentContextAttachmentsByAgentId = new Map([ + ['primary-agent', [primaryContext]], + ]); + client.agentConfigs = new Map(); + + await client.buildMessages( + [ + { + messageId: 'msg-1', + parentMessageId: null, + sender: 'User', + text: 'Answer from your context.', + isCreatedByUser: true, + }, + ], + 'msg-1', + {}, + ); + + expect(mockAgent.additional_instructions).toContain('Direct primary context'); + }); + }); + describe('runMemory method', () => { let client; let mockReq; diff --git a/api/server/services/Endpoints/agents/initialize.js b/api/server/services/Endpoints/agents/initialize.js index f1db6d325a..bbbcd535ce 100644 --- a/api/server/services/Endpoints/agents/initialize.js +++ b/api/server/services/Endpoints/agents/initialize.js @@ -10,6 +10,7 @@ const { getCustomEndpointConfig, discoverConnectedAgents, resolveAgentScopedSkillIds, + buildAgentContextAttachmentsByAgentId, } = require('@librechat/api'); const { ResourceType, @@ -796,6 +797,11 @@ const initializeClient = async ({ req, res, signal, endpointOption }) => { } } + const agentContextAttachmentsByAgentId = buildAgentContextAttachmentsByAgentId([ + primaryConfig, + ...agentConfigs.values(), + ]); + let endpointConfig = appConfig.endpoints?.[primaryConfig.endpoint]; if (!isAgentsEndpoint(primaryConfig.endpoint) && !endpointConfig) { try { @@ -851,7 +857,8 @@ const initializeClient = async ({ req, res, signal, endpointOption }) => { agent: primaryConfig, spec: endpointOption.spec, iconURL: endpointOption.iconURL, - attachments: primaryConfig.attachments, + attachments: primaryConfig.requestAttachments ?? primaryConfig.attachments, + agentContextAttachmentsByAgentId, endpointType: endpointOption.endpointType, resendFiles: primaryConfig.resendFiles ?? true, maxContextTokens: primaryConfig.maxContextTokens, diff --git a/api/server/services/Endpoints/agents/initialize.spec.js b/api/server/services/Endpoints/agents/initialize.spec.js index 3bf6a67e50..fd1eb7c897 100644 --- a/api/server/services/Endpoints/agents/initialize.spec.js +++ b/api/server/services/Endpoints/agents/initialize.spec.js @@ -183,6 +183,15 @@ describe('initializeClient — processAgent ACL gate', () => { }); const edges = [{ from: PRIMARY_ID, to: AUTHORIZED_ID, edgeType: 'handoff' }]; + const requestAttachment = { file_id: 'request_file', filename: 'request.txt' }; + const primaryContextAttachment = { file_id: 'primary_context', filename: 'primary.txt' }; + const handoffContextAttachment = { file_id: 'handoff_context', filename: 'handoff.txt' }; + const primaryConfig = { + ...makePrimaryConfig(edges), + attachments: [primaryContextAttachment, requestAttachment], + requestAttachments: [requestAttachment], + agentContextAttachments: [primaryContextAttachment], + }; const handoffConfig = { id: AUTHORIZED_ID, edges: [], @@ -190,14 +199,13 @@ describe('initializeClient — processAgent ACL gate', () => { toolRegistry: new Map(), userMCPAuthMap: null, tool_resources: {}, + agentContextAttachments: [handoffContextAttachment], }; let callCount = 0; mockInitializeAgent.mockImplementation(() => { callCount++; - return callCount === 1 - ? Promise.resolve(makePrimaryConfig(edges)) - : Promise.resolve(handoffConfig); + return callCount === 1 ? Promise.resolve(primaryConfig) : Promise.resolve(handoffConfig); }); await initializeClient({ @@ -210,6 +218,13 @@ describe('initializeClient — processAgent ACL gate', () => { expect(mockInitializeAgent).toHaveBeenCalledTimes(2); expect(agentClientArgs.agent.edges).toHaveLength(1); expect(agentClientArgs.agent.edges[0].to).toBe(AUTHORIZED_ID); + expect(agentClientArgs.attachments).toEqual([requestAttachment]); + expect(agentClientArgs.agentContextAttachmentsByAgentId.get(PRIMARY_ID)).toEqual([ + primaryContextAttachment, + ]); + expect(agentClientArgs.agentContextAttachmentsByAgentId.get(AUTHORIZED_ID)).toEqual([ + handoffContextAttachment, + ]); }); }); diff --git a/packages/api/src/agents/__tests__/initialize.test.ts b/packages/api/src/agents/__tests__/initialize.test.ts index d7fa396024..5b229b1f58 100644 --- a/packages/api/src/agents/__tests__/initialize.test.ts +++ b/packages/api/src/agents/__tests__/initialize.test.ts @@ -502,6 +502,45 @@ describe('initializeAgent — stable and dynamic instruction fields', () => { }); }); +describe('initializeAgent — attachment scoping', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + it('keeps request attachments separate from agent context attachments', async () => { + const { primeResources } = jest.requireMock('../resources') as { + primeResources: jest.Mock; + }; + const requestFile = { file_id: 'request-file', filename: 'request.txt' }; + const agentContextFile = { file_id: 'agent-context-file', filename: 'agent-context.txt' }; + primeResources.mockResolvedValueOnce({ + attachments: [agentContextFile, requestFile], + requestAttachments: [requestFile], + agentContextAttachments: [agentContextFile], + tool_resources: undefined, + }); + + const { agent, req, res, loadTools, db } = createMocks(); + + const result = await initializeAgent( + { + req, + res, + agent, + loadTools, + endpointOption: { endpoint: EModelEndpoint.agents }, + allowedProviders: new Set([Providers.OPENAI]), + isInitialAgent: true, + }, + db, + ); + + expect(result.attachments).toEqual([agentContextFile, requestFile]); + expect(result.requestAttachments).toEqual([requestFile]); + expect(result.agentContextAttachments).toEqual([agentContextFile]); + }); +}); + describe('initializeAgent — maxContextTokens', () => { beforeEach(() => { jest.clearAllMocks(); diff --git a/packages/api/src/agents/attachments.test.ts b/packages/api/src/agents/attachments.test.ts new file mode 100644 index 0000000000..88bc4d1105 --- /dev/null +++ b/packages/api/src/agents/attachments.test.ts @@ -0,0 +1,81 @@ +import { FileSources } from 'librechat-data-provider'; +import type { IMongoFile } from '@librechat/data-schemas'; +import type { ServerRequest } from '~/types'; +import { + collectFileIds, + buildAgentScopedContext, + getAgentContextAttachments, + buildAgentContextAttachmentsByAgentId, +} from './attachments'; + +const makeTextFile = (file_id: string, filename: string, text: string): IMongoFile => + ({ + file_id, + filename, + text, + source: FileSources.text, + }) as IMongoFile; + +describe('agent attachment helpers', () => { + it('collects file ids from attachment-like files', () => { + const fileIds = collectFileIds([ + { file_id: 'file-1' }, + null, + { file_id: '' }, + { file_id: 'file-2' }, + { file_id: 'file-1' }, + ]); + + expect(Array.from(fileIds)).toEqual(['file-1', 'file-2']); + }); + + it('builds an agent context attachment map from initialized configs', () => { + const file = makeTextFile('context-file', 'context.txt', 'context'); + const attachmentsByAgentId = buildAgentContextAttachmentsByAgentId([ + { id: 'agent-a', agentContextAttachments: [file] }, + { id: 'agent-b', agentContextAttachments: [] }, + { id: null, agentContextAttachments: [file] }, + undefined, + ]); + + expect(attachmentsByAgentId.size).toBe(1); + expect(attachmentsByAgentId.get('agent-a')).toEqual([file]); + }); + + it('filters shared request files out of scoped context attachments', () => { + const shared = makeTextFile('shared-file', 'shared.txt', 'shared'); + const scoped = makeTextFile('scoped-file', 'scoped.txt', 'scoped'); + + const attachments = getAgentContextAttachments({ + agentId: 'agent-a', + attachmentsByAgentId: new Map([['agent-a', [shared, scoped]]]), + excludeFileIds: new Set(['shared-file']), + }); + + expect(attachments).toEqual([scoped]); + }); + + it('builds scoped context only from non-shared context documents', async () => { + const shared = makeTextFile('shared-file', 'shared.txt', 'Shared duplicate context'); + const scoped = makeTextFile('scoped-file', 'scoped.txt', 'Scoped private context'); + const req = { + body: { fileTokenLimit: 1000 }, + config: {}, + } as ServerRequest; + + const scopedContext = await buildAgentScopedContext({ + agentIds: ['agent-a', 'agent-b'], + attachmentsByAgentId: new Map([ + ['agent-a', [shared, scoped]], + ['agent-b', [shared]], + ]), + sharedRunAttachmentIds: new Set(['shared-file']), + req, + tokenCountFn: (text) => text.length, + }); + + expect(scopedContext.get('agent-a')).toContain('Scoped private context'); + expect(scopedContext.get('agent-a')).not.toContain('Shared duplicate context'); + expect(scopedContext.has('agent-b')).toBe(false); + }); +}); diff --git a/packages/api/src/agents/attachments.ts b/packages/api/src/agents/attachments.ts new file mode 100644 index 0000000000..4513d5b6e0 --- /dev/null +++ b/packages/api/src/agents/attachments.ts @@ -0,0 +1,112 @@ +import type { IMongoFile } from '@librechat/data-schemas'; +import type { ServerRequest } from '~/types'; +import type { TokenCountFn } from '~/utils/text'; +import { countTokens } from '~/utils/tokenizer'; +import { extractFileContext } from '~/files'; + +type FileWithId = { + file_id?: string | null; +}; + +export type AgentContextAttachmentCarrier = { + id?: string | null; + agentContextAttachments?: TFile[] | null; +}; + +export type AgentContextAttachmentsByAgentId = + | Map + | Record + | null + | undefined; + +export function collectFileIds( + files?: Array | null, +): Set { + const fileIds = new Set(); + for (const file of files ?? []) { + if (file?.file_id) { + fileIds.add(file.file_id); + } + } + return fileIds; +} + +export function buildAgentContextAttachmentsByAgentId( + configs: Iterable | null | undefined>, +): Map { + const attachmentsByAgentId = new Map(); + + for (const config of configs) { + if (!config?.id || !Array.isArray(config.agentContextAttachments)) { + continue; + } + if (config.agentContextAttachments.length === 0) { + continue; + } + attachmentsByAgentId.set(config.id, config.agentContextAttachments); + } + + return attachmentsByAgentId; +} + +export function getAgentContextAttachments({ + agentId, + attachmentsByAgentId, + excludeFileIds, +}: { + agentId: string; + attachmentsByAgentId: AgentContextAttachmentsByAgentId; + excludeFileIds?: Set; +}): TFile[] { + if (!attachmentsByAgentId) { + return []; + } + + const attachments: TFile[] = + attachmentsByAgentId instanceof Map + ? (attachmentsByAgentId.get(agentId) ?? []) + : (attachmentsByAgentId[agentId] ?? []); + + if (!excludeFileIds || excludeFileIds.size === 0) { + return attachments; + } + + return attachments.filter((file) => !file?.file_id || !excludeFileIds.has(file.file_id)); +} + +export async function buildAgentScopedContext({ + agentIds, + attachmentsByAgentId, + sharedRunAttachmentIds, + req, + tokenCountFn = countTokens, +}: { + agentIds: string[]; + attachmentsByAgentId: AgentContextAttachmentsByAgentId; + sharedRunAttachmentIds?: Set; + req?: ServerRequest; + tokenCountFn?: TokenCountFn; +}): Promise> { + const uniqueAgentIds = Array.from(new Set(agentIds.filter(Boolean))); + const entries = await Promise.all( + uniqueAgentIds.map(async (agentId) => { + const attachments = getAgentContextAttachments({ + agentId, + attachmentsByAgentId, + excludeFileIds: sharedRunAttachmentIds, + }); + if (attachments.length === 0) { + return [agentId, ''] as const; + } + + const context = await extractFileContext({ + attachments, + req, + tokenCountFn, + }); + return [agentId, context ?? ''] as const; + }), + ); + + return new Map(entries.filter(([, context]) => Boolean(context))); +} diff --git a/packages/api/src/agents/index.ts b/packages/api/src/agents/index.ts index 1a6ae696a0..31d2485a01 100644 --- a/packages/api/src/agents/index.ts +++ b/packages/api/src/agents/index.ts @@ -1,4 +1,5 @@ export * from './avatars'; +export * from './attachments'; export * from './chain'; export * from './client'; export * from './config'; diff --git a/packages/api/src/agents/initialize.ts b/packages/api/src/agents/initialize.ts index 1b668a9eb0..cfe32fedd3 100644 --- a/packages/api/src/agents/initialize.ts +++ b/packages/api/src/agents/initialize.ts @@ -117,7 +117,12 @@ function resolveAnthropicToolConflicts({ */ export type InitializedAgent = Agent & { tools: GenericTool[]; + /** @deprecated use requestAttachments or agentContextAttachments based on sharing semantics. */ attachments: IMongoFile[]; + /** Files attached to the current user message/run and safe to share across run agents. */ + requestAttachments: IMongoFile[]; + /** Files attached to this agent's permanent context via tool_resources. */ + agentContextAttachments: IMongoFile[]; toolContextMap: Record; dynamicToolContextMap?: Record; maxContextTokens: number; @@ -535,7 +540,12 @@ export async function initializeAgent( }); } - const { attachments: primedAttachments, tool_resources } = await primeResources({ + const { + attachments: primedAttachments, + requestAttachments: primedRequestAttachments, + agentContextAttachments: primedAgentContextAttachments, + tool_resources, + } = await primeResources({ req: req as never, getFiles: db.getFiles as never, filterFiles: db.filterFilesByAgentAccess, @@ -959,9 +969,17 @@ export async function initializeAgent( const maxOutputTokensNum = Number(maxOutputTokens) || 0; const baseContextTokens = Math.max(0, agentMaxContextNum - maxOutputTokensNum); - const finalAttachments: IMongoFile[] = (primedAttachments ?? []) - .filter((a): a is TFile => a != null) - .map((a) => a as unknown as IMongoFile); + const toMongoFiles = (files: Array | undefined): IMongoFile[] => + (files ?? []).filter((a): a is TFile => a != null).map((a) => a as unknown as IMongoFile); + + const finalAttachments: IMongoFile[] = toMongoFiles(primedAttachments); + const requestAttachments: IMongoFile[] = toMongoFiles(primedRequestAttachments); + const agentContextAttachments: IMongoFile[] = toMongoFiles(primedAgentContextAttachments); + + const compatibilityAttachments = + finalAttachments.length > 0 + ? finalAttachments + : requestAttachments.concat(agentContextAttachments); const endpointConfigs = req.config?.endpoints; const providerConfig = @@ -992,7 +1010,9 @@ export async function initializeAgent( activeSkillNames, manualSkillPrimes, alwaysApplySkillPrimes, - attachments: finalAttachments, + attachments: compatibilityAttachments, + requestAttachments, + agentContextAttachments, toolContextMap: toolContextMap ?? {}, dynamicToolContextMap: dynamicToolContextMap ?? {}, useLegacyContent: !!options.useLegacyContent, diff --git a/packages/api/src/agents/resources.test.ts b/packages/api/src/agents/resources.test.ts index 718e5cbdec..04849a5177 100644 --- a/packages/api/src/agents/resources.test.ts +++ b/packages/api/src/agents/resources.test.ts @@ -84,6 +84,8 @@ describe('primeResources', () => { agentId: 'agent_test', }); expect(result.attachments).toEqual(mockOcrFiles); + expect(result.agentContextAttachments).toEqual(mockOcrFiles); + expect(result.requestAttachments).toBeUndefined(); expect(result.tool_resources).toEqual({}); }); }); @@ -423,6 +425,8 @@ describe('primeResources', () => { expect(result.attachments).toHaveLength(2); expect(result.attachments?.[0]?.file_id).toBe('ocr-file-1'); expect(result.attachments?.[1]?.file_id).toBe('file1'); + expect(result.agentContextAttachments).toEqual(mockOcrFiles); + expect(result.requestAttachments).toEqual(mockAttachmentFiles); }); it('should include both context (as `ocr` resource) files and attachment files', async () => { @@ -475,6 +479,8 @@ describe('primeResources', () => { expect(result.attachments).toHaveLength(2); expect(result.attachments?.[0]?.file_id).toBe('ocr-file-1'); expect(result.attachments?.[1]?.file_id).toBe('file1'); + expect(result.agentContextAttachments).toEqual(mockOcrFiles); + expect(result.requestAttachments).toEqual(mockAttachmentFiles); }); it('should prevent duplicate files when same file exists in context tool_resource and attachments', async () => { @@ -528,6 +534,8 @@ describe('primeResources', () => { expect(result.attachments).toHaveLength(2); expect(result.attachments?.filter((f) => f?.file_id === 'shared-file-id')).toHaveLength(1); expect(result.attachments?.find((f) => f?.file_id === 'unique-file')).toBeDefined(); + expect(result.agentContextAttachments).toEqual(mockOcrFiles); + expect(result.requestAttachments).toEqual(mockAttachmentFiles); }); it('should still categorize duplicate files for tool_resources', async () => { diff --git a/packages/api/src/agents/resources.ts b/packages/api/src/agents/resources.ts index 56cd9b6c0d..47239e54bb 100644 --- a/packages/api/src/agents/resources.ts +++ b/packages/api/src/agents/resources.ts @@ -174,8 +174,12 @@ export const primeResources = async ({ agentId?: string; }): Promise<{ attachments: Array | undefined; + requestAttachments: Array | undefined; + agentContextAttachments: Array | undefined; tool_resources: AgentToolResources | undefined; }> => { + const requestAttachments: Array = []; + const agentContextAttachments: Array = []; try { /** * Array to collect all unique files that will be returned as attachments @@ -269,6 +273,7 @@ export const primeResources = async ({ // Add to attachments attachments.push(file); + agentContextAttachments.push(file); attachmentFileIds.add(file.file_id); // Categorize for tool resources @@ -282,10 +287,17 @@ export const primeResources = async ({ } if (!_attachments) { - return { attachments: attachments.length > 0 ? attachments : undefined, tool_resources }; + return { + attachments: attachments.length > 0 ? attachments : undefined, + requestAttachments: undefined, + agentContextAttachments: + agentContextAttachments.length > 0 ? agentContextAttachments : undefined, + tool_resources, + }; } const files = await _attachments; + const requestAttachmentFileIds = new Set(); for (const file of files) { if (!file) { @@ -300,16 +312,30 @@ export const primeResources = async ({ }); if (file.file_id && attachmentFileIds.has(file.file_id)) { + if (!requestAttachmentFileIds.has(file.file_id)) { + requestAttachments.push(file); + requestAttachmentFileIds.add(file.file_id); + } continue; } attachments.push(file); + if (!file.file_id || !requestAttachmentFileIds.has(file.file_id)) { + requestAttachments.push(file); + } if (file.file_id) { attachmentFileIds.add(file.file_id); + requestAttachmentFileIds.add(file.file_id); } } - return { attachments: attachments.length > 0 ? attachments : [], tool_resources }; + return { + attachments: attachments.length > 0 ? attachments : [], + requestAttachments, + agentContextAttachments: + agentContextAttachments.length > 0 ? agentContextAttachments : undefined, + tool_resources, + }; } catch (error) { logger.error('Error priming resources', error); @@ -328,6 +354,9 @@ export const primeResources = async ({ return { attachments: safeAttachments, + requestAttachments: safeAttachments, + agentContextAttachments: + agentContextAttachments.length > 0 ? agentContextAttachments : undefined, tool_resources: _tool_resources, }; } diff --git a/packages/api/src/files/context.ts b/packages/api/src/files/context.ts index 1da6c387ce..36209f34c1 100644 --- a/packages/api/src/files/context.ts +++ b/packages/api/src/files/context.ts @@ -3,6 +3,7 @@ import { FileSources, mergeFileConfig } from 'librechat-data-provider'; import type { IMongoFile } from '@librechat/data-schemas'; import type { ServerRequest } from '~/types'; import { processTextWithTokenLimit } from '~/utils/text'; +import type { TokenCountFn } from '~/utils/text'; /** * Extracts text context from attachments and returns formatted text. @@ -20,7 +21,7 @@ export async function extractFileContext({ }: { attachments: IMongoFile[]; req?: ServerRequest; - tokenCountFn: (text: string) => number; + tokenCountFn: TokenCountFn; }): Promise { if (!attachments || attachments.length === 0) { return undefined;