@lobehub/lobehub 2.0.0-next.45 → 2.0.0-next.47

package/src/server/modules/ContentChunk/index.test.ts

@@ -0,0 +1,372 @@
+import { beforeEach, describe, expect, it, vi } from 'vitest';
+
+import type { NewChunkItem, NewUnstructuredChunkItem } from '@/database/schemas';
+import { knowledgeEnv } from '@/envs/knowledge';
+import { ChunkingLoader } from '@/libs/langchain';
+import { ChunkingStrategy, Unstructured } from '@/libs/unstructured';
+
+import { ContentChunk } from './index';
+
+// Mock the dependencies
+vi.mock('@/libs/unstructured');
+vi.mock('@/libs/langchain');
+vi.mock('@/envs/knowledge', () => ({
+  knowledgeEnv: {
+    FILE_TYPE_CHUNKING_RULES: '',
+    UNSTRUCTURED_API_KEY: 'test-api-key',
+    UNSTRUCTURED_SERVER_URL: 'https://test.unstructured.io',
+  },
+}));
+
+describe('ContentChunk', () => {
+  let contentChunk: ContentChunk;
+  let mockUnstructuredPartition: ReturnType<typeof vi.fn>;
+  let mockLangChainPartition: ReturnType<typeof vi.fn>;
+
+  beforeEach(() => {
+    vi.clearAllMocks();
+
+    // Setup Unstructured mock
+    mockUnstructuredPartition = vi.fn();
+    (Unstructured as unknown as ReturnType<typeof vi.fn>).mockImplementation(() => ({
+      partition: mockUnstructuredPartition,
+    }));
+
+    // Setup LangChain mock
+    mockLangChainPartition = vi.fn();
+    (ChunkingLoader as unknown as ReturnType<typeof vi.fn>).mockImplementation(() => ({
+      partitionContent: mockLangChainPartition,
+    }));
+
+    contentChunk = new ContentChunk();
+  });
+
+  describe('constructor', () => {
+    it('should initialize with Unstructured and LangChain clients', () => {
+      expect(Unstructured).toHaveBeenCalledTimes(1);
+      expect(ChunkingLoader).toHaveBeenCalledTimes(1);
+    });
+  });
+
+  describe('chunkContent', () => {
+    const mockFileContent = new Uint8Array([1, 2, 3, 4, 5]);
+    const mockFilename = 'test-document.pdf';
+
+    it('should use default langchain service when no rules are configured', async () => {
+      const mockLangChainResult = [
+        {
+          id: 'chunk-1',
+          metadata: { source: 'test' },
+          pageContent: 'Test content chunk 1',
+        },
+        {
+          id: 'chunk-2',
+          metadata: { source: 'test' },
+          pageContent: 'Test content chunk 2',
+        },
+      ];
+
+      mockLangChainPartition.mockResolvedValue(mockLangChainResult);
+
+      const result = await contentChunk.chunkContent({
+        content: mockFileContent,
+        fileType: 'application/pdf',
+        filename: mockFilename,
+      });
+
+      expect(mockLangChainPartition).toHaveBeenCalledWith(mockFilename, mockFileContent);
+      expect(result.chunks).toHaveLength(2);
+      expect(result.chunks[0]).toMatchObject({
+        id: 'chunk-1',
+        index: 0,
+        metadata: { source: 'test' },
+        text: 'Test content chunk 1',
+        type: 'LangChainElement',
+      });
+      expect(result.unstructuredChunks).toBeUndefined();
+    });
+
+    it('should use langchain when unstructured is not configured', async () => {
+      // Temporarily mock env to disable unstructured
+      vi.mocked(knowledgeEnv).UNSTRUCTURED_API_KEY = '';
+
+      const mockLangChainResult = [
+        {
+          id: 'chunk-1',
+          metadata: { source: 'test' },
+          pageContent: 'LangChain content',
+        },
+      ];
+
+      mockLangChainPartition.mockResolvedValue(mockLangChainResult);
+
+      const result = await contentChunk.chunkContent({
+        content: mockFileContent,
+        fileType: 'application/pdf',
+        filename: mockFilename,
+      });
+
+      expect(mockLangChainPartition).toHaveBeenCalledWith(mockFilename, mockFileContent);
+      expect(result.chunks).toHaveLength(1);
+      expect(result.chunks[0].text).toBe('LangChain content');
+      expect(result.unstructuredChunks).toBeUndefined();
+
+      // Restore mock
+      vi.mocked(knowledgeEnv).UNSTRUCTURED_API_KEY = 'test-api-key';
+    });
+
+    it('should handle langchain results with metadata', async () => {
+      const mockLangChainResult = [
+        {
+          id: 'chunk-1',
+          metadata: {
+            source: 'test-document.pdf',
+            page: 1,
+            loc: { lines: { from: 1, to: 10 } },
+          },
+          pageContent: 'First paragraph content',
+        },
+        {
+          id: 'chunk-2',
+          metadata: {
+            source: 'test-document.pdf',
+            page: 2,
+          },
+          pageContent: 'Second paragraph content',
+        },
+      ];
+
+      mockLangChainPartition.mockResolvedValue(mockLangChainResult);
+
+      const result = await contentChunk.chunkContent({
+        content: mockFileContent,
+        fileType: 'application/pdf',
+        filename: mockFilename,
+      });
+
+      expect(result.chunks).toHaveLength(2);
+      expect(result.chunks[0]).toMatchObject({
+        id: 'chunk-1',
+        index: 0,
+        metadata: {
+          source: 'test-document.pdf',
+          page: 1,
+          loc: { lines: { from: 1, to: 10 } },
+        },
+        text: 'First paragraph content',
+        type: 'LangChainElement',
+      });
+      expect(result.chunks[1]).toMatchObject({
+        id: 'chunk-2',
+        index: 1,
+        text: 'Second paragraph content',
+        type: 'LangChainElement',
+      });
+    });
+
+    it('should handle different file types', async () => {
+      const mockLangChainResult = [
+        {
+          id: 'docx-chunk-1',
+          metadata: { source: 'test.docx' },
+          pageContent: 'Word document content',
+        },
+      ];
+
+      mockLangChainPartition.mockResolvedValue(mockLangChainResult);
+
+      const result = await contentChunk.chunkContent({
+        content: mockFileContent,
+        fileType: 'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
+        filename: 'test.docx',
+      });
+
+      expect(mockLangChainPartition).toHaveBeenCalledWith('test.docx', mockFileContent);
+      expect(result.chunks).toHaveLength(1);
+      expect(result.chunks[0].text).toBe('Word document content');
+    });
+
+    it('should throw error when all services fail and its the last service', async () => {
+      mockLangChainPartition.mockRejectedValue(new Error('LangChain error'));
+
+      await expect(
+        contentChunk.chunkContent({
+          content: mockFileContent,
+          fileType: 'application/pdf',
+          filename: mockFilename,
+        }),
+      ).rejects.toThrow('LangChain error');
+    });
+
+    it('should handle empty langchain results', async () => {
+      mockLangChainPartition.mockResolvedValue([]);
+
+      const result = await contentChunk.chunkContent({
+        content: mockFileContent,
+        fileType: 'application/pdf',
+        filename: mockFilename,
+      });
+
+      expect(result.chunks).toHaveLength(0);
+      expect(result.unstructuredChunks).toBeUndefined();
+    });
+
+    it('should extract file extension correctly from MIME type', async () => {
+      const mockLangChainResult = [
+        {
+          id: 'chunk-1',
+          metadata: {},
+          pageContent: 'Content',
+        },
+      ];
+
+      mockLangChainPartition.mockResolvedValue(mockLangChainResult);
+
+      await contentChunk.chunkContent({
+        content: mockFileContent,
+        fileType: 'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
+        filename: 'test.docx',
+      });
+
+      expect(mockLangChainPartition).toHaveBeenCalledWith('test.docx', mockFileContent);
+    });
+
+    it('should handle langchain results with minimal metadata', async () => {
+      const mockLangChainResult = [
+        {
+          id: 'chunk-minimal',
+          metadata: {},
+          pageContent: 'Content with no metadata',
+        },
+      ];
+
+      mockLangChainPartition.mockResolvedValue(mockLangChainResult);
+
+      const result = await contentChunk.chunkContent({
+        content: mockFileContent,
+        fileType: 'text/plain',
+        filename: 'test.txt',
+      });
+
+      expect(result.chunks[0]).toMatchObject({
+        id: 'chunk-minimal',
+        index: 0,
+        metadata: {},
+        text: 'Content with no metadata',
+        type: 'LangChainElement',
+      });
+    });
+  });
+
+  describe('canUseUnstructured', () => {
+    it('should return true when API key and server URL are configured', () => {
+      const result = contentChunk['canUseUnstructured']();
+      expect(result).toBe(true);
+    });
+
+    it('should return false when API key is missing', () => {
+      const originalKey = knowledgeEnv.UNSTRUCTURED_API_KEY;
+      vi.mocked(knowledgeEnv).UNSTRUCTURED_API_KEY = '';
+
+      const result = contentChunk['canUseUnstructured']();
+      expect(result).toBe(false);
+
+      // Restore
+      vi.mocked(knowledgeEnv).UNSTRUCTURED_API_KEY = originalKey;
+    });
+
+    it('should return false when server URL is missing', () => {
+      const originalUrl = knowledgeEnv.UNSTRUCTURED_SERVER_URL;
+      vi.mocked(knowledgeEnv).UNSTRUCTURED_SERVER_URL = '';
+
+      const result = contentChunk['canUseUnstructured']();
+      expect(result).toBe(false);
+
+      // Restore
+      vi.mocked(knowledgeEnv).UNSTRUCTURED_SERVER_URL = originalUrl;
+    });
+  });
+
+  describe('getChunkingServices', () => {
+    it('should return default service for unknown file type', () => {
+      const services = contentChunk['getChunkingServices']('application/unknown');
+      expect(services).toEqual(['default']);
+    });
+
+    it('should extract extension from MIME type correctly', () => {
+      const services = contentChunk['getChunkingServices']('application/pdf');
+      expect(services).toEqual(['default']);
+    });
+
+    it('should handle MIME types with multiple slashes', () => {
+      const services = contentChunk['getChunkingServices'](
+        'application/vnd.openxmlformats-officedocument/wordprocessingml.document',
+      );
+      expect(services).toEqual(['default']);
+    });
+
+    it('should convert extension to lowercase', () => {
+      const services = contentChunk['getChunkingServices']('application/PDF');
+      expect(services).toEqual(['default']);
+    });
+  });
+
+  describe('integration scenarios', () => {
+    it('should handle multiple chunk items with correct indices', async () => {
+      const mockLangChainResult = Array.from({ length: 5 }, (_, i) => ({
+        id: `chunk-${i}`,
+        metadata: { index: i },
+        pageContent: `Content ${i}`,
+      }));
+
+      mockLangChainPartition.mockResolvedValue(mockLangChainResult);
+
+      const result = await contentChunk.chunkContent({
+        content: new Uint8Array([1, 2, 3]),
+        fileType: 'text/plain',
+        filename: 'test.txt',
+      });
+
+      expect(result.chunks).toHaveLength(5);
+      result.chunks.forEach((chunk, index) => {
+        expect(chunk.index).toBe(index);
+        expect(chunk.text).toBe(`Content ${index}`);
+      });
+    });
+
+    it('should preserve order of chunks from langchain response', async () => {
+      const mockLangChainResult = [
+        {
+          id: 'elem-3',
+          metadata: { source: 'test.txt' },
+          pageContent: 'Third',
+        },
+        {
+          id: 'elem-1',
+          metadata: { source: 'test.txt' },
+          pageContent: 'First',
+        },
+        {
+          id: 'elem-2',
+          metadata: { source: 'test.txt' },
+          pageContent: 'Second',
+        },
+      ];
+
+      mockLangChainPartition.mockResolvedValue(mockLangChainResult);
+
+      const result = await contentChunk.chunkContent({
+        content: new Uint8Array([1, 2, 3]),
+        fileType: 'text/plain',
+        filename: 'test.txt',
+      });
+
+      expect(result.chunks[0].text).toBe('Third');
+      expect(result.chunks[1].text).toBe('First');
+      expect(result.chunks[2].text).toBe('Second');
+      expect(result.chunks[0].index).toBe(0);
+      expect(result.chunks[1].index).toBe(1);
+      expect(result.chunks[2].index).toBe(2);
+    });
+  });
+});

