@lobehub/chat 1.12.9 → 1.12.11

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

@@ -1,7 +1,8 @@
 // @vitest-environment node
-import { eq } from 'drizzle-orm';
+import { eq, inArray } from 'drizzle-orm';
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
 
+import { getServerDBConfig, serverDBEnv } from '@/config/db';
 import { getTestDBInstance } from '@/database/server/core/dbForTest';
 import { FilesTabs, SortType } from '@/types/files';
 
@@ -22,6 +23,20 @@ vi.mock('@/database/server/core/db', async () => ({
   },
 }));
 
+let DISABLE_REMOVE_GLOBAL_FILE = false;
+
+vi.mock('@/config/db', async () => ({
+  get serverDBEnv() {
+    return {
+      get DISABLE_REMOVE_GLOBAL_FILE() {
+        return DISABLE_REMOVE_GLOBAL_FILE;
+      },
+      DATABASE_TEST_URL: process.env.DATABASE_TEST_URL,
+      DATABASE_DRIVER: 'node',
+    };
+  },
+}));
+
 const userId = 'file-model-test-user-id';
 const fileModel = new FileModel(userId);
 
@@ -73,18 +88,21 @@ describe('FileModel', () => {
     });
   });
 
-  it('should create a global file', async () => {
-    const globalFile = {
-      hashId: 'test-hash',
-      fileType: 'text/plain',
-      size: 100,
-      url: 'https://example.com/global-file.txt',
-      metadata: { key: 'value' },
-    };
+  describe('createGlobalFile', () => {
+    it('should create a global file', async () => {
+      const globalFile = {
+        hashId: 'test-hash',
+        fileType: 'text/plain',
+        size: 100,
+        url: 'https://example.com/global-file.txt',
+        metadata: { key: 'value' },
+      };
 
-    const result = await fileModel.createGlobalFile(globalFile);
-    expect(result[0]).toMatchObject(globalFile);
+      const result = await fileModel.createGlobalFile(globalFile);
+      expect(result[0]).toMatchObject(globalFile);
+    });
   });
+
   describe('checkHash', () => {
     it('should return isExist: false for non-existent hash', async () => {
       const result = await fileModel.checkHash('non-existent-hash');
@@ -113,58 +131,183 @@ describe('FileModel', () => {
     });
   });
 
-  it('should delete a file by id', async () => {
-    const { id } = await fileModel.create({
-      name: 'test-file.txt',
-      url: 'https://example.com/test-file.txt',
-      size: 100,
-      fileType: 'text/plain',
-    });
+  describe('delete', () => {
+    it('should delete a file by id', async () => {
+      await fileModel.createGlobalFile({
+        hashId: '1',
+        url: 'https://example.com/file1.txt',
+        size: 100,
+        fileType: 'text/plain',
+      });
 
-    await fileModel.delete(id);
+      const { id } = await fileModel.create({
+        name: 'test-file.txt',
+        url: 'https://example.com/test-file.txt',
+        size: 100,
+        fileType: 'text/plain',
+        fileHash: '1',
+      });
 
-    const file = await serverDB.query.files.findFirst({ where: eq(files.id, id) });
-    expect(file).toBeUndefined();
-  });
+      await fileModel.delete(id);
 
-  it('should delete multiple files', async () => {
-    const file1 = await fileModel.create({
-      name: 'file1.txt',
-      url: 'https://example.com/file1.txt',
-      size: 100,
-      fileType: 'text/plain',
-    });
-    const file2 = await fileModel.create({
-      name: 'file2.txt',
-      url: 'https://example.com/file2.txt',
-      size: 200,
-      fileType: 'text/plain',
+      const file = await serverDB.query.files.findFirst({ where: eq(files.id, id) });
+      const globalFile = await serverDB.query.globalFiles.findFirst({
+        where: eq(globalFiles.hashId, '1'),
+      });
+
+      expect(file).toBeUndefined();
+      expect(globalFile).toBeUndefined();
     });
+    it('should delete a file by id but global file not removed ', async () => {
+      DISABLE_REMOVE_GLOBAL_FILE = true;
+      await fileModel.createGlobalFile({
+        hashId: '1',
+        url: 'https://example.com/file1.txt',
+        size: 100,
+        fileType: 'text/plain',
+      });
 
-    await fileModel.deleteMany([file1.id, file2.id]);
+      const { id } = await fileModel.create({
+        name: 'test-file.txt',
+        url: 'https://example.com/test-file.txt',
+        size: 100,
+        fileType: 'text/plain',
+        fileHash: '1',
+      });
+
+      await fileModel.delete(id);
+
+      const file = await serverDB.query.files.findFirst({ where: eq(files.id, id) });
+      const globalFile = await serverDB.query.globalFiles.findFirst({
+        where: eq(globalFiles.hashId, '1'),
+      });
 
-    const remainingFiles = await serverDB.query.files.findMany({ where: eq(files.userId, userId) });
-    expect(remainingFiles).toHaveLength(0);
+      expect(file).toBeUndefined();
+      expect(globalFile).toBeDefined();
+      DISABLE_REMOVE_GLOBAL_FILE = false;
+    });
   });
 
-  it('should clear all files for the user', async () => {
-    await fileModel.create({
-      name: 'test-file-1.txt',
-      url: 'https://example.com/test-file-1.txt',
-      size: 100,
-      fileType: 'text/plain',
+  describe('deleteMany', () => {
+    it('should delete multiple files', async () => {
+      await fileModel.createGlobalFile({
+        hashId: '1',
+        url: 'https://example.com/file1.txt',
+        size: 100,
+        fileType: 'text/plain',
+      });
+      await fileModel.createGlobalFile({
+        hashId: '2',
+        url: 'https://example.com/file2.txt',
+        size: 200,
+        fileType: 'text/plain',
+      });
+
+      const file1 = await fileModel.create({
+        name: 'file1.txt',
+        url: 'https://example.com/file1.txt',
+        size: 100,
+        fileHash: '1',
+        fileType: 'text/plain',
+      });
+      const file2 = await fileModel.create({
+        name: 'file2.txt',
+        url: 'https://example.com/file2.txt',
+        size: 200,
+        fileType: 'text/plain',
+        fileHash: '2',
+      });
+      const globalFilesResult = await serverDB.query.globalFiles.findMany({
+        where: inArray(globalFiles.hashId, ['1', '2']),
+      });
+      expect(globalFilesResult).toHaveLength(2);
+
+      await fileModel.deleteMany([file1.id, file2.id]);
+
+      const remainingFiles = await serverDB.query.files.findMany({
+        where: eq(files.userId, userId),
+      });
+      const globalFilesResult2 = await serverDB.query.globalFiles.findMany({
+        where: inArray(
+          globalFiles.hashId,
+          remainingFiles.map((i) => i.fileHash as string),
+        ),
+      });
+
+      expect(remainingFiles).toHaveLength(0);
+      expect(globalFilesResult2).toHaveLength(0);
     });
-    await fileModel.create({
-      name: 'test-file-2.txt',
-      url: 'https://example.com/test-file-2.txt',
-      size: 200,
-      fileType: 'text/plain',
+    it('should delete multiple files but not remove global files if DISABLE_REMOVE_GLOBAL_FILE=true', async () => {
+      DISABLE_REMOVE_GLOBAL_FILE = true;
+      await fileModel.createGlobalFile({
+        hashId: '1',
+        url: 'https://example.com/file1.txt',
+        size: 100,
+        fileType: 'text/plain',
+      });
+      await fileModel.createGlobalFile({
+        hashId: '2',
+        url: 'https://example.com/file2.txt',
+        size: 200,
+        fileType: 'text/plain',
+      });
+
+      const file1 = await fileModel.create({
+        name: 'file1.txt',
+        url: 'https://example.com/file1.txt',
+        size: 100,
+        fileType: 'text/plain',
+        fileHash: '1',
+      });
+      const file2 = await fileModel.create({
+        name: 'file2.txt',
+        url: 'https://example.com/file2.txt',
+        size: 200,
+        fileType: 'text/plain',
+        fileHash: '2',
+      });
+
+      const globalFilesResult = await serverDB.query.globalFiles.findMany({
+        where: inArray(globalFiles.hashId, ['1', '2']),
+      });
+
+      expect(globalFilesResult).toHaveLength(2);
+
+      await fileModel.deleteMany([file1.id, file2.id]);
+
+      const remainingFiles = await serverDB.query.files.findMany({
+        where: eq(files.userId, userId),
+      });
+      const globalFilesResult2 = await serverDB.query.globalFiles.findMany({
+        where: inArray(globalFiles.hashId, ['1', '2']),
+      });
+
+      expect(remainingFiles).toHaveLength(0);
+      expect(globalFilesResult2).toHaveLength(2);
+      DISABLE_REMOVE_GLOBAL_FILE = false;
     });
+  });
 
-    await fileModel.clear();
+  describe('clear', () => {
+    it('should clear all files for the user', async () => {
+      await fileModel.create({
+        name: 'test-file-1.txt',
+        url: 'https://example.com/test-file-1.txt',
+        size: 100,
+        fileType: 'text/plain',
+      });
+      await fileModel.create({
+        name: 'test-file-2.txt',
+        url: 'https://example.com/test-file-2.txt',
+        size: 200,
+        fileType: 'text/plain',
+      });
+
+      await fileModel.clear();
 
-    const userFiles = await serverDB.query.files.findMany({ where: eq(files.userId, userId) });
-    expect(userFiles).toHaveLength(0);
+      const userFiles = await serverDB.query.files.findMany({ where: eq(files.userId, userId) });
+      expect(userFiles).toHaveLength(0);
+    });
   });
 
   describe('Query', () => {
@@ -334,22 +477,24 @@ describe('FileModel', () => {
     });
   });
 
-  it('should find a file by id', async () => {
-    const { id } = await fileModel.create({
-      name: 'test-file.txt',
-      url: 'https://example.com/test-file.txt',
-      size: 100,
-      fileType: 'text/plain',
-    });
+  describe('findById', () => {
+    it('should find a file by id', async () => {
+      const { id } = await fileModel.create({
+        name: 'test-file.txt',
+        url: 'https://example.com/test-file.txt',
+        size: 100,
+        fileType: 'text/plain',
+      });
 
-    const file = await fileModel.findById(id);
-    expect(file).toMatchObject({
-      id,
-      name: 'test-file.txt',
-      url: 'https://example.com/test-file.txt',
-      size: 100,
-      fileType: 'text/plain',
-      userId,
+      const file = await fileModel.findById(id);
+      expect(file).toMatchObject({
+        id,
+        name: 'test-file.txt',
+        url: 'https://example.com/test-file.txt',
+        size: 100,
+        fileType: 'text/plain',
+        userId,
+      });
     });
   });
 

package/src/database/server/models/__tests__/knowledgeBase.test.ts

@@ -1,11 +1,18 @@
 // @vitest-environment node
 import { eq } from 'drizzle-orm';
-import { desc } from 'drizzle-orm/expressions';
+import { and, desc } from 'drizzle-orm/expressions';
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
 
 import { getTestDBInstance } from '@/database/server/core/dbForTest';
 
-import { NewKnowledgeBase, knowledgeBases, users } from '../../schemas/lobechat';
+import {
+  NewKnowledgeBase,
+  files,
+  globalFiles,
+  knowledgeBaseFiles,
+  knowledgeBases,
+  users,
+} from '../../schemas/lobechat';
 import { KnowledgeBaseModel } from '../knowledgeBase';
 
 let serverDB = await getTestDBInstance();
@@ -21,6 +28,7 @@ const knowledgeBaseModel = new KnowledgeBaseModel(userId);
 
 beforeEach(async () => {
   await serverDB.delete(users);
+  await serverDB.delete(globalFiles);
   await serverDB.insert(users).values([{ id: userId }, { id: 'user2' }]);
 });
 
@@ -129,4 +137,122 @@ describe('KnowledgeBaseModel', () => {
       });
     });
   });
+
+  const fileList = [
+    {
+      id: 'file1',
+      name: 'document.pdf',
+      url: 'https://example.com/document.pdf',
+      fileHash: 'hash1',
+      size: 1000,
+      fileType: 'application/pdf',
+      userId,
+    },
+    {
+      id: 'file2',
+      name: 'image.jpg',
+      url: 'https://example.com/image.jpg',
+      fileHash: 'hash2',
+      size: 500,
+      fileType: 'image/jpeg',
+      userId,
+    },
+  ];
+
+  describe('addFilesToKnowledgeBase', () => {
+    it('should add files to a knowledge base', async () => {
+      await serverDB.insert(globalFiles).values([
+        {
+          hashId: 'hash1',
+          url: 'https://example.com/document.pdf',
+          size: 1000,
+          fileType: 'application/pdf',
+        },
+        {
+          hashId: 'hash2',
+          url: 'https://example.com/image.jpg',
+          size: 500,
+          fileType: 'image/jpeg',
+        },
+      ]);
+
+      await serverDB.insert(files).values(fileList);
+
+      const { id: knowledgeBaseId } = await knowledgeBaseModel.create({ name: 'Test Group' });
+      const fileIds = ['file1', 'file2'];
+
+      const result = await knowledgeBaseModel.addFilesToKnowledgeBase(knowledgeBaseId, fileIds);
+
+      expect(result).toHaveLength(2);
+      expect(result).toEqual(
+        expect.arrayContaining(
+          fileIds.map((fileId) => expect.objectContaining({ fileId, knowledgeBaseId })),
+        ),
+      );
+
+      const addedFiles = await serverDB.query.knowledgeBaseFiles.findMany({
+        where: eq(knowledgeBaseFiles.knowledgeBaseId, knowledgeBaseId),
+      });
+      expect(addedFiles).toHaveLength(2);
+    });
+  });
+
+  describe('removeFilesFromKnowledgeBase', () => {
+    it('should remove files from a knowledge base', async () => {
+      await serverDB.insert(globalFiles).values([
+        {
+          hashId: 'hash1',
+          url: 'https://example.com/document.pdf',
+          size: 1000,
+          fileType: 'application/pdf',
+        },
+        {
+          hashId: 'hash2',
+          url: 'https://example.com/image.jpg',
+          size: 500,
+          fileType: 'image/jpeg',
+        },
+      ]);
+
+      await serverDB.insert(files).values(fileList);
+
+      const { id: knowledgeBaseId } = await knowledgeBaseModel.create({ name: 'Test Group' });
+      const fileIds = ['file1', 'file2'];
+      await knowledgeBaseModel.addFilesToKnowledgeBase(knowledgeBaseId, fileIds);
+
+      const filesToRemove = ['file1'];
+      await knowledgeBaseModel.removeFilesFromKnowledgeBase(knowledgeBaseId, filesToRemove);
+
+      const remainingFiles = await serverDB.query.knowledgeBaseFiles.findMany({
+        where: and(eq(knowledgeBaseFiles.knowledgeBaseId, knowledgeBaseId)),
+      });
+      expect(remainingFiles).toHaveLength(1);
+      expect(remainingFiles[0].fileId).toBe('file2');
+    });
+  });
+
+  describe('static findById', () => {
+    it('should find a knowledge base by id without user restriction', async () => {
+      const { id } = await knowledgeBaseModel.create({ name: 'Test Group' });
+
+      const group = await KnowledgeBaseModel.findById(id);
+      expect(group).toMatchObject({
+        id,
+        name: 'Test Group',
+        userId,
+      });
+    });
+
+    it('should find a knowledge base created by another user', async () => {
+      const anotherKnowledgeBaseModel = new KnowledgeBaseModel('user2');
+      const { id } = await anotherKnowledgeBaseModel.create({ name: 'Another User Group' });
+
+      const group = await KnowledgeBaseModel.findById(id);
+      expect(group).toMatchObject({
+        id,
+        name: 'Another User Group',
+        userId: 'user2',
+      });
+    });
+  });
 });

package/src/database/server/models/__tests__/user.test.ts

@@ -4,7 +4,7 @@ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
 import { INBOX_SESSION_ID } from '@/const/session';
 import { getTestDBInstance } from '@/database/server/core/dbForTest';
 import { KeyVaultsGateKeeper } from '@/server/modules/KeyVaultsEncrypt';
-import { UserPreference } from '@/types/user';
+import { UserGuide, UserPreference } from '@/types/user';
 import { UserSettings } from '@/types/user/settings';
 
 import { userSettings, users } from '../../schemas/lobechat';
@@ -167,6 +167,24 @@ describe('UserModel', () => {
       const { plaintext } = await gateKeeper.decrypt(updatedSettings!.keyVaults!);
       expect(JSON.parse(plaintext)).toEqual(settings.keyVaults);
     });
+
+    it('should update user settings with encrypted keyVaults', async () => {
+      const settings = {
+        general: { language: 'en-US' },
+      } as UserSettings;
+      await serverDB.insert(users).values({ id: userId });
+      await serverDB.insert(userSettings).values({ ...settings, keyVaults: '', id: userId });
+
+      const newSettings = {
+        general: { fontSize: 16, language: 'zh-CN', themeMode: 'dark' },
+      } as UserSettings;
+      await userModel.updateSetting(userId, newSettings);
+
+      const updatedSettings = await serverDB.query.userSettings.findFirst({
+        where: eq(users.id, userId),
+      });
+      expect(updatedSettings?.general).toEqual(newSettings.general);
+    });
   });
 
   describe('updatePreference', () => {
@@ -183,4 +201,21 @@ describe('UserModel', () => {
       expect(updatedUser?.preference).toEqual({ ...preference, ...newPreference });
     });
   });
+
+  describe('updateGuide', () => {
+    it('should update user guide', async () => {
+      const preference = { guide: { topic: false } } as UserGuide;
+      await serverDB.insert(users).values({ id: userId, preference });
+
+      const newGuide: Partial<UserGuide> = {
+        topic: true,
+        moveSettingsToAvatar: true,
+        uploadFileInKnowledgeBase: true,
+      };
+      await userModel.updateGuide(userId, newGuide);
+
+      const updatedUser = await serverDB.query.users.findFirst({ where: eq(users.id, userId) });
+      expect(updatedUser?.preference).toEqual({ ...preference, guide: newGuide });
+    });
+  });
 });

package/src/database/server/models/chunk.ts

@@ -1,5 +1,5 @@
 import { asc, cosineDistance, count, eq, inArray, sql } from 'drizzle-orm';
-import { and, desc } from 'drizzle-orm/expressions';
+import { and, desc, isNull } from 'drizzle-orm/expressions';
 
 import { serverDB } from '@/database/server';
 import { ChunkMetadata, FileChunk, SemanticSearchChunk } from '@/types/chunk';
@@ -43,6 +43,19 @@ export class ChunkModel {
     return serverDB.delete(chunks).where(and(eq(chunks.id, id), eq(chunks.userId, this.userId)));
   };
 
+  deleteOrphanChunks = async () => {
+    const orphanedChunks = await serverDB
+      .select({ chunkId: chunks.id })
+      .from(chunks)
+      .leftJoin(fileChunks, eq(chunks.id, fileChunks.chunkId))
+      .where(isNull(fileChunks.fileId));
+
+    const ids = orphanedChunks.map((chunk) => chunk.chunkId);
+    if (ids.length === 0) return;
+
+    await serverDB.delete(chunks).where(inArray(chunks.id, ids));
+  };
+
   findById = async (id: string) => {
     return serverDB.query.chunks.findFirst({
       where: and(eq(chunks.id, id)),

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

@@ -1,6 +1,7 @@
 import { asc, count, eq, ilike, inArray, notExists } from 'drizzle-orm';
 import { and, desc } from 'drizzle-orm/expressions';
 
+import { serverDBEnv } from '@/config/db';
 import { serverDB } from '@/database/server/core/db';
 import { FilesTabs, QueryFileListParams, SortType } from '@/types/files';
 
@@ -77,7 +78,8 @@ export class FileModel {
       const fileCount = result[0].count;
 
       // delete the file from global file if it is not used by other files
-      if (fileCount === 0) {
+      // if `DISABLE_REMOVE_GLOBAL_FILE` is true, we will not remove the global file
+      if (fileCount === 0 && !serverDBEnv.DISABLE_REMOVE_GLOBAL_FILE) {
         await trx.delete(globalFiles).where(eq(globalFiles.hashId, fileHash));
 
         return file;
@@ -118,7 +120,7 @@ export class FileModel {
 
       const needToDeleteList = fileHashCounts.filter((item) => item.count === 0);
 
-      if (needToDeleteList.length === 0) return;
+      if (needToDeleteList.length === 0 || serverDBEnv.DISABLE_REMOVE_GLOBAL_FILE) return;
 
       // delete the file from global file if it is not used by other files
       await trx.delete(globalFiles).where(

package/src/features/FileManager/FileList/FileListItem/DropdownMenu.tsx

@@ -15,7 +15,7 @@ import { useTranslation } from 'react-i18next';
 import { useAddFilesToKnowledgeBaseModal } from '@/features/KnowledgeBaseModal';
 import { useFileStore } from '@/store/file';
 import { useKnowledgeBaseStore } from '@/store/knowledgeBase';
-import { downloadFile } from '@/utils/downloadFile';
+import { downloadFile } from '@/utils/client/downloadFile';
 
 interface DropdownMenuProps {
   filename: string;

package/src/features/FileViewer/NotSupport/index.tsx

@@ -5,7 +5,7 @@ import React, { ComponentType, useState } from 'react';
 import { useTranslation } from 'react-i18next';
 import { Center, Flexbox } from 'react-layout-kit';
 
-import { downloadFile } from '@/utils/downloadFile';
+import { downloadFile } from '@/utils/client/downloadFile';
 
 const useStyles = createStyles(({ css, token }) => ({
   page: css`

package/src/server/routers/lambda/file.ts

@@ -66,7 +66,10 @@ export const fileRouter = router({
       }),
     )
     .query(async ({ ctx, input }) => {
-      return ctx.fileModel.findById(input.id);
+      const item = await ctx.fileModel.findById(input.id);
+      if (!item) throw new TRPCError({ code: 'BAD_REQUEST', message: 'File not found' });
+
+      return { ...item, url: getFullFileUrl(item?.url) };
     }),
 
   getFileItemById: fileProcedure
@@ -147,6 +150,8 @@ export const fileRouter = router({
   removeFile: fileProcedure.input(z.object({ id: z.string() })).mutation(async ({ input, ctx }) => {
     const file = await ctx.fileModel.delete(input.id);
 
+    // delete the orphan chunks
+    await ctx.chunkModel.deleteOrphanChunks();
     if (!file) return;
 
     // delele the file from remove from S3 if it is not used by other files
@@ -159,6 +164,9 @@ export const fileRouter = router({
     .mutation(async ({ input, ctx }) => {
       const needToRemoveFileList = await ctx.fileModel.deleteMany(input.ids);
 
+      // delete the orphan chunks
+      await ctx.chunkModel.deleteOrphanChunks();
+
       if (!needToRemoveFileList || needToRemoveFileList.length === 0) return;
 
       // remove from S3