@lobehub/chat 1.98.2 → 1.99.0

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

@@ -0,0 +1,134 @@
+import { StateCreator } from 'zustand';
+
+import { imageService } from '@/services/image';
+
+import { ImageStore } from '../../store';
+import { generationBatchSelectors } from '../generationBatch/selectors';
+import { imageGenerationConfigSelectors } from '../generationConfig/selectors';
+import { generationTopicSelectors } from '../generationTopic';
+
+// ====== action interface ====== //
+
+export interface CreateImageAction {
+  createImage: () => Promise<void>;
+  /**
+   * eg: invalid api key, recreate image
+   */
+  recreateImage: (generationBatchId: string) => Promise<void>;
+}
+
+// ====== helper functions ====== //
+
+// ====== action implementation ====== //
+
+export const createCreateImageSlice: StateCreator<
+  ImageStore,
+  [['zustand/devtools', never]],
+  [],
+  CreateImageAction
+> = (set, get) => ({
+  async createImage() {
+    set({ isCreating: true }, false, 'createImage/startCreateImage');
+
+    const store = get();
+    const imageNum = imageGenerationConfigSelectors.imageNum(store);
+    const parameters = imageGenerationConfigSelectors.parameters(store);
+    const provider = imageGenerationConfigSelectors.provider(store);
+    const model = imageGenerationConfigSelectors.model(store);
+    const activeGenerationTopicId = generationTopicSelectors.activeGenerationTopicId(store);
+    const { createGenerationTopic, switchGenerationTopic, setTopicBatchLoaded } = store;
+
+    if (!parameters) {
+      throw new TypeError('parameters is not initialized');
+    }
+
+    if (!parameters.prompt) {
+      throw new TypeError('prompt is empty');
+    }
+
+    // Track the final topic ID to use for image creation
+    let finalTopicId = activeGenerationTopicId;
+
+    // 1. Create generation topic if not exists
+    let generationTopicId = activeGenerationTopicId;
+    let isNewTopic = false;
+
+    if (!generationTopicId) {
+      isNewTopic = true;
+      const prompts = [parameters.prompt];
+      const newGenerationTopicId = await createGenerationTopic(prompts);
+      finalTopicId = newGenerationTopicId;
+
+      // 2. Initialize empty batch array to avoid skeleton screen
+      setTopicBatchLoaded(newGenerationTopicId);
+
+      // 3. Switch to the new topic (now it has empty data, so no skeleton screen)
+      switchGenerationTopic(newGenerationTopicId);
+    }
+
+    try {
+      // 4. If it's a new topic, set the creating state after topic creation
+      if (isNewTopic) {
+        set({ isCreatingWithNewTopic: true }, false, 'createImage/startCreateImageWithNewTopic');
+      }
+
+      // 5. Create image via service
+      await imageService.createImage({
+        generationTopicId: finalTopicId!,
+        provider,
+        model,
+        imageNum,
+        params: parameters as any,
+      });
+
+      // 6. Only refresh generation batches if it's not a new topic
+      if (!isNewTopic) {
+        await get().refreshGenerationBatches();
+      }
+    } finally {
+      // 7. Reset all creating states
+      if (isNewTopic) {
+        set(
+          { isCreating: false, isCreatingWithNewTopic: false },
+          false,
+          'createImage/endCreateImageWithNewTopic',
+        );
+      } else {
+        set({ isCreating: false }, false, 'createImage/endCreateImage');
+      }
+    }
+  },
+
+  async recreateImage(generationBatchId: string) {
+    set({ isCreating: true }, false, 'recreateImage/startCreateImage');
+
+    const store = get();
+    const imageNum = imageGenerationConfigSelectors.imageNum(store);
+    const activeGenerationTopicId = generationTopicSelectors.activeGenerationTopicId(store);
+    const batch = generationBatchSelectors.getGenerationBatchByBatchId(generationBatchId)(store)!;
+    const { removeGenerationBatch } = store;
+
+    if (!activeGenerationTopicId) {
+      throw new Error('No active generation topic');
+    }
+
+    try {
+      // 1. Delete generation batch
+      await removeGenerationBatch(generationBatchId, activeGenerationTopicId);
+
+      // 2. Create image via service
+      await imageService.createImage({
+        generationTopicId: activeGenerationTopicId,
+        provider: batch.provider,
+        model: batch.model,
+        imageNum,
+        params: batch.config as any,
+      });
+
+      // 3. Refresh generation batches to show the real data
+      await store.refreshGenerationBatches();
+    } finally {
+      set({ isCreating: false }, false, 'recreateImage/endCreateImage');
+    }
+  },
+});

