@lobehub/lobehub 2.0.0-next.54 → 2.0.0-next.55

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@lobehub/lobehub",
-  "version": "2.0.0-next.54",
+  "version": "2.0.0-next.55",
   "description": "LobeHub - an open-source,comprehensive AI Agent framework that supports speech synthesis, multimodal, and extensible Function Call plugin system. Supports one-click free deployment of your private ChatGPT/LLM web application.",
   "keywords": [
     "framework",

package/packages/database/src/models/__tests__/file.test.ts

@@ -4,7 +4,16 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
 
 import { FilesTabs, SortType } from '@/types/files';
 
-import { chunks, embeddings, fileChunks, files, globalFiles, knowledgeBaseFiles, knowledgeBases, users } from '../../schemas';
+import {
+  chunks,
+  embeddings,
+  fileChunks,
+  files,
+  globalFiles,
+  knowledgeBaseFiles,
+  knowledgeBases,
+  users,
+} from '../../schemas';
 import { LobeChatDatabase } from '../../type';
 import { FileModel } from '../file';
 import { getTestDB } from './_util';
@@ -340,13 +349,15 @@ describe('FileModel', () => {
     ];
 
     it('should query files for the user', async () => {
-      await fileModel.create({
+      const file1 = await fileModel.create({
         name: 'test-file-1.txt',
         url: 'https://example.com/test-file-1.txt',
         size: 100,
         fileType: 'text/plain',
       });
-      await fileModel.create({
+      // Add a small delay to ensure different timestamps
+      await new Promise((resolve) => setTimeout(resolve, 10));
+      const file2 = await fileModel.create({
         name: 'test-file-2.txt',
         url: 'https://example.com/test-file-2.txt',
         size: 200,
@@ -362,8 +373,9 @@ describe('FileModel', () => {
 
       const userFiles = await fileModel.query();
       expect(userFiles).toHaveLength(2);
-      expect(userFiles[0].name).toBe('test-file-2.txt');
-      expect(userFiles[1].name).toBe('test-file-1.txt');
+      // file2 should be first since it was created more recently
+      expect(userFiles[0].id).toBe(file2.id);
+      expect(userFiles[1].id).toBe(file1.id);
     });
 
     it('should filter files by name', async () => {
@@ -381,6 +393,14 @@ describe('FileModel', () => {
       expect(imageFiles[0].name).toBe('image.jpg');
     });
 
+    it('should filter audio files by category', async () => {
+      await serverDB.insert(files).values(sharedFileList);
+
+      const audioFiles = await fileModel.query({ category: FilesTabs.Audios });
+      expect(audioFiles).toHaveLength(1);
+      expect(audioFiles[0].name).toBe('audio.mp3');
+    });
+
     it('should sort files by name in ascending order', async () => {
       await serverDB.insert(files).values(sharedFileList);
 
@@ -1022,8 +1042,58 @@ describe('FileModel', () => {
   });
 
   describe('private getFileTypePrefix method', () => {
+    beforeEach(async () => {
+      // Create test files for all categories
+      await serverDB.insert(files).values([
+        {
+          id: 'video-file',
+          name: 'video.mp4',
+          url: 'https://example.com/video.mp4',
+          size: 1000,
+          fileType: 'video/mp4',
+          userId,
+        },
+        {
+          id: 'page-file',
+          name: 'page.html',
+          url: 'https://example.com/page.html',
+          size: 500,
+          fileType: 'text/html',
+          userId,
+        },
+        {
+          id: 'unknown-file',
+          name: 'unknown.xyz',
+          url: 'https://example.com/unknown.xyz',
+          size: 200,
+          fileType: 'application/xyz',
+          userId,
+        },
+      ]);
+    });
+
+    it('should filter video files correctly', async () => {
+      const result = await fileModel.query({ category: FilesTabs.Videos });
+      expect(result).toHaveLength(1);
+      expect(result[0].id).toBe('video-file');
+    });
+
+    it('should filter website/page files correctly', async () => {
+      const result = await fileModel.query({ category: FilesTabs.Websites });
+      expect(result).toHaveLength(1);
+      expect(result[0].id).toBe('page-file');
+    });
+
+    it('should handle Pages category (should use text/html like Websites)', async () => {
+      // FilesTabs.Pages is not explicitly handled in switch, falls to default
+      // which returns empty string, so it won't filter by file type
+      const result = await fileModel.query({ category: FilesTabs.Pages });
+      // Should return all files since default case returns empty string
+      expect(result.length).toBeGreaterThan(0);
+    });
+
     it('should handle unknown file category', async () => {
-      // This tests the default case in switch statement (line 312-313)
+      // This tests the default case in switch statement
       const unknownCategory = 'unknown' as FilesTabs;
 
       // We need to access the private method indirectly by testing the query method
@@ -1059,7 +1129,7 @@ describe('FileModel', () => {
       // Note: This is a simplified test since we can't easily create 3000+ chunks
       // But it will still exercise the batch deletion code path
 
-      // Insert chunks (this might need to be done through proper API)  
+      // Insert chunks (this might need to be done through proper API)
       // For testing purposes, we'll delete the file which should trigger the batch deletion
       await fileModel.delete(fileId, true);
 
@@ -1112,9 +1182,9 @@ describe('FileModel', () => {
 
       // 插入embeddings (1024维向量)
       const testEmbedding = new Array(1024).fill(0.1);
-      await serverDB.insert(embeddings).values([
-        { chunkId: chunkId1, embeddings: testEmbedding, model: 'test-model', userId },
-      ]);
+      await serverDB
+        .insert(embeddings)
+        .values([{ chunkId: chunkId1, embeddings: testEmbedding, model: 'test-model', userId }]);
 
       // 跳过 documentChunks 测试，因为需要先创建 documents 记录
 
@@ -1163,20 +1233,18 @@ describe('FileModel', () => {
       const chunkId = '550e8400-e29b-41d4-a716-446655440003';
 
       // 插入chunk
-      await serverDB.insert(chunks).values([
-        { id: chunkId, text: 'complete test chunk', userId, type: 'text' },
-      ]);
+      await serverDB
+        .insert(chunks)
+        .values([{ id: chunkId, text: 'complete test chunk', userId, type: 'text' }]);
 
       // 插入fileChunks关联
-      await serverDB.insert(fileChunks).values([
-        { fileId, chunkId, userId },
-      ]);
+      await serverDB.insert(fileChunks).values([{ fileId, chunkId, userId }]);
 
       // 插入embeddings
       const testEmbedding = new Array(1024).fill(0.1);
-      await serverDB.insert(embeddings).values([
-        { chunkId, embeddings: testEmbedding, model: 'test-model', userId },
-      ]);
+      await serverDB
+        .insert(embeddings)
+        .values([{ chunkId, embeddings: testEmbedding, model: 'test-model', userId }]);
 
       // 删除文件
       await fileModel.delete(fileId, true);
@@ -1206,7 +1274,6 @@ describe('FileModel', () => {
       expect(remainingFileChunks).toHaveLength(0);
     });
 
-
     it('should delete files that are in knowledge bases (removed protection)', async () => {
       // 测试修复后的逻辑：知识库中的文件也应该被删除
       const testFile = {
@@ -1223,19 +1290,17 @@ describe('FileModel', () => {
       const chunkId = '550e8400-e29b-41d4-a716-446655440007';
 
       // 插入chunk和关联数据
-      await serverDB.insert(chunks).values([
-        { id: chunkId, text: 'knowledge base chunk', userId, type: 'text' },
-      ]);
+      await serverDB
+        .insert(chunks)
+        .values([{ id: chunkId, text: 'knowledge base chunk', userId, type: 'text' }]);
 
-      await serverDB.insert(fileChunks).values([
-        { fileId, chunkId, userId },
-      ]);
+      await serverDB.insert(fileChunks).values([{ fileId, chunkId, userId }]);
 
       // 插入embeddings (1024维向量)
       const testEmbedding = new Array(1024).fill(0.1);
-      await serverDB.insert(embeddings).values([
-        { chunkId, embeddings: testEmbedding, model: 'test-model', userId },
-      ]);
+      await serverDB
+        .insert(embeddings)
+        .values([{ chunkId, embeddings: testEmbedding, model: 'test-model', userId }]);
 
       // 验证文件确实在知识库中
       const kbFile = await serverDB.query.knowledgeBaseFiles.findFirst({

package/packages/database/src/models/file.ts

@@ -211,9 +211,17 @@ export class FileModel {
       q ? ilike(files.name, `%${q}%`) : undefined,
       eq(files.userId, this.userId),
     );
-    if (category && category !== FilesTabs.All) {
+    if (category && category !== FilesTabs.All && category !== FilesTabs.Home) {
       const fileTypePrefix = this.getFileTypePrefix(category as FilesTabs);
-      whereClause = and(whereClause, ilike(files.fileType, `${fileTypePrefix}%`));
+      if (Array.isArray(fileTypePrefix)) {
+        // For multiple file types (e.g., Documents includes 'application' and 'custom')
+        whereClause = and(
+          whereClause,
+          or(...fileTypePrefix.map((prefix) => ilike(files.fileType, `${prefix}%`))),
+        );
+      } else {
+        whereClause = and(whereClause, ilike(files.fileType, `${fileTypePrefix}%`));
+      }
     }
 
     // 2. order part
@@ -308,13 +316,13 @@ export class FileModel {
   /**
    * get the corresponding file type prefix according to FilesTabs
    */
-  private getFileTypePrefix = (category: FilesTabs): string => {
+  private getFileTypePrefix = (category: FilesTabs): string | string[] => {
     switch (category) {
       case FilesTabs.Audios: {
         return 'audio';
       }
       case FilesTabs.Documents: {
-        return 'application';
+        return ['application', 'custom'];
       }
       case FilesTabs.Images: {
         return 'image';
@@ -322,6 +330,9 @@ export class FileModel {
       case FilesTabs.Videos: {
         return 'video';
       }
+      case FilesTabs.Websites: {
+        return 'text/html';
+      }
       default: {
         return '';
       }

package/packages/database/src/repositories/knowledge/index.test.ts

@@ -0,0 +1,300 @@
+// @vitest-environment node
+import { FilesTabs } from '@lobechat/types';
+import { beforeEach, describe, expect, it } from 'vitest';
+
+import { getTestDB } from '../../models/__tests__/_util';
+import { NewDocument, documents } from '../../schemas/document';
+import { NewFile, files } from '../../schemas/file';
+import { users } from '../../schemas/user';
+import { LobeChatDatabase } from '../../type';
+import { KnowledgeRepo } from './index';
+
+const userId = 'knowledge-test-user';
+const otherUserId = 'other-knowledge-user';
+
+let knowledgeRepo: KnowledgeRepo;
+
+const serverDB: LobeChatDatabase = await getTestDB();
+
+beforeEach(async () => {
+  // Clean up
+  await serverDB.delete(users);
+
+  // Create test users
+  await serverDB.insert(users).values([{ id: userId }, { id: otherUserId }]);
+
+  // Initialize repo
+  knowledgeRepo = new KnowledgeRepo(serverDB, userId);
+});
+
+describe('KnowledgeRepo', () => {
+  describe('query - Documents category filtering', () => {
+    beforeEach(async () => {
+      // Create test files
+      const testFiles: NewFile[] = [
+        {
+          fileType: 'application/pdf',
+          name: 'regular-pdf-file.pdf',
+          size: 1024,
+          url: 'file-pdf-url',
+          userId,
+        },
+        {
+          fileType: 'custom/other',
+          name: 'custom-file.txt',
+          size: 512,
+          url: 'custom-file-url',
+          userId,
+        },
+      ];
+
+      await serverDB.insert(files).values(testFiles);
+
+      // Create test documents
+      const testDocuments: NewDocument[] = [
+        // This should be EXCLUDED (sourceType='file')
+        {
+          content: 'PDF from file upload',
+          fileType: 'application/pdf',
+          filename: 'uploaded-pdf.pdf',
+          source: 'upload-source',
+          sourceType: 'file',
+          totalCharCount: 100,
+          totalLineCount: 10,
+          userId,
+        },
+        // This should be EXCLUDED (fileType='custom/document')
+        {
+          content: 'Editor document',
+          fileType: 'custom/document',
+          filename: 'editor-doc.md',
+          source: 'editor-source',
+          sourceType: 'file',
+          totalCharCount: 200,
+          totalLineCount: 20,
+          userId,
+        },
+        // This should be INCLUDED (application/pdf with sourceType='api')
+        {
+          content: 'PDF from API',
+          fileType: 'application/pdf',
+          filename: 'api-pdf.pdf',
+          source: 'api-source',
+          sourceType: 'api',
+          totalCharCount: 300,
+          totalLineCount: 30,
+          userId,
+        },
+        // This should be INCLUDED (custom/other with sourceType='web')
+        {
+          content: 'Custom web document',
+          fileType: 'custom/other',
+          filename: 'web-doc.txt',
+          source: 'web-source',
+          sourceType: 'web',
+          totalCharCount: 400,
+          totalLineCount: 40,
+          userId,
+        },
+      ];
+
+      await serverDB.insert(documents).values(testDocuments);
+    });
+
+    it('should exclude documents with fileType="custom/document" from Documents category', async () => {
+      const results = await knowledgeRepo.query({ category: FilesTabs.Documents });
+
+      // Should not include editor document (custom/document)
+      const editorDoc = results.find((item) => item.name === 'editor-doc.md');
+      expect(editorDoc).toBeUndefined();
+    });
+
+    it('should exclude documents with sourceType="file" from Documents category', async () => {
+      const results = await knowledgeRepo.query({ category: FilesTabs.Documents });
+
+      // Should not include uploaded PDF document (sourceType='file')
+      const uploadedPdf = results.find((item) => item.name === 'uploaded-pdf.pdf');
+      expect(uploadedPdf).toBeUndefined();
+    });
+
+    it('should include documents with sourceType="api" in Documents category', async () => {
+      const results = await knowledgeRepo.query({ category: FilesTabs.Documents });
+
+      // Should include API PDF (application/pdf with sourceType='api')
+      const apiPdf = results.find((item) => item.name === 'api-pdf.pdf');
+      expect(apiPdf).toBeDefined();
+      expect(apiPdf?.sourceType).toBe('document');
+      expect(apiPdf?.fileType).toBe('application/pdf');
+    });
+
+    it('should include documents with sourceType="web" in Documents category', async () => {
+      const results = await knowledgeRepo.query({ category: FilesTabs.Documents });
+
+      // Should include web document (custom/other with sourceType='web')
+      const webDoc = results.find((item) => item.name === 'web-doc.txt');
+      expect(webDoc).toBeDefined();
+      expect(webDoc?.sourceType).toBe('document');
+      expect(webDoc?.fileType).toBe('custom/other');
+    });
+
+    it('should include files from files table in Documents category', async () => {
+      const results = await knowledgeRepo.query({ category: FilesTabs.Documents });
+
+      // Should include regular files
+      const regularFile = results.find((item) => item.name === 'regular-pdf-file.pdf');
+      expect(regularFile).toBeDefined();
+      expect(regularFile?.sourceType).toBe('file');
+    });
+
+    it('should show all documents in All category (no filtering)', async () => {
+      const results = await knowledgeRepo.query({ category: FilesTabs.All });
+
+      // All category should include everything
+      expect(results.length).toBeGreaterThanOrEqual(6); // 2 files + 4 documents
+
+      const editorDoc = results.find((item) => item.name === 'editor-doc.md');
+      const uploadedPdf = results.find((item) => item.name === 'uploaded-pdf.pdf');
+
+      expect(editorDoc).toBeDefined();
+      expect(uploadedPdf).toBeDefined();
+    });
+
+    it('should apply both filters together in Documents category', async () => {
+      const results = await knowledgeRepo.query({ category: FilesTabs.Documents });
+
+      // Count documents with sourceType='document'
+      const documentTypeItems = results.filter((item) => item.sourceType === 'document');
+
+      // Should have exactly 2 documents (api-pdf and web-doc)
+      // Excluded: uploaded-pdf (sourceType='file') and editor-doc (fileType='custom/document')
+      expect(documentTypeItems).toHaveLength(2);
+
+      const names = documentTypeItems.map((item) => item.name).sort();
+      expect(names).toEqual(['api-pdf.pdf', 'web-doc.txt']);
+    });
+  });
+
+  describe('query - user isolation', () => {
+    beforeEach(async () => {
+      // Create files for current user
+      await serverDB.insert(files).values({
+        fileType: 'application/pdf',
+        name: 'user-file.pdf',
+        size: 1024,
+        url: 'user-file-url',
+        userId,
+      });
+
+      // Create files for other user
+      await serverDB.insert(files).values({
+        fileType: 'application/pdf',
+        name: 'other-user-file.pdf',
+        size: 1024,
+        url: 'other-file-url',
+        userId: otherUserId,
+      });
+
+      // Create documents for current user
+      await serverDB.insert(documents).values({
+        content: 'User document',
+        fileType: 'application/pdf',
+        filename: 'user-doc.pdf',
+        source: 'user-source',
+        sourceType: 'api',
+        totalCharCount: 100,
+        totalLineCount: 10,
+        userId,
+      });
+
+      // Create documents for other user
+      await serverDB.insert(documents).values({
+        content: 'Other user document',
+        fileType: 'application/pdf',
+        filename: 'other-doc.pdf',
+        source: 'other-source',
+        sourceType: 'api',
+        totalCharCount: 100,
+        totalLineCount: 10,
+        userId: otherUserId,
+      });
+    });
+
+    it('should only return current user items', async () => {
+      const results = await knowledgeRepo.query({ category: FilesTabs.All });
+
+      // Should only have items from current user
+      expect(results).toHaveLength(2);
+
+      const names = results.map((item) => item.name).sort();
+      expect(names).toEqual(['user-doc.pdf', 'user-file.pdf']);
+
+      // Should not include other user's items
+      const otherUserFile = results.find((item) => item.name === 'other-user-file.pdf');
+      const otherUserDoc = results.find((item) => item.name === 'other-doc.pdf');
+
+      expect(otherUserFile).toBeUndefined();
+      expect(otherUserDoc).toBeUndefined();
+    });
+  });
+
+  describe('query - search filtering', () => {
+    beforeEach(async () => {
+      await serverDB.insert(files).values([
+        {
+          fileType: 'application/pdf',
+          name: 'report-2024.pdf',
+          size: 1024,
+          url: 'report-url',
+          userId,
+        },
+        {
+          fileType: 'application/pdf',
+          name: 'invoice.pdf',
+          size: 512,
+          url: 'invoice-url',
+          userId,
+        },
+      ]);
+
+      await serverDB.insert(documents).values([
+        {
+          content: 'Annual report content',
+          fileType: 'application/pdf',
+          filename: 'annual-report.pdf',
+          source: 'api-source',
+          sourceType: 'api',
+          title: 'Annual Report',
+          totalCharCount: 1000,
+          totalLineCount: 100,
+          userId,
+        },
+        {
+          content: 'Meeting notes',
+          fileType: 'custom/other',
+          filename: 'notes.txt',
+          source: 'web-source',
+          sourceType: 'web',
+          title: 'Meeting Notes',
+          totalCharCount: 500,
+          totalLineCount: 50,
+          userId,
+        },
+      ]);
+    });
+
+    it('should filter by search query in file names', async () => {
+      const results = await knowledgeRepo.query({ category: FilesTabs.All, q: 'report' });
+
+      expect(results).toHaveLength(2);
+      const names = results.map((item) => item.name).sort();
+      expect(names).toEqual(['Annual Report', 'report-2024.pdf']);
+    });
+
+    it('should filter by search query in document titles', async () => {
+      const results = await knowledgeRepo.query({ category: FilesTabs.All, q: 'meeting' });
+
+      expect(results).toHaveLength(1);
+      expect(results[0].name).toBe('Meeting Notes');
+    });
+  });
+});