@lobehub/chat 0.100.5 → 0.101.0

package/src/store/files/slices/images/action.test.ts

@@ -0,0 +1,186 @@
+import { act, renderHook } from '@testing-library/react';
+import useSWR from 'swr';
+import { Mock, afterEach, beforeAll, beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { fileService } from '@/services/file';
+
+import { useFileStore as useStore } from '../../store';
+
+vi.mock('zustand/traditional');
+
+// Mocks for fileService
+vi.mock('@/services/file', () => ({
+  fileService: {
+    removeFile: vi.fn(),
+    uploadFile: vi.fn(),
+    getFile: vi.fn(),
+  },
+}));
+
+// Mock for useSWR
+vi.mock('swr', () => ({
+  default: vi.fn(),
+}));
+
+//  mock the arrayBuffer
+beforeAll(() => {
+  Object.defineProperty(File.prototype, 'arrayBuffer', {
+    writable: true,
+    value: function () {
+      return new Promise((resolve) => {
+        const reader = new FileReader();
+        reader.onloadend = () => {
+          resolve(reader.result);
+        };
+        reader.readAsArrayBuffer(this);
+      });
+    },
+  });
+});
+
+beforeEach(() => {
+  // Reset all mocks before each test
+  vi.resetAllMocks();
+});
+
+describe('useFileStore:images', () => {
+  it('clearImageList should clear the inputFilesList', () => {
+    const { result } = renderHook(() => useStore());
+
+    // Populate the list to clear it later
+    act(() => {
+      useStore.setState({ inputFilesList: ['test-id'] });
+    });
+
+    expect(result.current.inputFilesList).toEqual(['test-id']);
+
+    act(() => {
+      result.current.clearImageList();
+    });
+
+    expect(result.current.inputFilesList).toEqual([]);
+  });
+
+  it('removeFile should call fileService.removeFile and update the store', async () => {
+    const { result } = renderHook(() => useStore());
+
+    const fileId = 'test-id';
+
+    // Mock the fileService.removeFile to resolve
+    (fileService.removeFile as Mock).mockResolvedValue(undefined);
+
+    // Populate the list to remove an item later
+    act(() => {
+      useStore.setState(({ inputFilesList }) => ({ inputFilesList: [...inputFilesList, fileId] }));
+      //   // result.current.inputFilesList.push(fileId);
+    });
+
+    await act(async () => {
+      await result.current.removeFile(fileId);
+    });
+
+    expect(fileService.removeFile).toHaveBeenCalledWith(fileId);
+    expect(result.current.inputFilesList).toEqual([]);
+  });
+
+  // Test for useFetchFile
+  it('useFetchFile should call useSWR and update the store', async () => {
+    const fileId = 'test-id';
+    const fileData = {
+      id: fileId,
+      name: 'test',
+      url: 'blob:test',
+      fileType: 'image/png',
+      base64Url: '',
+      saveMode: 'local',
+    };
+
+    // Mock the fileService.getFile to resolve with fileData
+    (fileService.getFile as Mock).mockResolvedValue(fileData);
+
+    // Mock useSWR to call the fetcher function immediately
+    const useSWRMock = vi.mocked(useSWR);
+    useSWRMock.mockImplementation(((key: string, fetcher: any) => {
+      const data = fetcher(key);
+      return { data, error: undefined, isValidating: false, mutate: vi.fn() };
+    }) as any);
+
+    const { result } = renderHook(() => useStore().useFetchFile(fileId));
+
+    await act(async () => {
+      await result.current.data;
+    });
+
+    expect(fileService.getFile).toHaveBeenCalledWith(fileId);
+
+    // Since we are not rendering a component with the hook, we cannot test the state update here
+    // Instead, we would need to use a test renderer that can work with hooks, like @testing-library/react
+  });
+
+  describe('uploadFile', () => {
+    it('uploadFile should handle errors', async () => {
+      const { result } = renderHook(() => useStore());
+      const testFile = new File(['content'], 'test.png', { type: 'image/png' });
+
+      // 模拟 fileService.uploadFile 抛出错误
+      const errorMessage = 'Upload failed';
+      (fileService.uploadFile as Mock).mockRejectedValue(new Error(errorMessage));
+
+      // Mock console.error for testing
+      const consoleErrorMock = vi.spyOn(console, 'error').mockImplementation(() => {});
+
+      await act(async () => {
+        await result.current.uploadFile(testFile);
+      });
+
+      expect(fileService.uploadFile).toHaveBeenCalledWith({
+        createdAt: testFile.lastModified,
+        data: await testFile.arrayBuffer(),
+        fileType: testFile.type,
+        name: testFile.name,
+        saveMode: 'local',
+        size: testFile.size,
+      });
+      // 由于上传失败，inputFilesList 应该没有变化
+      expect(result.current.inputFilesList).toEqual([]);
+      // 确保错误被正确记录
+      expect(consoleErrorMock).toHaveBeenCalledWith('upload error:', expect.any(Error));
+
+      // Cleanup mock
+      consoleErrorMock.mockRestore();
+    });
+
+    it('uploadFile should upload the file and update inputFilesList', async () => {
+      const { result } = renderHook(() => useStore());
+      const testFile = new File(['content'], 'test.png', { type: 'image/png' });
+
+      // 模拟 fileService.uploadFile 返回的数据
+      const uploadedFileData = {
+        id: 'new-file-id',
+        createdAt: testFile.lastModified,
+        data: await testFile.arrayBuffer(),
+        fileType: testFile.type,
+        name: testFile.name,
+        saveMode: 'local',
+        size: testFile.size,
+      };
+
+      // Mock the fileService.uploadFile to resolve with uploadedFileData
+      (fileService.uploadFile as Mock).mockResolvedValue(uploadedFileData);
+
+      await act(async () => {
+        await result.current.uploadFile(testFile);
+      });
+
+      expect(fileService.uploadFile).toHaveBeenCalledWith({
+        createdAt: testFile.lastModified,
+        data: await testFile.arrayBuffer(),
+        fileType: testFile.type,
+        name: testFile.name,
+        saveMode: 'local',
+        size: testFile.size,
+      });
+      expect(result.current.inputFilesList).toContain(uploadedFileData.id);
+    });
+  });
+});

package/src/store/files/slices/images/action.ts

@@ -0,0 +1,83 @@
+import { produce } from 'immer';
+import useSWR, { SWRResponse } from 'swr';
+import { StateCreator } from 'zustand/vanilla';
+
+import { fileService } from '@/services/file';
+import { FilePreview } from '@/types/files';
+import { setNamespace } from '@/utils/storeDebug';
+
+import { FileStore } from '../../store';
+
+const t = setNamespace('files/image');
+
+/**
+ * 代理行为接口
+ */
+export interface FileAction {
+  clearImageList: () => void;
+  removeFile: (id: string) => Promise<void>;
+  setImageMapItem: (id: string, item: FilePreview) => void;
+  uploadFile: (file: File) => Promise<void>;
+
+  useFetchFile: (id: string) => SWRResponse<FilePreview>;
+}
+
+export const createFileSlice: StateCreator<
+  FileStore,
+  [['zustand/devtools', never]],
+  [],
+  FileAction
+> = (set, get) => ({
+  clearImageList: () => {
+    set({ inputFilesList: [] }, false, t('clearImageList'));
+  },
+  removeFile: async (id) => {
+    await fileService.removeFile(id);
+
+    set(
+      ({ inputFilesList }) => ({ inputFilesList: inputFilesList.filter((i) => i !== id) }),
+      false,
+      t('removeFile'),
+    );
+  },
+  setImageMapItem: (id, item) => {
+    set(
+      produce((draft) => {
+        if (draft.imagesMap[id]) return;
+
+        draft.imagesMap[id] = item;
+      }),
+      false,
+      t('setImageMapItem'),
+    );
+  },
+  uploadFile: async (file) => {
+    try {
+      const data = await fileService.uploadFile({
+        createdAt: file.lastModified,
+        data: await file.arrayBuffer(),
+        fileType: file.type,
+        name: file.name,
+        saveMode: 'local',
+        size: file.size,
+      });
+
+      set(
+        ({ inputFilesList }) => ({ inputFilesList: [...inputFilesList, data.id] }),
+        false,
+        t('uploadFile'),
+      );
+    } catch (error) {
+      // 提示用户上传失败
+      console.error('upload error:', error);
+    }
+  },
+  useFetchFile: (id) =>
+    useSWR(id, async (id) => {
+      const item = await fileService.getFile(id);
+
+      get().setImageMapItem(id, item);
+
+      return item;
+    }),
+});

package/src/store/files/slices/images/index.ts

@@ -0,0 +1,3 @@
+export * from './action';
+export * from './initialState';
+// export * from './selectors';

package/src/store/files/slices/images/initialState.ts

@@ -0,0 +1,11 @@
+import { FilePreview } from '@/types/files';
+
+export interface ImageFileState {
+  imagesMap: Record<string, FilePreview>;
+  inputFilesList: string[];
+}
+
+export const initialImageFileState: ImageFileState = {
+  imagesMap: {},
+  inputFilesList: [],
+};

package/src/store/files/store.ts

@@ -0,0 +1,27 @@
+import { devtools } from 'zustand/middleware';
+import { shallow } from 'zustand/shallow';
+import { createWithEqualityFn } from 'zustand/traditional';
+import { StateCreator } from 'zustand/vanilla';
+
+import { isDev } from '@/utils/env';
+
+import { FilesStoreState, initialState } from './initialState';
+import { FileAction, createFileSlice } from './slices/images';
+
+//  ===============  聚合 createStoreFn ============ //
+
+export type FileStore = FilesStoreState & FileAction;
+
+const createStore: StateCreator<FileStore, [['zustand/devtools', never]]> = (...parameters) => ({
+  ...initialState,
+  ...createFileSlice(...parameters),
+});
+
+//  ===============  实装 useStore ============ //
+
+export const useFileStore = createWithEqualityFn<FileStore>()(
+  devtools(createStore, {
+    name: 'LobeChat_File' + (isDev ? '_DEV' : ''),
+  }),
+  shallow,
+);

package/src/store/global/selectors/settings.test.ts

@@ -75,7 +75,7 @@ describe('settingsSelectors', () => {
       });
     });
   });
