@lobehub/chat 1.84.27 → 1.85.1

package/src/server/routers/lambda/__tests__/message.test.ts

@@ -0,0 +1,213 @@
+import { describe, expect, it, vi } from 'vitest';
+
+import { MessageModel } from '@/database/models/message';
+import { FileService } from '@/server/services/file';
+import { ChatMessage, CreateMessageParams } from '@/types/message';
+
+vi.mock('@/database/models/message', () => ({
+  MessageModel: vi.fn(),
+}));
+
+vi.mock('@/server/services/file', () => ({
+  FileService: vi.fn(),
+}));
+
+vi.mock('@/database/server', () => ({
+  getServerDB: vi.fn(),
+}));
+
+describe('messageRouter', () => {
+  it('should handle batchCreateMessages', async () => {
+    const mockBatchCreate = vi.fn().mockResolvedValue({ rowCount: 2 });
+    vi.mocked(MessageModel).mockImplementation(
+      () =>
+        ({
+          batchCreate: mockBatchCreate,
+        }) as any,
+    );
+
+    const input = [
+      {
+        id: '1',
+        role: 'user',
+        content: 'test',
+        sessionId: 'session1',
+        createdAt: new Date(),
+        updatedAt: new Date(),
+        agentId: 'agent1',
+        clientId: 'client1',
+        parentId: null,
+        quotaId: null,
+        model: null,
+        provider: null,
+        topicId: null,
+        error: null,
+        favorite: false,
+        observationId: null,
+        reasoning: null,
+        pluginState: null,
+        translate: null,
+        tts: null,
+        search: null,
+        threadId: null,
+        tools: null,
+        traceId: null,
+        userId: 'user1',
+      } as any,
+    ];
+
+    const ctx = {
+      messageModel: new MessageModel({} as any, 'user1'),
+    };
+
+    const result = await ctx.messageModel.batchCreate(input);
+
+    expect(mockBatchCreate).toHaveBeenCalledWith(input);
+    expect(result.rowCount).toBe(2);
+  });
+
+  it('should handle count', async () => {
+    const mockCount = vi.fn().mockResolvedValue(5);
+    vi.mocked(MessageModel).mockImplementation(
+      () =>
+        ({
+          count: mockCount,
+        }) as any,
+    );
+
+    const input = { startDate: '2024-01-01' };
+    const ctx = {
+      messageModel: new MessageModel({} as any, 'user1'),
+    };
+
+    const result = await ctx.messageModel.count(input);
+
+    expect(mockCount).toHaveBeenCalledWith(input);
+    expect(result).toBe(5);
+  });
+
+  it('should handle createMessage', async () => {
+    const mockCreate = vi.fn().mockResolvedValue({ id: 'msg1' });
+    vi.mocked(MessageModel).mockImplementation(
+      () =>
+        ({
+          create: mockCreate,
+        }) as any,
+    );
+
+    const input: CreateMessageParams = {
+      content: 'test',
+      role: 'user',
+      sessionId: 'session1',
+    };
+
+    const ctx = {
+      messageModel: new MessageModel({} as any, 'user1'),
+    };
+
+    const result = await ctx.messageModel.create(input);
+
+    expect(mockCreate).toHaveBeenCalledWith(input);
+    expect(result.id).toBe('msg1');
+  });
+
+  it('should handle getMessages', async () => {
+    const mockQuery = vi.fn().mockResolvedValue([{ id: 'msg1' }]);
+    const mockGetFullFileUrl = vi
+      .fn()
+      .mockImplementation((path: string | null, file: { fileType: string }) => {
+        return Promise.resolve('url');
+      });
+
+    vi.mocked(MessageModel).mockImplementation(
+      () =>
+        ({
+          query: mockQuery,
+        }) as any,
+    );
+
+    vi.mocked(FileService).mockImplementation(
+      () =>
+        ({
+          getFullFileUrl: mockGetFullFileUrl,
+        }) as any,
+    );
+
+    const input = { sessionId: 'session1' };
+    const ctx = {
+      messageModel: new MessageModel({} as any, 'user1'),
+      fileService: new FileService({} as any, 'user1'),
+      userId: 'user1',
+    };
+
+    const result = await ctx.messageModel.query(input, {
+      postProcessUrl: mockGetFullFileUrl,
+    });
+
+    expect(mockQuery).toHaveBeenCalledWith(input, expect.any(Object));
+    expect(result).toEqual([{ id: 'msg1' }]);
+  });
+
+  it('should handle getAllMessages', async () => {
+    const mockQueryAll = vi.fn().mockResolvedValue([
+      {
+        id: 'msg1',
+        meta: {},
+      } as ChatMessage,
+    ]);
+    vi.mocked(MessageModel).mockImplementation(
+      () =>
+        ({
+          queryAll: mockQueryAll,
+        }) as any,
+    );
+
+    const ctx = {
+      messageModel: new MessageModel({} as any, 'user1'),
+    };
+
+    const result = await ctx.messageModel.queryAll();
+
+    expect(mockQueryAll).toHaveBeenCalled();
+    expect(result).toEqual([{ id: 'msg1', meta: {} }]);
+  });
+
+  it('should handle removeMessage', async () => {
+    const mockDelete = vi.fn().mockResolvedValue(undefined);
+    vi.mocked(MessageModel).mockImplementation(
+      () =>
+        ({
+          deleteMessage: mockDelete,
+        }) as any,
+    );
+
+    const input = { id: 'msg1' };
+    const ctx = {
+      messageModel: new MessageModel({} as any, 'user1'),
+    };
+
+    await ctx.messageModel.deleteMessage(input.id);
+
+    expect(mockDelete).toHaveBeenCalledWith(input.id);
+  });
+
+  it('should handle updateMessage', async () => {
+    const mockUpdate = vi.fn().mockResolvedValue({ success: true });
+    vi.mocked(MessageModel).mockImplementation(
+      () =>
+        ({
+          update: mockUpdate,
+        }) as any,
+    );
+
+    const input = { id: 'msg1', value: { content: 'updated' } };
+    const ctx = {
+      messageModel: new MessageModel({} as any, 'user1'),
+    };
+
+    const result = await ctx.messageModel.update(input.id, input.value);
+
+    expect(mockUpdate).toHaveBeenCalledWith(input.id, input.value);
+    expect(result).toEqual({ success: true });
+  });
+});

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

