diff --git a/src/database/_deprecated/models/file.ts b/src/database/_deprecated/models/file.ts index f4a3244d89148..87ef02257d0bd 100644 --- a/src/database/_deprecated/models/file.ts +++ b/src/database/_deprecated/models/file.ts @@ -1,5 +1,6 @@ import { DBModel } from '@/database/_deprecated/core/types/db'; import { DB_File, DB_FileSchema } from '@/database/_deprecated/schemas/files'; +import { clientS3Storage } from '@/services/file/ClientS3'; import { nanoid } from '@/utils/uuid'; import { BaseModel } from '../core'; @@ -20,7 +21,13 @@ class _FileModel extends BaseModel<'files'> { if (!item) return; // arrayBuffer to url - const base64 = Buffer.from(item.data!).toString('base64'); + let base64; + if (!item.data) { + const hash = (item.url as string).replace('client-s3://', ''); + base64 = await this.getBase64ByFileHash(hash); + } else { + base64 = Buffer.from(item.data).toString('base64'); + } return { ...item, url: `data:${item.fileType};base64,${base64}` }; } @@ -32,6 +39,13 @@ class _FileModel extends BaseModel<'files'> { async clear() { return this.table.clear(); } + + private async getBase64ByFileHash(hash: string) { + const fileItem = await clientS3Storage.getObject(hash); + if (!fileItem) throw new Error('file not found'); + + return Buffer.from(await fileItem.arrayBuffer()).toString('base64'); + } } export const FileModel = new _FileModel(); diff --git a/src/server/routers/lambda/file.ts b/src/server/routers/lambda/file.ts index a030c1ddf163a..f31fc4d1aee8d 100644 --- a/src/server/routers/lambda/file.ts +++ b/src/server/routers/lambda/file.ts @@ -32,9 +32,7 @@ export const fileRouter = router({ }), createFile: fileProcedure - .input( - UploadFileSchema.omit({ data: true, saveMode: true, url: true }).extend({ url: z.string() }), - ) + .input(UploadFileSchema.omit({ url: true }).extend({ url: z.string() })) .mutation(async ({ ctx, input }) => { const { isExist } = await ctx.fileModel.checkHash(input.hash!); diff --git a/src/services/file/ClientS3/index.test.ts b/src/services/file/ClientS3/index.test.ts new file mode 100644 index 0000000000000..266c82078ac49 --- /dev/null +++ b/src/services/file/ClientS3/index.test.ts @@ -0,0 +1,115 @@ +import { createStore, del, get, set } from 'idb-keyval'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; + +import { BrowserS3Storage } from './index'; + +// Mock idb-keyval +vi.mock('idb-keyval', () => ({ + createStore: vi.fn(), + set: vi.fn(), + get: vi.fn(), + del: vi.fn(), +})); + +let storage: BrowserS3Storage; +let mockStore = {}; + +beforeEach(() => { + // Reset all mocks before each test + vi.clearAllMocks(); + mockStore = {}; + (createStore as any).mockReturnValue(mockStore); + storage = new BrowserS3Storage(); +}); + +describe('BrowserS3Storage', () => { + describe('constructor', () => { + it('should create store when in browser environment', () => { + expect(createStore).toHaveBeenCalledWith('lobechat-local-s3', 'objects'); + }); + }); + + describe('putObject', () => { + it('should successfully put a file object', async () => { + const mockFile = new File(['test content'], 'test.txt', { type: 'text/plain' }); + const mockArrayBuffer = new ArrayBuffer(8); + vi.spyOn(mockFile, 'arrayBuffer').mockResolvedValue(mockArrayBuffer); + (set as any).mockResolvedValue(undefined); + + await storage.putObject('1-test-key', mockFile); + + expect(set).toHaveBeenCalledWith( + '1-test-key', + { + data: mockArrayBuffer, + name: 'test.txt', + type: 'text/plain', + }, + mockStore, + ); + }); + + it('should throw error when put operation fails', async () => { + const mockFile = new File(['test content'], 'test.txt', { type: 'text/plain' }); + const mockError = new Error('Storage error'); + (set as any).mockRejectedValue(mockError); + + await expect(storage.putObject('test-key', mockFile)).rejects.toThrow( + 'Failed to put file test.txt: Storage error', + ); + }); + }); + + describe('getObject', () => { + it('should successfully get a file object', async () => { + const mockData = { + data: new ArrayBuffer(8), + name: 'test.txt', + type: 'text/plain', + }; + (get as any).mockResolvedValue(mockData); + + const result = await storage.getObject('test-key'); + + expect(result).toBeInstanceOf(File); + expect(result?.name).toBe('test.txt'); + expect(result?.type).toBe('text/plain'); + }); + + it('should return undefined when file not found', async () => { + (get as any).mockResolvedValue(undefined); + + const result = await storage.getObject('test-key'); + + expect(result).toBeUndefined(); + }); + + it('should throw error when get operation fails', async () => { + const mockError = new Error('Storage error'); + (get as any).mockRejectedValue(mockError); + + await expect(storage.getObject('test-key')).rejects.toThrow( + 'Failed to get object (key=test-key): Storage error', + ); + }); + }); + + describe('deleteObject', () => { + it('should successfully delete a file object', async () => { + (del as any).mockResolvedValue(undefined); + + await storage.deleteObject('test-key2'); + + expect(del).toHaveBeenCalledWith('test-key2', {}); + }); + + it('should throw error when delete operation fails', async () => { + const mockError = new Error('Storage error'); + (del as any).mockRejectedValue(mockError); + + await expect(storage.deleteObject('test-key')).rejects.toThrow( + 'Failed to delete object (key=test-key): Storage error', + ); + }); + }); +}); diff --git a/src/services/file/ClientS3/index.ts b/src/services/file/ClientS3/index.ts new file mode 100644 index 0000000000000..94692f0ecb054 --- /dev/null +++ b/src/services/file/ClientS3/index.ts @@ -0,0 +1,58 @@ +import { createStore, del, get, set } from 'idb-keyval'; + +const BROWSER_S3_DB_NAME = 'lobechat-local-s3'; + +export class BrowserS3Storage { + private store; + + constructor() { + // skip server-side rendering + if (typeof window === 'undefined') return; + + this.store = createStore(BROWSER_S3_DB_NAME, 'objects'); + } + + /** + * 上传文件 + * @param key 文件 hash + * @param file File 对象 + */ + async putObject(key: string, file: File): Promise { + try { + const data = await file.arrayBuffer(); + await set(key, { data, name: file.name, type: file.type }, this.store); + } catch (e) { + throw new Error(`Failed to put file ${file.name}: ${(e as Error).message}`); + } + } + + /** + * 获取文件 + * @param key 文件 hash + * @returns File 对象 + */ + async getObject(key: string): Promise { + try { + const res = await get<{ data: ArrayBuffer; name: string; type: string }>(key, this.store); + if (!res) return; + + return new File([res.data], res!.name, { type: res?.type }); + } catch (e) { + throw new Error(`Failed to get object (key=${key}): ${(e as Error).message}`); + } + } + + /** + * 删除文件 + * @param key 文件 hash + */ + async deleteObject(key: string): Promise { + try { + await del(key, this.store); + } catch (e) { + throw new Error(`Failed to delete object (key=${key}): ${(e as Error).message}`); + } + } +} + +export const clientS3Storage = new BrowserS3Storage(); diff --git a/src/services/file/client.ts b/src/services/file/client.ts index babfa78322c31..71c1317f6c119 100644 --- a/src/services/file/client.ts +++ b/src/services/file/client.ts @@ -1,16 +1,27 @@ import { FileModel } from '@/database/_deprecated/models/file'; -import { DB_File } from '@/database/_deprecated/schemas/files'; -import { FileItem } from '@/types/files'; +import { clientS3Storage } from '@/services/file/ClientS3'; +import { FileItem, UploadFileParams } from '@/types/files'; import { IFileService } from './type'; export class ClientService implements IFileService { - async createFile(file: DB_File) { + async createFile(file: UploadFileParams) { // save to local storage // we may want to save to a remote server later - const res = await FileModel.create(file); - // arrayBuffer to url - const base64 = Buffer.from(file.data!).toString('base64'); + const res = await FileModel.create({ + createdAt: Date.now(), + data: undefined, + fileHash: file.hash, + fileType: file.fileType, + metadata: file.metadata, + name: file.name, + saveMode: 'url', + size: file.size, + url: file.url, + } as any); + + // get file to base64 url + const base64 = await this.getBase64ByFileHash(file.hash!); return { id: res.id, @@ -18,6 +29,11 @@ export class ClientService implements IFileService { }; } + // eslint-disable-next-line @typescript-eslint/no-unused-vars + async checkFileHash(_hash: string) { + return { isExist: false, metadata: {} }; + } + async getFile(id: string): Promise { const item = await FileModel.findById(id); if (!item) { @@ -49,4 +65,11 @@ export class ClientService implements IFileService { async removeAllFiles() { return FileModel.clear(); } + + private async getBase64ByFileHash(hash: string) { + const fileItem = await clientS3Storage.getObject(hash); + if (!fileItem) throw new Error('file not found'); + + return Buffer.from(await fileItem.arrayBuffer()).toString('base64'); + } } diff --git a/src/services/upload.ts b/src/services/upload.ts index 38f6cdd2959b6..64fec769ecad4 100644 --- a/src/services/upload.ts +++ b/src/services/upload.ts @@ -1,7 +1,8 @@ import { fileEnv } from '@/config/file'; import { edgeClient } from '@/libs/trpc/client'; import { API_ENDPOINTS } from '@/services/_url'; -import { FileMetadata, UploadFileParams } from '@/types/files'; +import { clientS3Storage } from '@/services/file/ClientS3'; +import { FileMetadata } from '@/types/files'; import { FileUploadState, FileUploadStatus } from '@/types/files/upload'; import { uuid } from '@/utils/uuid'; @@ -66,23 +67,14 @@ class UploadService { return result; }; - uploadToClientDB = async (params: UploadFileParams, file: File) => { - const { FileModel } = await import('@/database/_deprecated/models/file'); - const fileArrayBuffer = await file.arrayBuffer(); - - // save to local storage - // we may want to save to a remote server later - const res = await FileModel.create({ - createdAt: Date.now(), - ...params, - data: fileArrayBuffer, - }); - // arrayBuffer to url - const base64 = Buffer.from(fileArrayBuffer).toString('base64'); + uploadToClientS3 = async (hash: string, file: File): Promise => { + await clientS3Storage.putObject(hash, file); return { - id: res.id, - url: `data:${params.fileType};base64,${base64}`, + date: (Date.now() / 1000 / 60 / 60).toFixed(0), + dirname: '', + filename: file.name, + path: `client-s3://${hash}`, }; }; diff --git a/src/store/chat/slices/builtinTool/action.test.ts b/src/store/chat/slices/builtinTool/action.test.ts index 795806b441955..63bc5b990a55d 100644 --- a/src/store/chat/slices/builtinTool/action.test.ts +++ b/src/store/chat/slices/builtinTool/action.test.ts @@ -2,6 +2,8 @@ import { act, renderHook } from '@testing-library/react'; import { describe, expect, it, vi } from 'vitest'; import { fileService } from '@/services/file'; +import { ClientService } from '@/services/file/client'; +import { messageService } from '@/services/message'; import { imageGenerationService } from '@/services/textToImage'; import { uploadService } from '@/services/upload'; import { chatSelectors } from '@/store/chat/selectors'; @@ -39,17 +41,23 @@ describe('chatToolSlice', () => { vi.spyOn(uploadService, 'getImageFileByUrlWithCORS').mockResolvedValue( new File(['1'], 'file.png', { type: 'image/png' }), ); - vi.spyOn(uploadService, 'uploadToClientDB').mockResolvedValue({} as any); - vi.spyOn(fileService, 'createFile').mockResolvedValue({ id: mockId, url: '' }); + vi.spyOn(uploadService, 'uploadToClientS3').mockResolvedValue({} as any); + vi.spyOn(ClientService.prototype, 'createFile').mockResolvedValue({ + id: mockId, + url: '', + }); vi.spyOn(result.current, 'toggleDallEImageLoading'); + vi.spyOn(ClientService.prototype, 'checkFileHash').mockImplementation(async () => ({ + isExist: false, + metadata: {}, + })); await act(async () => { await result.current.generateImageFromPrompts(prompts, messageId); }); // For each prompt, loading is toggled on and then off expect(imageGenerationService.generateImage).toHaveBeenCalledTimes(prompts.length); - expect(uploadService.uploadToClientDB).toHaveBeenCalledTimes(prompts.length); - + expect(uploadService.uploadToClientS3).toHaveBeenCalledTimes(prompts.length); expect(result.current.toggleDallEImageLoading).toHaveBeenCalledTimes(prompts.length * 2); }); }); @@ -75,6 +83,7 @@ describe('chatToolSlice', () => { content: initialMessageContent, }) as ChatMessage, ); + vi.spyOn(messageService, 'updateMessage').mockResolvedValueOnce(undefined); await act(async () => { await result.current.updateImageItem(messageId, updateFunction); diff --git a/src/store/file/slices/upload/action.ts b/src/store/file/slices/upload/action.ts index 1134d7907f430..3f73de77b08b4 100644 --- a/src/store/file/slices/upload/action.ts +++ b/src/store/file/slices/upload/action.ts @@ -6,14 +6,11 @@ import { message } from '@/components/AntdStaticMethods'; import { LOBE_CHAT_CLOUD } from '@/const/branding'; import { isServerMode } from '@/const/version'; import { fileService } from '@/services/file'; -import { ServerService } from '@/services/file/server'; import { uploadService } from '@/services/upload'; import { FileMetadata, UploadFileItem } from '@/types/files'; import { FileStore } from '../../store'; -const serverFileService = new ServerService(); - interface UploadWithProgressParams { file: File; knowledgeBaseId?: string; @@ -43,10 +40,6 @@ interface UploadWithProgressResult { } export interface FileUploadAction { - internal_uploadToClientDB: ( - params: Omit, - ) => Promise; - internal_uploadToServer: (params: UploadWithProgressParams) => Promise; uploadWithProgress: ( params: UploadWithProgressParams, ) => Promise; @@ -57,51 +50,14 @@ export const createFileUploadSlice: StateCreator< [['zustand/devtools', never]], [], FileUploadAction -> = (set, get) => ({ - internal_uploadToClientDB: async ({ file, onStatusUpdate, skipCheckFileType }) => { - if (!skipCheckFileType && !file.type.startsWith('image')) { - onStatusUpdate?.({ id: file.name, type: 'removeFile' }); - message.info({ - content: t('upload.fileOnlySupportInServerMode', { - cloud: LOBE_CHAT_CLOUD, - ext: file.name.split('.').pop(), - ns: 'error', - }), - duration: 5, - }); - return; - } - - const fileArrayBuffer = await file.arrayBuffer(); - - const hash = sha256(fileArrayBuffer); - - const data = await uploadService.uploadToClientDB( - { fileType: file.type, hash, name: file.name, saveMode: 'local', size: file.size }, - file, - ); - - onStatusUpdate?.({ - id: file.name, - type: 'updateFile', - value: { - fileUrl: data.url, - id: data.id, - status: 'success', - uploadState: { progress: 100, restTime: 0, speed: 0 }, - }, - }); - - return data; - }, - - internal_uploadToServer: async ({ file, onStatusUpdate, knowledgeBaseId }) => { +> = () => ({ + uploadWithProgress: async ({ file, onStatusUpdate, knowledgeBaseId, skipCheckFileType }) => { const fileArrayBuffer = await file.arrayBuffer(); // 1. check file hash const hash = sha256(fileArrayBuffer); - const checkStatus = await serverFileService.checkFileHash(hash); + const checkStatus = await fileService.checkFileHash(hash); let metadata: FileMetadata; // 2. if file exist, just skip upload @@ -112,17 +68,37 @@ export const createFileUploadSlice: StateCreator< type: 'updateFile', value: { status: 'processing', uploadState: { progress: 100, restTime: 0, speed: 0 } }, }); - } else { - // 2. if file don't exist, need upload files - metadata = await uploadService.uploadWithProgress(file, { - onProgress: (status, upload) => { - onStatusUpdate?.({ - id: file.name, - type: 'updateFile', - value: { status: status === 'success' ? 'processing' : status, uploadState: upload }, + } + // 2. if file don't exist, need upload files + else { + // if is server mode, upload to server s3, or upload to client s3 + if (isServerMode) { + metadata = await uploadService.uploadWithProgress(file, { + onProgress: (status, upload) => { + onStatusUpdate?.({ + id: file.name, + type: 'updateFile', + value: { status: status === 'success' ? 'processing' : status, uploadState: upload }, + }); + }, + }); + } else { + if (!skipCheckFileType && !file.type.startsWith('image')) { + onStatusUpdate?.({ id: file.name, type: 'removeFile' }); + message.info({ + content: t('upload.fileOnlySupportInServerMode', { + cloud: LOBE_CHAT_CLOUD, + ext: file.name.split('.').pop(), + ns: 'error', + }), + duration: 5, }); - }, - }); + return; + } + + // Upload to the indexeddb in the browser + metadata = await uploadService.uploadToClientS3(hash, file); + } } // 3. use more powerful file type detector to get file type @@ -138,12 +114,10 @@ export const createFileUploadSlice: StateCreator< // 4. create file to db const data = await fileService.createFile( { - createdAt: Date.now(), fileType, hash, metadata, name: file.name, - saveMode: 'url', size: file.size, url: metadata.path, }, @@ -163,12 +137,4 @@ export const createFileUploadSlice: StateCreator< return data; }, - - uploadWithProgress: async (payload) => { - const { internal_uploadToServer, internal_uploadToClientDB } = get(); - - if (isServerMode) return internal_uploadToServer(payload); - - return internal_uploadToClientDB(payload); - }, }); diff --git a/src/types/files/upload.ts b/src/types/files/upload.ts index 83aa3d5be9ce9..b94b69a5fdee3 100644 --- a/src/types/files/upload.ts +++ b/src/types/files/upload.ts @@ -53,7 +53,6 @@ export const FileMetadataSchema = z.object({ export type FileMetadata = z.infer; export const UploadFileSchema = z.object({ - data: z.instanceof(ArrayBuffer).optional(), /** * file type * @example 'image/png' @@ -77,7 +76,6 @@ export const UploadFileSchema = z.object({ * local mean save the raw file into data * url mean upload the file to a cdn and then save the url */ - saveMode: z.enum(['local', 'url']), /** * file size */ @@ -89,3 +87,11 @@ export const UploadFileSchema = z.object({ }); export type UploadFileParams = z.infer; + +export interface CheckFileHashResult { + fileType?: string; + isExist: boolean; + metadata?: unknown; + size?: number; + url?: string; +}