🔗 fix: File Citation Processing to Use Tool Artifacts

api/test/services/Files/processFileCitations.test.js

@@ -0,0 +1,337 @@
+const { Tools } = require('librechat-data-provider');
+const {
+  processFileCitations,
+  applyCitationLimits,
+  enhanceSourcesWithMetadata,
+} = require('~/server/services/Files/Citations');
+
+// Mock dependencies
+jest.mock('~/models', () => ({
+  Files: {
+    find: jest.fn().mockResolvedValue([]),
+  },
+}));
+
+jest.mock('~/models/Role', () => ({
+  getRoleByName: jest.fn(),
+}));
+
+jest.mock('@librechat/api', () => ({
+  checkAccess: jest.fn().mockResolvedValue(true),
+}));
+
+jest.mock('~/server/services/Config/getCustomConfig', () => ({
+  getCustomConfig: jest.fn().mockResolvedValue({
+    endpoints: {
+      agents: {
+        maxCitations: 30,
+        maxCitationsPerFile: 5,
+        minRelevanceScore: 0.45,
+      },
+    },
+    fileStrategy: 'local',
+  }),
+}));
+
+jest.mock('~/config', () => ({
+  logger: {
+    debug: jest.fn(),
+    error: jest.fn(),
+    warn: jest.fn(),
+  },
+}));
+
+describe('processFileCitations', () => {
+  const mockReq = {
+    user: {
+      id: 'user123',
+    },
+  };
+
+  const mockMetadata = {
+    run_id: 'run123',
+    thread_id: 'conv123',
+  };
+
+  describe('file search artifact processing', () => {
+    it('should process file search artifacts correctly', async () => {
+      const toolArtifact = {
+        [Tools.file_search]: {
+          sources: [
+            {
+              fileId: 'file_123',
+              fileName: 'example.pdf',
+              pages: [5],
+              relevance: 0.85,
+              type: 'file',
+              pageRelevance: { 5: 0.85 },
+              content: 'This is the content',
+            },
+            {
+              fileId: 'file_456',
+              fileName: 'document.txt',
+              pages: [],
+              relevance: 0.72,
+              type: 'file',
+              pageRelevance: {},
+              content: 'Another document',
+            },
+          ],
+        },
+      };
+
+      const result = await processFileCitations({
+        toolArtifact,
+        toolCallId: 'call_123',
+        metadata: mockMetadata,
+        user: mockReq.user,
+      });
+
+      expect(result).toBeTruthy();
+      expect(result.type).toBe('file_search');
+      expect(result.file_search.sources).toHaveLength(2);
+      expect(result.file_search.sources[0].fileId).toBe('file_123');
+      expect(result.file_search.sources[0].relevance).toBe(0.85);
+    });
+
+    it('should return null for non-file_search tools', async () => {
+      const result = await processFileCitations({
+        toolArtifact: { other_tool: {} },
+        toolCallId: 'call_123',
+        metadata: mockMetadata,
+        user: mockReq.user,
+      });
+
+      expect(result).toBeNull();
+    });
+
+    it('should filter results below relevance threshold', async () => {
+      const toolArtifact = {
+        [Tools.file_search]: {
+          sources: [
+            {
+              fileId: 'file_789',
+              fileName: 'low_relevance.pdf',
+              pages: [],
+              relevance: 0.2,
+              type: 'file',
+              pageRelevance: {},
+              content: 'Low relevance content',
+            },
+          ],
+        },
+      };
+
+      const result = await processFileCitations({
+        toolArtifact,
+        toolCallId: 'call_123',
+        metadata: mockMetadata,
+        user: mockReq.user,
+      });
+
+      expect(result).toBeNull();
+    });
+
+    it('should return null when artifact is missing file_search data', async () => {
+      const result = await processFileCitations({
+        toolArtifact: {},
+        toolCallId: 'call_123',
+        metadata: mockMetadata,
+        user: mockReq.user,
+      });
+
+      expect(result).toBeNull();
+    });
+  });
+
+  describe('applyCitationLimits', () => {
+    it('should limit citations per file and total', () => {
+      const sources = [
+        { fileId: 'file1', relevance: 0.9 },
+        { fileId: 'file1', relevance: 0.8 },
+        { fileId: 'file1', relevance: 0.7 },
+        { fileId: 'file2', relevance: 0.85 },
+        { fileId: 'file2', relevance: 0.75 },
+      ];
+
+      const result = applyCitationLimits(sources, 3, 2);
+
+      expect(result).toHaveLength(3);
+      expect(result[0].relevance).toBe(0.9);
+      expect(result[1].relevance).toBe(0.85);
+      expect(result[2].relevance).toBe(0.8);
+    });
+  });
+
+  describe('enhanceSourcesWithMetadata', () => {
+    const { Files } = require('~/models');
+    const mockCustomConfig = {
+      fileStrategy: 'local',
+    };
+
+    beforeEach(() => {
+      jest.clearAllMocks();
+    });
+
+    it('should enhance sources with file metadata from database', async () => {
+      const sources = [
+        {
+          fileId: 'file_123',
+          fileName: 'example.pdf',
+          relevance: 0.85,
+          type: 'file',
+        },
+        {
+          fileId: 'file_456',
+          fileName: 'document.txt',
+          relevance: 0.72,
+          type: 'file',
+        },
+      ];
+
+      Files.find.mockResolvedValue([
+        {
+          file_id: 'file_123',
+          filename: 'example_from_db.pdf',
+          source: 's3',
+        },
+        {
+          file_id: 'file_456',
+          filename: 'document_from_db.txt',
+          source: 'local',
+        },
+      ]);
+
+      const result = await enhanceSourcesWithMetadata(sources, mockCustomConfig);
+
+      expect(Files.find).toHaveBeenCalledWith({ file_id: { $in: ['file_123', 'file_456'] } });
+      expect(result).toHaveLength(2);
+
+      expect(result[0]).toEqual({
+        fileId: 'file_123',
+        fileName: 'example_from_db.pdf',
+        relevance: 0.85,
+        type: 'file',
+        metadata: {
+          storageType: 's3',
+        },
+      });
+
+      expect(result[1]).toEqual({
+        fileId: 'file_456',
+        fileName: 'document_from_db.txt',
+        relevance: 0.72,
+        type: 'file',
+        metadata: {
+          storageType: 'local',
+        },
+      });
+    });
+
+    it('should preserve existing metadata and source data', async () => {
+      const sources = [
+        {
+          fileId: 'file_123',
+          fileName: 'example.pdf',
+          relevance: 0.85,
+          type: 'file',
+          pages: [1, 2, 3],
+          content: 'Some content',
+          metadata: {
+            existingField: 'value',
+          },
+        },
+      ];
+
+      Files.find.mockResolvedValue([
+        {
+          file_id: 'file_123',
+          filename: 'example_from_db.pdf',
+          source: 'gcs',
+        },
+      ]);
+
+      const result = await enhanceSourcesWithMetadata(sources, mockCustomConfig);
+
+      expect(result[0]).toEqual({
+        fileId: 'file_123',
+        fileName: 'example_from_db.pdf',
+        relevance: 0.85,
+        type: 'file',
+        pages: [1, 2, 3],
+        content: 'Some content',
+        metadata: {
+          existingField: 'value',
+          storageType: 'gcs',
+        },
+      });
+    });
+
+    it('should handle missing file metadata gracefully', async () => {
+      const sources = [
+        {
+          fileId: 'file_789',
+          fileName: 'missing.pdf',
+          relevance: 0.9,
+          type: 'file',
+        },
+      ];
+
+      Files.find.mockResolvedValue([]);
+
+      const result = await enhanceSourcesWithMetadata(sources, mockCustomConfig);
+
+      expect(result[0]).toEqual({
+        fileId: 'file_789',
+        fileName: 'missing.pdf',
+        relevance: 0.9,
+        type: 'file',
+        metadata: {
+          storageType: 'local', // Falls back to customConfig.fileStrategy
+        },
+      });
+    });
+
+    it('should handle database errors gracefully', async () => {
+      const sources = [
+        {
+          fileId: 'file_123',
+          fileName: 'example.pdf',
+          relevance: 0.85,
+          type: 'file',
+        },
+      ];
+
+      Files.find.mockRejectedValue(new Error('Database error'));
+
+      const result = await enhanceSourcesWithMetadata(sources, mockCustomConfig);
+
+      expect(result[0]).toEqual({
+        fileId: 'file_123',
+        fileName: 'example.pdf',
+        relevance: 0.85,
+        type: 'file',
+        metadata: {
+          storageType: 'local',
+        },
+      });
+    });
+
+    it('should deduplicate file IDs when querying database', async () => {
+      const sources = [
+        { fileId: 'file_123', fileName: 'doc1.pdf', relevance: 0.9, type: 'file' },
+        { fileId: 'file_123', fileName: 'doc1.pdf', relevance: 0.8, type: 'file' },
+        { fileId: 'file_456', fileName: 'doc2.pdf', relevance: 0.7, type: 'file' },
+      ];
+
+      Files.find.mockResolvedValue([
+        { file_id: 'file_123', filename: 'document1.pdf', source: 's3' },
+        { file_id: 'file_456', filename: 'document2.pdf', source: 'local' },
+      ]);
+
+      await enhanceSourcesWithMetadata(sources, mockCustomConfig);
+
+      expect(Files.find).toHaveBeenCalledWith({ file_id: { $in: ['file_123', 'file_456'] } });
+    });
+  });
+});