@@ -0,0 +1,36 @@
+import { z } from 'zod';
+
+import { ChunkModel } from '@/database/models/chunk';
+import { FileModel } from '@/database/models/file';
+import { MessageModel } from '@/database/models/message';
+import { authedProcedure, router } from '@/libs/trpc/lambda';
+import { serverDatabase } from '@/libs/trpc/lambda/middleware';
+import { DocumentService } from '@/server/services/document';
+
+const documentProcedure = authedProcedure.use(serverDatabase).use(async (opts) => {
+  const { ctx } = opts;
+
+  return opts.next({
+    ctx: {
+      chunkModel: new ChunkModel(ctx.serverDB, ctx.userId),
+      documentService: new DocumentService(ctx.serverDB, ctx.userId),
+      fileModel: new FileModel(ctx.serverDB, ctx.userId),
+      messageModel: new MessageModel(ctx.serverDB, ctx.userId),
+    },
+  });
+});
+
+export const documentRouter = router({
+  parseFileContent: documentProcedure
+    .input(
+      z.object({
+        id: z.string(),
+        skipExist: z.boolean().optional(),
+      }),
+    )
+    .mutation(async ({ ctx, input }) => {
+      const lobeDocument = await ctx.documentService.parseFile(input.id);
+
+      return lobeDocument;
+    }),
+});

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

@@ -7,6 +7,7 @@ import { agentRouter } from './agent';
 import { aiModelRouter } from './aiModel';
 import { aiProviderRouter } from './aiProvider';
 import { chunkRouter } from './chunk';
+import { documentRouter } from './document';
 import { exporterRouter } from './exporter';
 import { fileRouter } from './file';
 import { importerRouter } from './importer';