-  //
+
   // describe('defaultAgent', () => {
   //   it('should merge DEFAULT_AGENT and s.settings.defaultAgent correctly', () => {
   //     const s: GlobalStore = {
@@ -95,8 +95,6 @@ describe('settingsSelectors', () => {
   //
   //     const result = settingsSelectors.defaultAgent(s);
   //
-  //
-  //
   //     expect(result).toEqual(expected);
   //   });
   // });

package/src/store/session/slices/agentConfig/selectors.ts

@@ -1,6 +1,6 @@
 import { t } from 'i18next';
 
-import { DEFAULT_OPENAI_MODEL_LIST } from '@/const/llm';
+import { DEFAULT_OPENAI_MODEL_LIST, VISION_MODEL_WHITE_LIST } from '@/const/llm';
 import { DEFAULT_AVATAR, DEFAULT_BACKGROUND_COLOR } from '@/const/meta';
 import { SessionStore } from '@/store/session';
 import { LanguageModel } from '@/types/llm';
@@ -49,6 +49,11 @@ const currentAgentModel = (s: SessionStore): LanguageModel | string => {
   return config?.model || LanguageModel.GPT3_5;
 };
 
+const modelHasVisionAbility = (s: SessionStore): boolean => {
+  const model = currentAgentModel(s);
+  return VISION_MODEL_WHITE_LIST.includes(model);
+};
+
 const currentAgentPlugins = (s: SessionStore) => {
   const config = currentAgentConfig(s);
 
@@ -87,5 +92,6 @@ export const agentSelectors = {
   getDescription,
   getTitle,
   hasSystemRole,
+  modelHasVisionAbility,
   showTokenTag,
 };

package/src/store/session/slices/chat/actions/message.ts

@@ -1,10 +1,14 @@
 import { template } from 'lodash-es';
 import { StateCreator } from 'zustand/vanilla';
 
+import { VISION_MODEL_WHITE_LIST } from '@/const/llm';
 import { LOADING_FLAT } from '@/const/message';
+import { VISION_MODEL_DEFAULT_MAX_TOKENS } from '@/const/settings';
 import { fetchChatModel } from '@/services/chatModel';
+import { filesSelectors, useFileStore } from '@/store/files';
 import { SessionStore } from '@/store/session';
 import { ChatMessage } from '@/types/chatMessage';
+import { OpenAIChatMessage, UserMessageContentPart } from '@/types/openai/chat';
 import { fetchSSE } from '@/utils/fetch';
 import { isFunctionMessageAtStart, testFunctionMessageAtEnd } from '@/utils/message';
 import { setNamespace } from '@/utils/storeDebug';
@@ -12,6 +16,7 @@ import { nanoid } from '@/utils/uuid';
 
 import { agentSelectors } from '../../agentConfig/selectors';
 import { sessionSelectors } from '../../session/selectors';
+import { FileDispatch, filesReducer } from '../reducers/files';
 import { MessageDispatch, messagesReducer } from '../reducers/message';
 import { chatSelectors } from '../selectors';
 import { getSlicedMessagesWithConfig } from '../utils';
@@ -37,6 +42,10 @@ export interface ChatMessageAction {
    * @param id - 消息 ID
    */
   deleteMessage: (id: string) => void;
+  /**
+   * agent files dispatch method
+   */
+  dispatchAgentFile: (payload: FileDispatch) => void;
   /**
    * 分发消息
    * @param payload - 消息分发参数
@@ -66,7 +75,7 @@ export interface ChatMessageAction {
    * 发送消息
    * @param text - 消息文本
    */
-  sendMessage: (text: string) => Promise<void>;
+  sendMessage: (text: string, images?: { id: string; url: string }[]) => Promise<void>;
   stopGenerateMessage: () => void;
   toggleChatLoading: (
     loading: boolean,
@@ -164,6 +173,15 @@ export const chatMessage: StateCreator<
     get().dispatchMessage({ id, type: 'deleteMessage' });
   },
 
+  dispatchAgentFile: (payload) => {
+    const { activeId } = get();
+    const session = sessionSelectors.currentSession(get());
+    if (!activeId || !session) return;
+
+    const files = filesReducer(session.files || [], payload);
+
+    get().dispatchSession({ files, id: activeId, type: 'updateSessionFiles' });
+  },
   dispatchMessage: (payload) => {
     const { activeId } = get();
     const session = sessionSelectors.currentSession(get());
@@ -173,6 +191,7 @@ export const chatMessage: StateCreator<
 
     get().dispatchSession({ chats, id: activeId, type: 'updateSessionChat' });
   },
+
   fetchAIChatMessage: async (messages, assistantId) => {
     const { dispatchMessage, toggleChatLoading } = get();
 
@@ -186,17 +205,17 @@ export const chatMessage: StateCreator<
 
     const compiler = template(config.inputTemplate, { interpolate: /{{([\S\s]+?)}}/g });
 
-    // ========================== //
-    //   对 messages 做统一预处理    //
-    // ========================== //
+    // ================================== //
+    //   messages uniformly preprocess    //
+    // ================================== //
 
-    // 1. 按参数设定截断长度
-    const slicedMessages = getSlicedMessagesWithConfig(messages, config);
+    // 1. slice messages with config
+    let preprocessMsgs = getSlicedMessagesWithConfig(messages, config);
 
-    // 2. 替换 inputMessage 模板
-    const postMessages = !config.inputTemplate
-      ? slicedMessages
-      : slicedMessages.map((m) => {
+    // 2. replace inputMessage template
+    preprocessMsgs = !config.inputTemplate
+      ? preprocessMsgs
+      : preprocessMsgs.map((m) => {
           if (m.role === 'user') {
             try {
               return { ...m, content: compiler({ text: m.content }) };
@@ -206,12 +225,41 @@ export const chatMessage: StateCreator<
               return m;
             }
           }
+
           return m;
         });
 
-    // 3. 添加 systemRole
+    // 3. add systemRole
     if (config.systemRole) {
-      postMessages.unshift({ content: config.systemRole, role: 'system' } as ChatMessage);
+      preprocessMsgs.unshift({ content: config.systemRole, role: 'system' } as ChatMessage);
+    }
+
+    let postMessages: OpenAIChatMessage[] = preprocessMsgs;
+
+    // 4. handle content type for vision model
+    // for the models with visual ability, add image url to content
+    // refs: https://platform.openai.com/docs/guides/vision/quick-start
+    if (VISION_MODEL_WHITE_LIST.includes(config.model)) {
+      postMessages = preprocessMsgs.map((m) => {
+        if (!m.files) return m;
+
+        const imageList = filesSelectors.getImageUrlOrBase64ByList(m.files)(
+          useFileStore.getState(),
+        );
+
+        if (imageList.length === 0) return m;
+
+        const content: UserMessageContentPart[] = [
+          { text: m.content, type: 'text' },
+          ...imageList.map(
+            (i) => ({ image_url: { detail: 'auto', url: i.url }, type: 'image_url' }) as const,
+          ),
+        ];
+        return { ...m, content };
+      });
+
+      // due to vision model's default max_tokens is very small, we need to set the max_tokens a larger one.
+      if (!config.params.max_tokens) config.params.max_tokens = VISION_MODEL_DEFAULT_MAX_TOKENS;
     }
 
     const fetcher = () =>
@@ -248,7 +296,7 @@ export const chatMessage: StateCreator<
 
     toggleChatLoading(false, undefined, t('generateMessage(end)') as string);
 
-    // also exist message like this: 请稍等，我帮您查询一下。{"function_call": {"name": "plugin-identifier____recommendClothes____standalone", "arguments": "{\n "mood": "",\n "gender": "man"\n}"}}
+    // also exist message like this:  请稍等，我帮您查询一下。{"function_call": {"name": "plugin-identifier____recommendClothes____standalone", "arguments": "{\n "mood": "",\n "gender": "man"\n}"}}
     if (!isFunctionCall) {
       const { content, valid } = testFunctionMessageAtEnd(output);
 
@@ -305,13 +353,27 @@ export const chatMessage: StateCreator<
     await coreProcessMessage(contextMessages, latestMsg.id);
   },
 
-  sendMessage: async (message) => {
-    const { dispatchMessage, coreProcessMessage, activeTopicId } = get();
+  sendMessage: async (message, files) => {
+    const { dispatchMessage, dispatchAgentFile, coreProcessMessage, activeTopicId } = get();
     const session = sessionSelectors.currentSession(get());
     if (!session || !message) return;
 
     const userId = nanoid();
-    dispatchMessage({ id: userId, message, role: 'user', type: 'addMessage' });
+
+    dispatchMessage({
+      id: userId,
+      message: message,
+      role: 'user',
+      type: 'addMessage',
+    });
+
+    // if message has attached with files, then add files to message and the agent
+    if (files && files.length > 0) {
+      const fileIdList = files.map((f) => f.id);
+      dispatchMessage({ id: userId, key: 'files', type: 'updateMessage', value: fileIdList });
+
+      dispatchAgentFile({ files: fileIdList, type: 'addFiles' });
+    }
 
     // if there is activeTopicId，then add topicId to message
     if (activeTopicId) {
@@ -341,7 +403,6 @@ export const chatMessage: StateCreator<
 
     toggleChatLoading(false);
   },
-
   toggleChatLoading: (loading, id, action) => {
     if (loading) {
       const abortController = new AbortController();

package/src/store/session/slices/chat/reducers/files.test.ts

@@ -0,0 +1,38 @@
+import { FileDispatch, FilesState, filesReducer } from './files';
+
+describe('filesReducer', () => {
+  it('should add a file to the state', () => {
+    const initialState: FilesState = ['file1', 'file2'];
+    const action: FileDispatch = { type: 'addFile', file: 'file3' };
+    const newState = filesReducer(initialState, action);
+    expect(newState).toEqual(['file1', 'file2', 'file3']);
+  });
+
+  it('should delete a file from the state by ID', () => {
+    const initialState: FilesState = ['file1', 'file2', 'file3'];
+    const action: FileDispatch = { type: 'deleteFile', id: 'file2' };
+    const newState = filesReducer(initialState, action);
+    expect(newState).toEqual(['file1', 'file3']);
+  });
+
+  it('should return the state unchanged if file ID does not exist for deletion', () => {
+    const initialState: FilesState = ['file1', 'file2', 'file3'];
+    const action: FileDispatch = { type: 'deleteFile', id: 'file4' };
+    const newState = filesReducer(initialState, action);
+    expect(newState).toEqual(['file1', 'file2', 'file3']);
+  });
+
+  it('should add multiple files to the state', () => {
+    const initialState: FilesState = ['file1', 'file2'];
+    const action: FileDispatch = { type: 'addFiles', files: ['file3', 'file4'] };
+    const newState = filesReducer(initialState, action);
+    expect(newState).toEqual(['file1', 'file2', 'file3', 'file4']);
+  });
+
+  it('should return the initial state if the action type is unknown', () => {
+    const initialState: FilesState = ['file1', 'file2'];
+    const action = { type: 'unknown', id: 'file1' };
+    const newState = filesReducer(initialState, action as FileDispatch);
+    expect(newState).toEqual(['file1', 'file2']);
+  });
+});

package/src/store/session/slices/chat/reducers/files.ts

@@ -0,0 +1,37 @@
+import { produce } from 'immer';
+
+export type FilesState = string[];
+
+export type FileDispatch =
+  | { file: string; type: 'addFile' }
+  | { id: string; type: 'deleteFile' }
+  | { files: string[]; type: 'addFiles' };
+
+export const filesReducer = (state: FilesState, payload: FileDispatch): FilesState => {
+  switch (payload.type) {
+    case 'addFile': {
+      return produce(state, (draftState) => {
+        draftState.push(payload.file);
+      });
+    }
+
+    case 'deleteFile': {
+      return produce(state, (draftState) => {
+        const index = draftState.indexOf(payload.id);
+        if (index !== -1) {
+          draftState.splice(index, 1);
+        }
+      });
+    }
+
+    case 'addFiles': {
+      return produce(state, (draftState) => {
+        draftState.push(...payload.files);
+      });
+    }
+
+    default: {
+      return state;
+    }
+  }
+};

package/src/store/session/slices/session/initialState.ts

@@ -25,6 +25,7 @@ export const initLobeSession: LobeAgentSession = {
   chats: {},
   config: initialLobeAgentConfig,
   createAt: Date.now(),
+  files: [],
   id: '',
   meta: DEFAULT_AGENT_META,
   type: LobeSessionType.Agent,