@lobehub/chat 1.98.1 → 1.99.0

package/src/store/aiInfra/slices/aiProvider/action.ts

@@ -7,7 +7,7 @@ import { isDeprecatedEdition, isDesktop, isUsePgliteDB } from '@/const/version';
 import { useClientDataSWR } from '@/libs/swr';
 import { aiProviderService } from '@/services/aiProvider';
 import { AiInfraStore } from '@/store/aiInfra/store';
-import { ModelAbilities } from '@/types/aiModel';
+import { AIImageModelCard, ModelAbilities } from '@/types/aiModel';
 import {
   AiProviderDetailItem,
   AiProviderListItem,
@@ -15,6 +15,7 @@ import {
   AiProviderSortMap,
   AiProviderSourceEnum,
   CreateAiProviderParams,
+  EnabledProvider,
   UpdateAiProviderConfigParams,
   UpdateAiProviderParams,
 } from '@/types/aiProvider';
@@ -175,11 +176,20 @@ export const createAiProviderSlice: StateCreator<
         if (isLogin) return aiProviderService.getAiProviderRuntimeState();
 
         const { LOBE_DEFAULT_MODEL_LIST } = await import('@/config/aiModels');
+        const enabledAiProviders: EnabledProvider[] = DEFAULT_MODEL_PROVIDER_LIST.filter(
+          (provider) => provider.enabled,
+        ).map((item) => ({ id: item.id, name: item.name, source: 'builtin' }));
+        const allModels = LOBE_DEFAULT_MODEL_LIST;
         return {
-          enabledAiModels: LOBE_DEFAULT_MODEL_LIST.filter((m) => m.enabled),
-          enabledAiProviders: DEFAULT_MODEL_PROVIDER_LIST.filter(
-            (provider) => provider.enabled,
-          ).map((item) => ({ id: item.id, name: item.name, source: 'builtin' })),
+          enabledAiModels: allModels.filter((m) => m.enabled),
+          enabledAiProviders: enabledAiProviders,
+          enabledImageAiProviders: enabledAiProviders
+            .filter((provider) => {
+              return allModels.some(
+                (model) => model.providerId === provider.id && model.type === 'image',
+              );
+            })
+            .map((item) => ({ id: item.id, name: item.name, source: 'builtin' })),
           runtimeConfig: {},
         };
       },
@@ -196,6 +206,9 @@ export const createAiProviderSlice: StateCreator<
                 contextWindowTokens: model.contextWindowTokens,
                 displayName: model.displayName ?? '',
                 id: model.id,
+                ...(model.type === 'image' && {
+                  parameters: (model as AIImageModelCard).parameters,
+                }),
               }));
 
             return uniqBy(models, 'id');
@@ -207,6 +220,12 @@ export const createAiProviderSlice: StateCreator<
             children: getModelListByType(provider.id, 'chat'),
             name: provider.name || provider.id,
           }));
+
+          const enabledImageModelList = data.enabledImageAiProviders.map((provider) => ({
+            ...provider,
+            children: getModelListByType(provider.id, 'image'),
+            name: provider.name || provider.id,
+          }));
           const { LOBE_DEFAULT_MODEL_LIST } = await import('@/config/aiModels');
 
           set(
@@ -216,6 +235,7 @@ export const createAiProviderSlice: StateCreator<
               enabledAiModels: data.enabledAiModels,
               enabledAiProviders: data.enabledAiProviders,
               enabledChatModelList,
+              enabledImageModelList,
             },
             false,
             'useFetchAiProviderRuntimeState',

package/src/store/aiInfra/slices/aiProvider/initialState.ts

@@ -19,6 +19,7 @@ export interface AIProviderState {
   enabledAiProviders?: EnabledProvider[];
   // used for select
   enabledChatModelList?: EnabledProviderWithModels[];
+  enabledImageModelList?: EnabledProviderWithModels[];
   initAiProviderList: boolean;
   providerSearchKeyword: string;
 }