@@ -25,6 +26,7 @@ export const lambdaRouter = router({
   aiModel: aiModelRouter,
   aiProvider: aiProviderRouter,
   chunk: chunkRouter,
+  document: documentRouter,
   exporter: exporterRouter,
   file: fileRouter,
   healthcheck: publicProcedure.query(() => "i'm live!"),

package/src/server/services/document/index.ts

@@ -0,0 +1,66 @@
+import { loadFile } from '@lobechat/file-loaders';
+import debug from 'debug';
+
+import { DocumentModel } from '@/database/models/document';
+import { FileModel } from '@/database/models/file';
+import { LobeChatDatabase } from '@/database/type';
+import { LobeDocument } from '@/types/document';
+
+import { FileService } from '../file';
+
+const log = debug('lobe-chat:service:document');
+
+export class DocumentService {
+  userId: string;
+  private fileModel: FileModel;
+  private documentModel: DocumentModel;
+  private fileService: FileService;
+
+  constructor(db: LobeChatDatabase, userId: string) {
+    this.userId = userId;
+    this.fileModel = new FileModel(db, userId);
+    this.fileService = new FileService(db, userId);
+    this.documentModel = new DocumentModel(db, userId);
+  }
+
+  /**
+   * 解析文件内容
+   *
+   */
+  async parseFile(fileId: string): Promise<LobeDocument> {
+    const { filePath, file, cleanup } = await this.fileService.downloadFileToLocal(fileId);
+
+    const logPrefix = `[${file.name}]`;
+    log(`${logPrefix} 开始解析文件, 路径: ${filePath}`);
+
+    try {
+      // 使用loadFile加载文件内容
+      const fileDocument = await loadFile(filePath);
+
+      log(`${logPrefix} 文件解析成功 %O`, {
+        fileType: fileDocument.fileType,
+        size: fileDocument.content.length,
+      });
+
+      const document = await this.documentModel.create({
+        content: fileDocument.content,
+        fileId,
+        fileType: file.fileType,
+        metadata: fileDocument.metadata,
+        pages: fileDocument.pages,
+        source: file.url,
+        sourceType: 'file',
+        title: fileDocument.metadata?.title,
+        totalCharCount: fileDocument.totalCharCount,
+        totalLineCount: fileDocument.totalLineCount,
+      });
+
+      return document as LobeDocument;
+    } catch (error) {
+      console.error(`${logPrefix} 文件解析失败:`, error);
+      throw error;
+    } finally {
+      cleanup();
+    }
+  }
+}

package/src/server/services/file/__tests__/index.test.ts

@@ -0,0 +1,115 @@
+import { TRPCError } from '@trpc/server';
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { FileModel } from '@/database/models/file';
+import { TempFileManager } from '@/server/utils/tempFileManager';
+
+import { FileService } from '../index';
+
+vi.mock('@/config/db', () => ({
+  serverDBEnv: {
+    REMOVE_GLOBAL_FILE: false,
+  },
+}));
+
+vi.mock('../impls', () => ({
+  createFileServiceModule: () => ({
+    deleteFile: vi.fn(),
+    deleteFiles: vi.fn(),
+    getFileContent: vi.fn(),
+    getFileByteArray: vi.fn(),
+    createPreSignedUrl: vi.fn(),
+    createPreSignedUrlForPreview: vi.fn(),
+    uploadContent: vi.fn(),
+    getFullFileUrl: vi.fn(),
+  }),
+}));
+
+vi.mock('@/database/models/file');
+
+vi.mock('@/server/utils/tempFileManager');
+
+vi.mock('@/utils/uuid', () => ({
+  nanoid: () => 'test-id',
+}));
+
+describe('FileService', () => {
+  let service: FileService;
+  const mockDb = {} as any;
+  const mockUserId = 'test-user';
+  let mockFileModel: any;
+  let mockTempManager: any;
+
+  beforeEach(() => {
+    mockFileModel = {
+      findById: vi.fn(),
+      delete: vi.fn(),
+    };
+    mockTempManager = {
+      writeTempFile: vi.fn(),
+      cleanup: vi.fn(),
+    };
+    vi.mocked(FileModel).mockImplementation(() => mockFileModel);
+    vi.mocked(TempFileManager).mockImplementation(() => mockTempManager);
+    service = new FileService(mockDb, mockUserId);
+  });
+
+  afterEach(() => {
+    vi.clearAllMocks();
+  });
+
+  describe('downloadFileToLocal', () => {
+    const mockFile = {
+      id: 'test-file-id',
+      name: 'test.txt',
+      url: 'test-url',
+    };
+
+    it('should throw error if file not found', async () => {
+      mockFileModel.findById.mockResolvedValue(undefined);
+
+      await expect(service.downloadFileToLocal('test-file-id')).rejects.toThrow(
+        new TRPCError({ code: 'BAD_REQUEST', message: 'File not found' }),
+      );
+    });
+
+    it('should throw error if file content is empty', async () => {
+      mockFileModel.findById.mockResolvedValue(mockFile);
+      vi.mocked(service['impl'].getFileByteArray).mockResolvedValue(undefined as any);
+
+      await expect(service.downloadFileToLocal('test-file-id')).rejects.toThrow(
+        new TRPCError({ code: 'BAD_REQUEST', message: 'File content is empty' }),
+      );
+    });
+
+    it('should delete file from db and throw error if file not found in storage', async () => {
+      mockFileModel.findById.mockResolvedValue(mockFile);
+      vi.mocked(service['impl'].getFileByteArray).mockRejectedValue({ Code: 'NoSuchKey' });
+
+      await expect(service.downloadFileToLocal('test-file-id')).rejects.toThrow(
+        new TRPCError({ code: 'BAD_REQUEST', message: 'File not found' }),
+      );
+
+      expect(mockFileModel.delete).toHaveBeenCalledWith('test-file-id', false);
+    });
+
+    it('should successfully download file to local', async () => {
+      const mockContent = new Uint8Array([1, 2, 3]);
+      const mockFilePath = '/tmp/test.txt';
+
+      mockFileModel.findById.mockResolvedValue(mockFile);
+      vi.mocked(service['impl'].getFileByteArray).mockResolvedValue(mockContent);
+      mockTempManager.writeTempFile.mockResolvedValue(mockFilePath);
+
+      const result = await service.downloadFileToLocal('test-file-id');
+
+      expect(result).toEqual({
+        cleanup: expect.any(Function),
+        file: mockFile,
+        filePath: mockFilePath,
+      });
+
+      expect(mockTempManager.writeTempFile).toHaveBeenCalledWith(mockContent, mockFile.name);
+    });
+  });
+});

package/src/server/services/mcp/index.ts

@@ -15,10 +15,6 @@ class MCPService {
   // Store instances of the custom MCPClient, keyed by serialized MCPClientParams
   private clients: Map<string, MCPClient> = new Map();
 
-  constructor() {
-    log('MCPService initialized');
-  }
-
   // --- MCP Interaction ---
 
   // listTools now accepts MCPClientParams

package/src/server/utils/__tests__/tempFileManager.test.ts

@@ -0,0 +1,94 @@
+import { existsSync, mkdtempSync, rmSync, writeFileSync } from 'node:fs';
+import { tmpdir } from 'node:os';
+import { join } from 'node:path';
+import { describe, expect, it, vi } from 'vitest';
+
+import { TempFileManager } from '../tempFileManager';
+
+// Mock node modules
+vi.mock('node:fs');
+vi.mock('node:os');
+vi.mock('node:path', () => ({
+  join: (...args: string[]) => args.join('/'),
+  default: {
+    join: (...args: string[]) => args.join('/'),
+  },
+}));
+
+describe('TempFileManager', () => {
+  const mockTmpDir = '/tmp';
+  const mockDirname = 'test-';
+  const mockFullTmpDir = '/tmp/test-xyz';
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+    vi.mocked(tmpdir).mockReturnValue(mockTmpDir);
+    vi.mocked(mkdtempSync).mockReturnValue(mockFullTmpDir);
+    vi.mocked(existsSync).mockReturnValue(true);
+  });
+
+  it('should create temp directory on initialization', () => {
+    new TempFileManager(mockDirname);
+
+    expect(tmpdir).toHaveBeenCalled();
+    expect(mkdtempSync).toHaveBeenCalledWith(`${mockTmpDir}/${mockDirname}`);
+  });
+
+  it('should write temp file successfully', async () => {
+    const manager = new TempFileManager(mockDirname);
+    const testData = new Uint8Array([1, 2, 3]);
+    const fileName = 'test.txt';
+
+    const filePath = await manager.writeTempFile(testData, fileName);
+
+    expect(writeFileSync).toHaveBeenCalledWith(`${mockFullTmpDir}/${fileName}`, testData);
+    expect(filePath).toBe(`${mockFullTmpDir}/${fileName}`);
+  });
+
+  it('should cleanup on write failure', async () => {
+    const manager = new TempFileManager(mockDirname);
+    const testData = new Uint8Array([1, 2, 3]);
+    const fileName = 'test.txt';
+
+    vi.mocked(writeFileSync).mockImplementation(() => {
+      throw new Error('Write failed');
+    });
+
+    await expect(manager.writeTempFile(testData, fileName)).rejects.toThrow(
+      'Failed to write temp file: Write failed',
+    );
+
+    expect(existsSync).toHaveBeenCalledWith(mockFullTmpDir);
+    expect(rmSync).toHaveBeenCalledWith(mockFullTmpDir, { force: true, recursive: true });
+  });
+
+  it('should cleanup temp directory', () => {
+    const manager = new TempFileManager(mockDirname);
+    vi.mocked(existsSync).mockReturnValue(true);
+
+    manager.cleanup();
+
+    expect(existsSync).toHaveBeenCalledWith(mockFullTmpDir);
+    expect(rmSync).toHaveBeenCalledWith(mockFullTmpDir, { force: true, recursive: true });
+  });
+
+  it('should skip cleanup if directory does not exist', () => {
+    const manager = new TempFileManager(mockDirname);
+    vi.mocked(existsSync).mockReturnValue(false);
+
+    manager.cleanup();
+
+    expect(existsSync).toHaveBeenCalledWith(mockFullTmpDir);
+    expect(rmSync).not.toHaveBeenCalled();
+  });
+
+  it('should register cleanup hooks on process events', () => {
+    const processOnSpy = vi.spyOn(process, 'on');
+    new TempFileManager(mockDirname);
+
+    expect(processOnSpy).toHaveBeenCalledWith('exit', expect.any(Function));
+    expect(processOnSpy).toHaveBeenCalledWith('uncaughtException', expect.any(Function));
+    expect(processOnSpy).toHaveBeenCalledWith('SIGINT', expect.any(Function));
+    expect(processOnSpy).toHaveBeenCalledWith('SIGTERM', expect.any(Function));
+  });
+});

package/src/services/rag.ts

@@ -2,6 +2,10 @@ import { lambdaClient } from '@/libs/trpc/client';
 import { SemanticSearchSchemaType } from '@/types/rag';
 
 class RAGService {
+  parseFileContent = async (id: string, skipExist?: boolean) => {
+    return lambdaClient.document.parseFileContent.mutate({ id, skipExist });
+  };
+
   createParseFileTask = async (id: string, skipExist?: boolean) => {
     return lambdaClient.chunk.createParseFileTask.mutate({ id, skipExist });
   };

package/src/store/chat/slices/aiChat/actions/__tests__/rag.test.ts

@@ -197,13 +197,13 @@ describe('chatRAG actions', () => {
       expect(result.current.internal_shouldUseRAG()).toBe(true);
     });
 
-    it('should return true if has user files', () => {
+    it('should return false if has user files', () => {
       const { result } = renderHook(() => useChatStore());
 
       vi.spyOn(agentSelectors, 'hasEnabledKnowledge').mockReturnValue(false);
       vi.spyOn(chatSelectors, 'currentUserFiles').mockReturnValue([{ id: 'file-1' }] as any);
 
-      expect(result.current.internal_shouldUseRAG()).toBe(true);
+      expect(result.current.internal_shouldUseRAG()).toBeFalsy();
     });
 
     it('should return false if no knowledge or files', () => {

package/src/store/chat/slices/aiChat/actions/rag.ts

@@ -130,9 +130,8 @@ export const chatRag: StateCreator<ChatStore, [['zustand/devtools', never]], [],
     return rewriteQuery;
   },
   internal_shouldUseRAG: () => {
-    const userFiles = chatSelectors.currentUserFiles(get()).map((f) => f.id);
-    //  if there is relative files or enabled knowledge, try with ragQuery
-    return hasEnabledKnowledge() || userFiles.length > 0;
+    //  if there is enabled knowledge, try with ragQuery
+    return hasEnabledKnowledge();
   },
 
   internal_toggleMessageRAGLoading: (loading, id) => {

package/src/store/file/slices/chat/action.ts

@@ -7,14 +7,10 @@ import { fileService } from '@/services/file';
 import { ServerService } from '@/services/file/server';
 import { ragService } from '@/services/rag';
 import { UPLOAD_NETWORK_ERROR } from '@/services/upload';
-import { userService } from '@/services/user';
-import { useAgentStore } from '@/store/agent';
 import {
   UploadFileListDispatch,
   uploadFileListReducer,
 } from '@/store/file/reducers/uploadFileList';
-import { useUserStore } from '@/store/user';
-import { preferenceSelectors } from '@/store/user/selectors';
 import { FileListItem } from '@/types/files';
 import { UploadFileItem } from '@/types/files/upload';
 import { isChunkingUnsupported } from '@/utils/isChunkingUnsupported';
@@ -97,7 +93,7 @@ export const createFileSlice: StateCreator<
   },
 
   uploadChatFiles: async (rawFiles) => {
-    const { dispatchChatUploadFileList, startAsyncTask } = get();
+    const { dispatchChatUploadFileList } = get();
     // 0. skip file in blacklist
     const files = rawFiles.filter((file) => !FILE_UPLOAD_BLACKLIST.includes(file.name));
     // 1. add files with base64
@@ -154,52 +150,8 @@ export const createFileSlice: StateCreator<
       // image don't need to be chunked and embedding
       if (isChunkingUnsupported(file.type)) return;
 
-      // 3. auto chunk and embedding
-      dispatchChatUploadFileList({
-        id: fileResult.id,
-        type: 'updateFile',
-        // make the taks empty to hint the user that the task is starting but not triggered
-        value: { tasks: {} },
-      });
-
-      await startAsyncTask(
-        fileResult.id,
-        async (id) => {
-          const data = await ragService.createParseFileTask(id);
-          if (!data || !data.id) throw new Error('failed to createParseFileTask');
-
-          // run the assignment
-          useAgentStore
-            .getState()
-            .addFilesToAgent([id], false)
-            .then(() => {
-              // trigger the tip if it's the first time
-              if (!preferenceSelectors.shouldTriggerFileInKnowledgeBaseTip(useUserStore.getState()))
-                return;
-
-              userService.updateGuide({ uploadFileInKnowledgeBase: true });
-            });
-
-          return data.id;
-        },
-
-        (fileItem) => {
-          dispatchChatUploadFileList({
-            id: fileResult.id,
-            type: 'updateFile',
-            value: {
-              tasks: {
-                chunkCount: fileItem.chunkCount,
-                chunkingError: fileItem.chunkingError,
-                chunkingStatus: fileItem.chunkingStatus,
-                embeddingError: fileItem.embeddingError,
-                embeddingStatus: fileItem.embeddingStatus,
-                finishEmbedding: fileItem.finishEmbedding,
-              },
-            },
-          });
-        },
-      );
+      const data = await ragService.parseFileContent(fileResult.id);
+      console.log(data);
     });
 
     await Promise.all(pools);