package/src/server/routers/desktop/mcp.ts

@@ -2,8 +2,12 @@ import { GetStreamableMcpServerManifestInputSchema } from '@lobechat/types';
 import debug from 'debug';
 import { z } from 'zod';
 
+import { ToolCallContent } from '@/libs/mcp';
 import { authedProcedure, router } from '@/libs/trpc/lambda';
+import { serverDatabase } from '@/libs/trpc/lambda/middleware';
+import { FileService } from '@/server/services/file';
 import { mcpService } from '@/server/services/mcp';
+import { processContentBlocks } from '@/server/services/mcp/contentProcessor';
 
 const log = debug('lobe-mcp:router');
 
@@ -22,7 +26,13 @@ const stdioParamsSchema = z.object({
   type: z.literal('stdio').default('stdio'),
 });
 
-const mcpProcedure = authedProcedure;
+const mcpProcedure = authedProcedure.use(serverDatabase).use(async ({ ctx, next }) => {
+  return next({
+    ctx: {
+      fileService: new FileService(ctx.serverDB, ctx.userId),
+    },
+  });
+});
 
 export const mcpRouter = router({
   getStdioMcpServerManifest: mcpProcedure.input(stdioParamsSchema).query(async ({ input }) => {
@@ -85,13 +95,18 @@ export const mcpRouter = router({
         toolName: z.string(),
       }),
     )
-    .mutation(async ({ input }) => {
-      // Pass the validated params, toolName, and args to the service
-      return await mcpService.callTool(
-        { ...input.params, env: input.env },
-        input.toolName,
-        input.args,
-      );
+    .mutation(async ({ input, ctx }) => {
+      // Create a closure that binds fileService and userId to processContentBlocks
+      const boundProcessContentBlocks = async (blocks: ToolCallContent[]) =>
+        processContentBlocks(blocks, ctx.fileService);
+
+      // Pass the validated params, toolName, args, and bound processContentBlocks to the service
+      return await mcpService.callTool({
+        clientParams: { ...input.params, env: input.env },
+        toolName: input.toolName,
+        argsStr: input.args,
+        processContentBlocks: boundProcessContentBlocks,
+      });
     }),
 
   validMcpServerInstallable: mcpProcedure

package/src/server/routers/tools/mcp.ts

@@ -6,8 +6,12 @@ import {
 import { TRPCError } from '@trpc/server';
 import { z } from 'zod';
 
+import { ToolCallContent } from '@/libs/mcp';
 import { authedProcedure, router } from '@/libs/trpc/lambda';
+import { serverDatabase } from '@/libs/trpc/lambda/middleware';
+import { FileService } from '@/server/services/file';
 import { mcpService } from '@/server/services/mcp';
+import { processContentBlocks } from '@/server/services/mcp/contentProcessor';
 
 // Define Zod schemas for MCP Client parameters
 const httpParamsSchema = z.object({
@@ -37,7 +41,13 @@ const checkStdioEnvironment = (params: z.infer<typeof mcpClientParamsSchema>) =>
   }
 };
 
-const mcpProcedure = authedProcedure;
+const mcpProcedure = authedProcedure.use(serverDatabase).use(async ({ ctx, next }) => {
+  return next({
+    ctx: {
+      fileService: new FileService(ctx.serverDB, ctx.userId),
+    },
+  });
+});
 
 export const mcpRouter = router({
   getStreamableMcpServerManifest: mcpProcedure
@@ -95,11 +105,21 @@ export const mcpRouter = router({
         toolName: z.string(),
       }),
     )
-    .mutation(async ({ input }) => {
+    .mutation(async ({ input, ctx }) => {
       // Stdio check can be done here or rely on the service/client layer
       checkStdioEnvironment(input.params);
 
-      // Pass the validated params, toolName, and args to the service
-      return await mcpService.callTool(input.params, input.toolName, input.args);
+      // Create a closure that binds fileService and userId to processContentBlocks
+      const boundProcessContentBlocks = async (blocks: ToolCallContent[]) => {
+        return processContentBlocks(blocks, ctx.fileService);
+      };
+
+      // Pass the validated params, toolName, args, and bound processContentBlocks to the service
+      return await mcpService.callTool({
+        clientParams: input.params,
+        toolName: input.toolName,
+        argsStr: input.args,
+        processContentBlocks: boundProcessContentBlocks,
+      });
     }),
 });

package/src/server/services/file/impls/local.ts

@@ -1,3 +1,4 @@
+import debug from 'debug';
 import { sha256 } from 'js-sha256';
 import { existsSync, readFileSync } from 'node:fs';
 import path from 'node:path';
@@ -8,6 +9,8 @@ import { inferContentTypeFromImageUrl } from '@/utils/url';
 import { FileServiceImpl } from './type';
 import { extractKeyFromUrlOrReturnOriginal } from './utils';
 
+const log = debug('lobe-file:desktop-local');
+
 /**
  * 桌面应用本地文件服务实现
  */
@@ -202,7 +205,7 @@ export class DesktopLocalFileImpl implements FileServiceImpl {
         throw new Error('Failed to upload file via Electron IPC');
       }
 
-      console.log('[DesktopLocalFileImpl] File uploaded successfully:', result.metadata);
+      log('File uploaded successfully: %O', result.metadata);
       return { key: result.metadata.path };
     } catch (error) {
       console.error('[DesktopLocalFileImpl] Failed to upload media file:', error);

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

@@ -1,11 +1,12 @@
 import { LobeChatDatabase } from '@lobechat/database';
+import { inferContentTypeFromImageUrl, nanoid, uuid } from '@lobechat/utils';
 import { TRPCError } from '@trpc/server';
+import { sha256 } from 'js-sha256';
 
 import { serverDBEnv } from '@/config/db';
 import { FileModel } from '@/database/models/file';
 import { FileItem } from '@/database/schemas';
 import { TempFileManager } from '@/server/utils/tempFileManager';
-import { nanoid } from '@/utils/uuid';
 
 import { FileServiceImpl, createFileServiceModule } from './impls';
 
@@ -94,6 +95,100 @@ export class FileService {
     return this.impl.uploadMedia(key, buffer);
   }
 
+  /**
+   * Create file record (common method)
+   * Automatically handles globalFiles deduplication logic
+   *
+   * @param params - File parameters
+   * @param params.id - Optional custom file ID (defaults to auto-generated)
+   * @returns File record and proxy URL
+   */
+  public async createFileRecord(params: {
+    fileHash: string;
+    fileType: string;
+    id?: string;
+    name: string;
+    size: number;
+    url: string;
+  }): Promise<{ fileId: string; url: string }> {
+    // Check if hash already exists in globalFiles
+    const { isExist } = await this.fileModel.checkHash(params.fileHash);
+
+    // Create database record
+    // If hash doesn't exist, also create globalFiles record
+    const { id } = await this.fileModel.create(
+      {
+        fileHash: params.fileHash,
+        fileType: params.fileType,
+        id: params.id, // Use custom ID if provided
+        name: params.name,
+        size: params.size,
+        url: params.url,
+      },
+      !isExist, // insertToGlobalFiles
+    );
+
+    // Return unified proxy URL: /f/:id
+    return {
+      fileId: id,
+      url: `/f/${id}`,
+    };
+  }
+
+  /**
+   * Upload base64 data and create database record
+   * @param base64Data - Base64 data (supports data URI format or pure base64)
+   * @param pathname - File storage path (must include file extension)
+   * @returns Contains key (storage path), fileId (database record ID) and url (proxy access path)
+   */
+  public async uploadBase64(
+    base64Data: string,
+    pathname: string,
+  ): Promise<{ fileId: string; key: string; url: string }> {
+    let base64String: string;
+
+    // If data URI format (data:image/png;base64,xxx)
+    if (base64Data.startsWith('data:')) {
+      const commaIndex = base64Data.indexOf(',');
+      if (commaIndex === -1) {
+        throw new TRPCError({ code: 'BAD_REQUEST', message: 'Invalid base64 data format' });
+      }
+      base64String = base64Data.slice(commaIndex + 1);
+    } else {
+      // Pure base64 string
+      base64String = base64Data;
+    }
+
+    // Convert to Buffer
+    const buffer = Buffer.from(base64String, 'base64');
+
+    // Upload to storage (S3 or local)
+    const { key } = await this.uploadMedia(pathname, buffer);
+
+    // Extract filename from pathname
+    const name = pathname.split('/').pop() || 'unknown';
+
+    // Calculate file metadata
+    const size = buffer.length;
+    const fileType = inferContentTypeFromImageUrl(pathname) || 'application/octet-stream';
+    const hash = sha256(buffer);
+
+    // Generate UUID for cleaner URLs
+    const fileId = uuid();
+
+    // Use common method to create file record
+    const { fileId: createdId, url } = await this.createFileRecord({
+      fileHash: hash,
+      fileType,
+      id: fileId, // Use UUID instead of auto-generated ID
+      name,
+      size,
+      url: key, // Store original key (S3 key or desktop://)
+    });
+
+    return { fileId: createdId, key, url };
+  }
+
   async downloadFileToLocal(
     fileId: string,
   ): Promise<{ cleanup: () => void; file: FileItem; filePath: string }> {