package/src/store/aiInfra/slices/aiProvider/selectors.ts

@@ -10,6 +10,8 @@ const enabledAiProviderList = (s: AIProviderStoreState) =>
 const disabledAiProviderList = (s: AIProviderStoreState) =>
   s.aiProviderList.filter((item) => !item.enabled);
 
+const enabledImageModelList = (s: AIProviderStoreState) => s.enabledImageModelList || [];
+
 const isProviderEnabled = (id: string) => (s: AIProviderStoreState) =>
   enabledAiProviderList(s).some((i) => i.id === id);
 
@@ -113,6 +115,7 @@ export const aiProviderSelectors = {
   activeProviderConfig,
   disabledAiProviderList,
   enabledAiProviderList,
+  enabledImageModelList,
   isActiveProviderApiKeyNotEmpty,
   isActiveProviderEndpointNotEmpty,
   isAiProviderConfigLoading,

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

@@ -24,7 +24,7 @@ import { WebBrowsingManifest } from '@/tools/web-browsing';
 import { ChatMessage, CreateMessageParams, SendMessageParams } from '@/types/message';
 import { ChatImageItem } from '@/types/message/image';
 import { MessageSemanticSearchChunk } from '@/types/rag';
-import { setNamespace } from '@/utils/storeDebug';
+import { Action, setNamespace } from '@/utils/storeDebug';
 
 import { chatSelectors, topicSelectors } from '../../../selectors';
 
@@ -104,7 +104,7 @@ export interface AIGenerateAction {
   internal_toggleChatLoading: (
     loading: boolean,
     id?: string,
-    action?: string,
+    action?: Action,
   ) => AbortController | undefined;
   /**
    * Controls the streaming state of tool calling processes, updating the UI accordingly
@@ -383,7 +383,7 @@ export const generateAIChat: StateCreator<
       const abortController = get().internal_toggleChatLoading(
         true,
         assistantId,
-        n('generateMessage(start)', { messageId: assistantId, messages }) as string,
+        n('generateMessage(start)', { messageId: assistantId, messages }),
       );
 
       get().internal_toggleSearchWorkflow(true, assistantId);
@@ -434,7 +434,7 @@ export const generateAIChat: StateCreator<
       get().internal_toggleChatLoading(
         false,
         assistantId,
-        n('generateMessage(start)', { messageId: assistantId, messages }) as string,
+        n('generateMessage(start)', { messageId: assistantId, messages }),
       );
       get().internal_toggleSearchWorkflow(false, assistantId);
 
@@ -502,7 +502,7 @@ export const generateAIChat: StateCreator<
     const abortController = internal_toggleChatLoading(
       true,
       messageId,
-      n('generateMessage(start)', { messageId, messages }) as string,
+      n('generateMessage(start)', { messageId, messages }),
     );
 
     const agentConfig = agentSelectors.currentAgentConfig(getAgentStoreState());

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

@@ -25,7 +25,7 @@ import {
 import { ChatImageItem } from '@/types/message/image';
 import { GroundingSearch } from '@/types/search';
 import { TraceEventPayloads } from '@/types/trace';
-import { setNamespace } from '@/utils/storeDebug';
+import { Action, setNamespace } from '@/utils/storeDebug';
 import { nanoid } from '@/utils/uuid';
 
 import type { ChatStoreState } from '../../initialState';
@@ -130,7 +130,7 @@ export interface ChatMessageAction {
     key: keyof ChatStoreState,
     loading: boolean,
     id?: string,
-    action?: string,
+    action?: Action,
   ) => AbortController | undefined;
 }
 

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

@@ -54,7 +54,7 @@ export const chatTranslate: StateCreator<
     // create translate extra
     await updateMessageTranslate(id, { content: '', from: '', to: targetLang });
 
-    internal_toggleChatLoading(true, id, n('translateMessage(start)', { id }) as string);
+    internal_toggleChatLoading(true, id, n('translateMessage(start)', { id }));
 
     let content = '';
     let from = '';

package/src/store/global/initialState.ts

@@ -10,6 +10,7 @@ export enum SidebarTabKey {
   Chat = 'chat',
   Discover = 'discover',
   Files = 'files',
+  Image = 'image',
   Me = 'me',
   Setting = 'settings',
 }
@@ -50,6 +51,8 @@ export interface SystemStatus {
   filePanelWidth: number;
   hidePWAInstaller?: boolean;
   hideThreadLimitAlert?: boolean;
+  imagePanelWidth: number;
+  imageTopicPanelWidth?: number;
   inputHeight: number;
   /**
    * 应用初始化时不启用 PGLite，只有当用户手动开启时才启用
@@ -65,6 +68,8 @@ export interface SystemStatus {
   showChatSideBar?: boolean;
   showFilePanel?: boolean;
   showHotkeyHelper?: boolean;
+  showImagePanel?: boolean;
+  showImageTopicPanel?: boolean;
   showSessionPanel?: boolean;
   showSystemRole?: boolean;
   systemRoleExpandedMap: Record<string, boolean>;
@@ -104,6 +109,8 @@ export const INITIAL_STATUS = {
   filePanelWidth: 320,
   hidePWAInstaller: false,
   hideThreadLimitAlert: false,
+  imagePanelWidth: 320,
+  imageTopicPanelWidth: 80,
   inputHeight: 200,
   mobileShowTopic: false,
   portalWidth: 400,
@@ -111,6 +118,8 @@ export const INITIAL_STATUS = {
   showChatSideBar: true,
   showFilePanel: true,
   showHotkeyHelper: false,
+  showImagePanel: true,
+  showImageTopicPanel: true,
   showSessionPanel: true,
   showSystemRole: false,
   systemRoleExpandedMap: {},

package/src/store/global/selectors/systemStatus.ts

@@ -14,6 +14,8 @@ const mobileShowPortal = (s: GlobalState) => s.status.mobileShowPortal;
 const showChatSideBar = (s: GlobalState) => !s.status.zenMode && s.status.showChatSideBar;
 const showSessionPanel = (s: GlobalState) => !s.status.zenMode && s.status.showSessionPanel;
 const showFilePanel = (s: GlobalState) => s.status.showFilePanel;
+const showImagePanel = (s: GlobalState) => s.status.showImagePanel;
+const showImageTopicPanel = (s: GlobalState) => s.status.showImageTopicPanel;
 const hidePWAInstaller = (s: GlobalState) => s.status.hidePWAInstaller;
 const isShowCredit = (s: GlobalState) => s.status.isShowCredit;
 const themeMode = (s: GlobalState) => s.status.themeMode || 'auto';
@@ -24,6 +26,8 @@ const inZenMode = (s: GlobalState) => s.status.zenMode;
 const sessionWidth = (s: GlobalState) => s.status.sessionsWidth;
 const portalWidth = (s: GlobalState) => s.status.portalWidth || 400;
 const filePanelWidth = (s: GlobalState) => s.status.filePanelWidth;
+const imagePanelWidth = (s: GlobalState) => s.status.imagePanelWidth;
+const imageTopicPanelWidth = (s: GlobalState) => s.status.imageTopicPanelWidth;
 const inputHeight = (s: GlobalState) => s.status.inputHeight;
 const threadInputHeight = (s: GlobalState) => s.status.threadInputHeight;
 
@@ -62,6 +66,8 @@ export const systemStatusSelectors = {
   filePanelWidth,
   getAgentSystemRoleExpanded,
   hidePWAInstaller,
+  imagePanelWidth,
+  imageTopicPanelWidth,
   inZenMode,
   inputHeight,
   isDBInited,
@@ -78,6 +84,8 @@ export const systemStatusSelectors = {
   showChatHeader,
   showChatSideBar,
   showFilePanel,
+  showImagePanel,
+  showImageTopicPanel,
   showSessionPanel,
   showSystemRole,
   systemStatus,

package/src/store/image/index.ts

@@ -0,0 +1,2 @@
+export type { ImageStore } from './store';
+export { getImageStoreState, useImageStore } from './store';

package/src/store/image/initialState.ts

@@ -0,0 +1,25 @@
+import { CreateImageState, initialCreateImageState } from './slices/createImage/initialState';
+import {
+  GenerationBatchState,
+  initialGenerationBatchState,
+} from './slices/generationBatch/initialState';
+import {
+  GenerationConfigState,
+  initialGenerationConfigState,
+} from './slices/generationConfig/initialState';
+import {
+  GenerationTopicState,
+  initialGenerationTopicState,
+} from './slices/generationTopic/initialState';
+
+export type ImageStoreState = GenerationConfigState &
+  GenerationTopicState &
+  GenerationBatchState &
+  CreateImageState;
+
+export const initialState: ImageStoreState = {
+  ...initialGenerationConfigState,
+  ...initialGenerationTopicState,
+  ...initialGenerationBatchState,
+  ...initialCreateImageState,
+};

package/src/store/image/selectors.ts

@@ -0,0 +1,4 @@
+export * from './slices/createImage/selectors';
+export * from './slices/generationBatch/selectors';
+export * from './slices/generationConfig/selectors';
+export * from './slices/generationTopic/selectors';

package/src/store/image/slices/createImage/action.test.ts

@@ -0,0 +1,330 @@
+import { act, renderHook } from '@testing-library/react';
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { imageService } from '@/services/image';
+import { useImageStore } from '@/store/image';
+
+// Mock external dependencies
+vi.mock('@/services/image', () => ({
+  imageService: {
+    createImage: vi.fn().mockResolvedValue({
+      success: true,
+      data: {
+        batch: {
+          generationTopicId: 'test-topic-id',
+          provider: 'test-provider',
+          model: 'test-model',
+          prompt: 'test prompt',
+          width: 1024,
+          height: 1024,
+          userId: 'test-user',
+          id: 'batch-id',
+          accessedAt: new Date(),
+          createdAt: new Date(),
+          updatedAt: new Date(),
+          ratio: null,
+          config: {},
+        },
+        generations: [],
+      },
+    }),
+  },
+}));
+
+const mockImageService = vi.mocked(imageService);
+
+describe('CreateImageAction', () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    // Reset to initial state with proper defaults
+    const initialState = useImageStore.getState();
+    useImageStore.setState({
+      ...initialState,
+      isCreating: false,
+      isCreatingWithNewTopic: false,
+      activeGenerationTopicId: 'active-topic-id',
+      parameters: { prompt: 'test prompt', width: 1024, height: 1024 },
+      provider: 'test-provider',
+      model: 'test-model',
+      imageNum: 4,
+      generationBatchesMap: {
+        'active-topic-id': [
+          {
+            id: 'batch-id',
+            provider: 'batch-provider',
+            model: 'batch-model',
+            config: { prompt: 'batch prompt' },
+          } as any,
+        ],
+      },
+    });
+  });
+
+  afterEach(() => {
+    vi.restoreAllMocks();
+  });
+
+  describe('createImage', () => {
+    it('should create image with existing topic', async () => {
+      const { result } = renderHook(() => useImageStore());
+      const mockRefreshGenerationBatches = vi.fn().mockResolvedValue(undefined);
+
+      // Set up store state
+      act(() => {
+        useImageStore.setState({
+          refreshGenerationBatches: mockRefreshGenerationBatches,
+        });
+      });
+
+      await act(async () => {
+        await result.current.createImage();
+      });
+
+      // Verify state changes
+      expect(result.current.isCreating).toBe(false);
+      expect(result.current.isCreatingWithNewTopic).toBe(false);
+
+      // Verify service calls
+      expect(mockImageService.createImage).toHaveBeenCalledWith({
+        generationTopicId: 'active-topic-id',
+        provider: 'test-provider',
+        model: 'test-model',
+        imageNum: 4,
+        params: { prompt: 'test prompt', width: 1024, height: 1024 },
+      });
+
+      // Verify refresh was called
+      expect(mockRefreshGenerationBatches).toHaveBeenCalled();
+    });
+
+    it('should create new topic when no active topic exists', async () => {
+      const mockCreateGenerationTopic = vi.fn().mockResolvedValue('new-topic-id');
+      const mockSwitchGenerationTopic = vi.fn();
+      const mockSetTopicBatchLoaded = vi.fn();
+
+      const { result } = renderHook(() => useImageStore());
+
+      act(() => {
+        useImageStore.setState({
+          activeGenerationTopicId: '', // No active topic
+          createGenerationTopic: mockCreateGenerationTopic,
+          switchGenerationTopic: mockSwitchGenerationTopic,
+          setTopicBatchLoaded: mockSetTopicBatchLoaded,
+        });
+      });
+
+      await act(async () => {
+        await result.current.createImage();
+      });
+
+      // Verify state changes
+      expect(result.current.isCreating).toBe(false);
+      expect(result.current.isCreatingWithNewTopic).toBe(false);
+
+      // Verify topic creation
+      expect(mockCreateGenerationTopic).toHaveBeenCalledWith(['test prompt']);
+      expect(mockSetTopicBatchLoaded).toHaveBeenCalledWith('new-topic-id');
+      expect(mockSwitchGenerationTopic).toHaveBeenCalledWith('new-topic-id');
+
+      // Verify service call with new topic id
+      expect(mockImageService.createImage).toHaveBeenCalledWith({
+        generationTopicId: 'new-topic-id',
+        provider: 'test-provider',
+        model: 'test-model',
+        imageNum: 4,
+        params: { prompt: 'test prompt', width: 1024, height: 1024 },
+      });
+    });
+
+    it('should throw error when parameters is not initialized', async () => {
+      const { result } = renderHook(() => useImageStore());
+
+      act(() => {
+        useImageStore.setState({
+          parameters: undefined, // Set parameters to undefined
+        });
+      });
+
+      await expect(
+        act(async () => {
+          await result.current.createImage();
+        }),
+      ).rejects.toThrow('parameters is not initialized');
+    });
+
+    it('should throw error when prompt is empty', async () => {
+      const { result } = renderHook(() => useImageStore());
+
+      act(() => {
+        useImageStore.setState({
+          parameters: {
+            prompt: '', // Empty prompt
+            width: 1024,
+            height: 1024,
+          },
+        });
+      });
+
+      await expect(
+        act(async () => {
+          await result.current.createImage();
+        }),
+      ).rejects.toThrow('prompt is empty');
+    });
+
+    it('should handle service error', async () => {
+      const error = new Error('Service error');
+      mockImageService.createImage.mockRejectedValueOnce(error);
+
+      const mockRefreshGenerationBatches = vi.fn();
+      const { result } = renderHook(() => useImageStore());
+
+      act(() => {
+        useImageStore.setState({
+          refreshGenerationBatches: mockRefreshGenerationBatches,
+        });
+      });
+
+      await expect(
+        act(async () => {
+          await result.current.createImage();
+        }),
+      ).rejects.toThrow('Service error');
+
+      // The service should have been called before the error
+      expect(mockImageService.createImage).toHaveBeenCalled();
+    });
+
+    it('should handle service error with new topic', async () => {
+      const error = new Error('Service error');
+      mockImageService.createImage.mockRejectedValueOnce(error);
+
+      const mockCreateGenerationTopic = vi.fn().mockResolvedValue('new-topic-id');
+      const mockSwitchGenerationTopic = vi.fn();
+      const mockSetTopicBatchLoaded = vi.fn();
+
+      const { result } = renderHook(() => useImageStore());
+
+      act(() => {
+        useImageStore.setState({
+          activeGenerationTopicId: '', // No active topic
+          createGenerationTopic: mockCreateGenerationTopic,
+          switchGenerationTopic: mockSwitchGenerationTopic,
+          setTopicBatchLoaded: mockSetTopicBatchLoaded,
+        });
+      });
+
+      await expect(
+        act(async () => {
+          await result.current.createImage();
+        }),
+      ).rejects.toThrow('Service error');
+
+      // Verify topic was created before the error
+      expect(mockCreateGenerationTopic).toHaveBeenCalled();
+      expect(mockSwitchGenerationTopic).toHaveBeenCalled();
+    });
+  });
+
+  describe('recreateImage', () => {
+    it('should recreate image successfully', async () => {
+      const mockRefreshGenerationBatches = vi.fn().mockResolvedValue(undefined);
+      const mockRemoveGenerationBatch = vi.fn().mockResolvedValue(undefined);
+
+      const { result } = renderHook(() => useImageStore());
+
+      act(() => {
+        useImageStore.setState({
+          refreshGenerationBatches: mockRefreshGenerationBatches,
+          removeGenerationBatch: mockRemoveGenerationBatch,
+        });
+      });
+
+      await act(async () => {
+        await result.current.recreateImage('batch-id');
+      });
+
+      // Verify state changes
+      expect(result.current.isCreating).toBe(false);
+
+      // Verify batch removal
+      expect(mockRemoveGenerationBatch).toHaveBeenCalledWith('batch-id', 'active-topic-id');
+
+      // Verify service call
+      expect(mockImageService.createImage).toHaveBeenCalledWith({
+        generationTopicId: 'active-topic-id',
+        provider: 'batch-provider',
+        model: 'batch-model',
+        imageNum: 4,
+        params: { prompt: 'batch prompt' },
+      });
+
+      // Verify refresh was called
+      expect(mockRefreshGenerationBatches).toHaveBeenCalled();
+    });
+
+    it('should throw error when no active topic', async () => {
+      const { result } = renderHook(() => useImageStore());
+
+      act(() => {
+        useImageStore.setState({
+          activeGenerationTopicId: '', // No active topic
+        });
+      });
+
+      await expect(
+        act(async () => {
+          await result.current.recreateImage('batch-id');
+        }),
+      ).rejects.toThrow('No active generation topic');
+    });
+
+    it('should handle service error', async () => {
+      const error = new Error('Service error');
+      mockImageService.createImage.mockRejectedValueOnce(error);
+
+      const mockRefreshGenerationBatches = vi.fn();
+      const mockRemoveGenerationBatch = vi.fn().mockResolvedValue(undefined);
+
+      const { result } = renderHook(() => useImageStore());
+
+      act(() => {
+        useImageStore.setState({
+          refreshGenerationBatches: mockRefreshGenerationBatches,
+          removeGenerationBatch: mockRemoveGenerationBatch,
+        });
+      });
+
+      await expect(
+        act(async () => {
+          await result.current.recreateImage('batch-id');
+        }),
+      ).rejects.toThrow('Service error');
+
+      // Verify batch was removed before the error
+      expect(mockRemoveGenerationBatch).toHaveBeenCalledWith('batch-id', 'active-topic-id');
+    });
+
+    it('should handle batch removal error', async () => {
+      const error = new Error('Removal error');
+      const mockRemoveGenerationBatch = vi.fn().mockRejectedValueOnce(error);
+
+      const { result } = renderHook(() => useImageStore());
+
+      act(() => {
+        useImageStore.setState({
+          removeGenerationBatch: mockRemoveGenerationBatch,
+        });
+      });
+
+      await expect(
+        act(async () => {
+          await result.current.recreateImage('batch-id');
+        }),
+      ).rejects.toThrow('Removal error');
+
+      // Verify image service was not called after removal error
+      expect(mockImageService.createImage).not.toHaveBeenCalled();
+    });
+  });
+});