package/src/store/image/slices/createImage/initialState.ts

@@ -0,0 +1,9 @@
+export interface CreateImageState {
+  isCreating: boolean;
+  isCreatingWithNewTopic: boolean;
+}
+
+export const initialCreateImageState: CreateImageState = {
+  isCreating: false,
+  isCreatingWithNewTopic: false,
+};

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

@@ -0,0 +1,114 @@
+import { describe, expect, it } from 'vitest';
+
+import { merge } from '@/utils/merge';
+
+import { ImageStore } from '../../store';
+import { initialCreateImageState } from './initialState';
+import { createImageSelectors } from './selectors';
+
+// 创建一个最小的 ImageStore 模拟对象
+const createMockImageStore = (overrides?: Partial<ImageStore>): ImageStore => {
+  return merge(
+    {
+      ...initialCreateImageState,
+      // 其他必要的初始状态
+    } as ImageStore,
+    overrides || {},
+  );
+};
+
+describe('createImageSelectors', () => {
+  describe('isCreating', () => {
+    it('should return false from initial state', () => {
+      const state = createMockImageStore();
+
+      const result = createImageSelectors.isCreating(state);
+
+      expect(result).toBe(false);
+    });
+
+    it('should return true when isCreating is true', () => {
+      const state = createMockImageStore({ isCreating: true });
+
+      const result = createImageSelectors.isCreating(state);
+
+      expect(result).toBe(true);
+    });
+
+    it('should return false when isCreating is explicitly set to false', () => {
+      const state = createMockImageStore({ isCreating: false });
+
+      const result = createImageSelectors.isCreating(state);
+
+      expect(result).toBe(false);
+    });
+  });
+
+  describe('isCreatingWithNewTopic', () => {
+    it('should return false from initial state', () => {
+      const state = createMockImageStore();
+
+      const result = createImageSelectors.isCreatingWithNewTopic(state);
+
+      expect(result).toBe(false);
+    });
+
+    it('should return true when isCreatingWithNewTopic is true', () => {
+      const state = createMockImageStore({ isCreatingWithNewTopic: true });
+
+      const result = createImageSelectors.isCreatingWithNewTopic(state);
+
+      expect(result).toBe(true);
+    });
+
+    it('should return false when isCreatingWithNewTopic is explicitly set to false', () => {
+      const state = createMockImageStore({ isCreatingWithNewTopic: false });
+
+      const result = createImageSelectors.isCreatingWithNewTopic(state);
+
+      expect(result).toBe(false);
+    });
+  });
+
+  describe('combined states', () => {
+    it('should handle both states being true', () => {
+      const state = createMockImageStore({
+        isCreating: true,
+        isCreatingWithNewTopic: true,
+      });
+
+      expect(createImageSelectors.isCreating(state)).toBe(true);
+      expect(createImageSelectors.isCreatingWithNewTopic(state)).toBe(true);
+    });
+
+    it('should handle mixed states', () => {
+      const state = createMockImageStore({
+        isCreating: true,
+        isCreatingWithNewTopic: false,
+      });
+
+      expect(createImageSelectors.isCreating(state)).toBe(true);
+      expect(createImageSelectors.isCreatingWithNewTopic(state)).toBe(false);
+    });
+
+    it('should handle opposite mixed states', () => {
+      const state = createMockImageStore({
+        isCreating: false,
+        isCreatingWithNewTopic: true,
+      });
+
+      expect(createImageSelectors.isCreating(state)).toBe(false);
+      expect(createImageSelectors.isCreatingWithNewTopic(state)).toBe(true);
+    });
+
+    it('should handle both states being false', () => {
+      const state = createMockImageStore({
+        isCreating: false,
+        isCreatingWithNewTopic: false,
+      });
+
+      expect(createImageSelectors.isCreating(state)).toBe(false);
+      expect(createImageSelectors.isCreatingWithNewTopic(state)).toBe(false);
+    });
+  });
+});

package/src/store/image/slices/createImage/selectors.ts

@@ -0,0 +1,9 @@
+import { ImageStore } from '../../store';
+
+const isCreating = (state: ImageStore) => state.isCreating;
+const isCreatingWithNewTopic = (state: ImageStore) => state.isCreatingWithNewTopic;
+
+export const createImageSelectors = {
+  isCreating,
+  isCreatingWithNewTopic,
+};