@lobehub/chat 1.98.1 → 1.99.0

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

@@ -0,0 +1,307 @@
+import { describe, expect, it } from 'vitest';
+
+import { ImageStore } from '@/store/image';
+import { initialState } from '@/store/image/initialState';
+import { AsyncTaskStatus } from '@/types/asyncTask';
+import { Generation, GenerationBatch } from '@/types/generation';
+import { merge } from '@/utils/merge';
+
+import { generationBatchSelectors } from './selectors';
+
+const initialStore = initialState as ImageStore;
+
+// Mock generation data
+const mockGenerations: Generation[] = [
+  {
+    id: 'gen-1',
+    seed: 12345,
+    createdAt: new Date('2024-01-01'),
+    asyncTaskId: null,
+    asset: {
+      type: 'image',
+      url: 'image1.jpg',
+      width: 1024,
+      height: 1024,
+    },
+    task: {
+      id: 'task-1',
+      status: AsyncTaskStatus.Success,
+    },
+  },
+  {
+    id: 'gen-2',
+    seed: 12346,
+    createdAt: new Date('2024-01-01'),
+    asyncTaskId: null,
+    asset: {
+      type: 'image',
+      url: 'image2.jpg',
+      width: 1024,
+      height: 1024,
+    },
+    task: {
+      id: 'task-2',
+      status: AsyncTaskStatus.Success,
+    },
+  },
+  {
+    id: 'gen-3',
+    seed: 12347,
+    createdAt: new Date('2024-01-02'),
+    asyncTaskId: null,
+    asset: {
+      type: 'image',
+      url: 'image3.jpg',
+      width: 1024,
+      height: 1024,
+    },
+    task: {
+      id: 'task-3',
+      status: AsyncTaskStatus.Success,
+    },
+  },
+];
+
+// Mock data
+const mockGenerationBatches: GenerationBatch[] = [
+  {
+    id: 'batch-1',
+    provider: 'dall-e-3',
+    model: 'dall-e-3',
+    prompt: 'A beautiful sunset',
+    width: 1024,
+    height: 1024,
+    generations: [mockGenerations[0], mockGenerations[1]],
+    config: {
+      prompt: 'A beautiful sunset',
+      width: 1024,
+      height: 1024,
+      size: '1024x1024',
+    },
+    createdAt: new Date('2024-01-01'),
+  },
+  {
+    id: 'batch-2',
+    provider: 'dall-e-3',
+    model: 'dall-e-3',
+    prompt: 'A mountain landscape',
+    width: 1024,
+    height: 1024,
+    generations: [mockGenerations[2]],
+    config: {
+      prompt: 'A mountain landscape',
+      width: 1024,
+      height: 1024,
+      size: '1024x1024',
+    },
+    createdAt: new Date('2024-01-02'),
+  },
+  {
+    id: 'batch-3',
+    provider: 'midjourney',
+    model: 'midjourney-v6',
+    prompt: 'A futuristic city',
+    width: 1024,
+    height: 1024,
+    generations: [],
+    config: {
+      prompt: 'A futuristic city',
+      width: 1024,
+      height: 1024,
+      size: '1024x1024',
+    },
+    createdAt: new Date('2024-01-03'),
+  },
+];
+
+describe('generationBatchSelectors', () => {
+  describe('getGenerationBatchesByTopicId', () => {
+    it('should return batches for a specific topic', () => {
+      const state = merge(initialStore, {
+        generationBatchesMap: {
+          'topic-1': [mockGenerationBatches[0], mockGenerationBatches[1]],
+          'topic-2': [mockGenerationBatches[2]],
+        },
+      });
+
+      const batches = generationBatchSelectors.getGenerationBatchesByTopicId('topic-1')(state);
+      expect(batches).toHaveLength(2);
+      expect(batches[0].id).toBe('batch-1');
+      expect(batches[1].id).toBe('batch-2');
+    });
+
+    it('should return empty array if topic has no batches', () => {
+      const state = merge(initialStore, {
+        generationBatchesMap: {
+          'topic-1': [mockGenerationBatches[0]],
+        },
+      });
+
+      const batches = generationBatchSelectors.getGenerationBatchesByTopicId('topic-2')(state);
+      expect(batches).toEqual([]);
+    });
+
+    it('should return empty array if topic does not exist', () => {
+      const state = merge(initialStore, {
+        generationBatchesMap: {},
+      });
+
+      const batches = generationBatchSelectors.getGenerationBatchesByTopicId('non-existent')(state);
+      expect(batches).toEqual([]);
+    });
+  });
+
+  describe('currentGenerationBatches', () => {
+    it('should return batches for the active topic', () => {
+      const state = merge(initialStore, {
+        activeGenerationTopicId: 'topic-1',
+        generationBatchesMap: {
+          'topic-1': [mockGenerationBatches[0], mockGenerationBatches[1]],
+          'topic-2': [mockGenerationBatches[2]],
+        },
+      });
+
+      const batches = generationBatchSelectors.currentGenerationBatches(state);
+      expect(batches).toHaveLength(2);
+      expect(batches).toEqual([mockGenerationBatches[0], mockGenerationBatches[1]]);
+    });
+
+    it('should return empty array if no active topic', () => {
+      const state = merge(initialStore, {
+        activeGenerationTopicId: undefined,
+        generationBatchesMap: {
+          'topic-1': [mockGenerationBatches[0]],
+        },
+      });
+
+      const batches = generationBatchSelectors.currentGenerationBatches(state);
+      expect(batches).toEqual([]);
+    });
+
+    it('should return empty array if active topic has no batches', () => {
+      const state = merge(initialStore, {
+        activeGenerationTopicId: 'topic-3',
+        generationBatchesMap: {
+          'topic-1': [mockGenerationBatches[0]],
+        },
+      });
+
+      const batches = generationBatchSelectors.currentGenerationBatches(state);
+      expect(batches).toEqual([]);
+    });
+  });
+
+  describe('getGenerationBatchByBatchId', () => {
+    it('should return the batch with matching id', () => {
+      const state = merge(initialStore, {
+        activeGenerationTopicId: 'topic-1',
+        generationBatchesMap: {
+          'topic-1': [mockGenerationBatches[0], mockGenerationBatches[1]],
+        },
+      });
+
+      const batch = generationBatchSelectors.getGenerationBatchByBatchId('batch-2')(state);
+      expect(batch).toEqual(mockGenerationBatches[1]);
+    });
+
+    it('should return undefined if batch not found', () => {
+      const state = merge(initialStore, {
+        activeGenerationTopicId: 'topic-1',
+        generationBatchesMap: {
+          'topic-1': [mockGenerationBatches[0]],
+        },
+      });
+
+      const batch = generationBatchSelectors.getGenerationBatchByBatchId('non-existent')(state);
+      expect(batch).toBeUndefined();
+    });
+
+    it('should return undefined if no active topic', () => {
+      const state = merge(initialStore, {
+        activeGenerationTopicId: undefined,
+        generationBatchesMap: {
+          'topic-1': [mockGenerationBatches[0]],
+        },
+      });
+
+      const batch = generationBatchSelectors.getGenerationBatchByBatchId('batch-1')(state);
+      expect(batch).toBeUndefined();
+    });
+
+    it('should only search in current topic batches', () => {
+      const state = merge(initialStore, {
+        activeGenerationTopicId: 'topic-1',
+        generationBatchesMap: {
+          'topic-1': [mockGenerationBatches[0], mockGenerationBatches[1]],
+          'topic-2': [mockGenerationBatches[2]],
+        },
+      });
+
+      // batch-3 belongs to topic-2, but active topic is topic-1
+      const batch = generationBatchSelectors.getGenerationBatchByBatchId('batch-3')(state);
+      expect(batch).toBeUndefined();
+    });
+  });
+
+  describe('isCurrentGenerationTopicLoaded', () => {
+    it('should return true if current topic has batches array', () => {
+      const state = merge(initialStore, {
+        activeGenerationTopicId: 'topic-1',
+        generationBatchesMap: {
+          'topic-1': [mockGenerationBatches[0]],
+        },
+      });
+
+      const isLoaded = generationBatchSelectors.isCurrentGenerationTopicLoaded(state);
+      expect(isLoaded).toBe(true);
+    });
+
+    it('should return true even if batches array is empty', () => {
+      const state = merge(initialStore, {
+        activeGenerationTopicId: 'topic-1',
+        generationBatchesMap: {
+          'topic-1': [],
+        },
+      });
+
+      const isLoaded = generationBatchSelectors.isCurrentGenerationTopicLoaded(state);
+      expect(isLoaded).toBe(true);
+    });
+
+    it('should return false if current topic has no batches entry', () => {
+      const state = merge(initialStore, {
+        activeGenerationTopicId: 'topic-2',
+        generationBatchesMap: {
+          'topic-1': [mockGenerationBatches[0]],
+        },
+      });
+
+      const isLoaded = generationBatchSelectors.isCurrentGenerationTopicLoaded(state);
+      expect(isLoaded).toBe(false);
+    });
+
+    it('should return false if no active topic', () => {
+      const state = merge(initialStore, {
+        activeGenerationTopicId: undefined,
+        generationBatchesMap: {
+          'topic-1': [mockGenerationBatches[0]],
+        },
+      });
+
+      const isLoaded = generationBatchSelectors.isCurrentGenerationTopicLoaded(state);
+      expect(isLoaded).toBe(false);
+    });
+
+    it('should return false if batches map has non-array value', () => {
+      const state = merge(initialStore, {
+        activeGenerationTopicId: 'topic-1',
+        generationBatchesMap: {
+          'topic-1': null as any, // Simulating incorrect data type
+        },
+      });
+
+      const isLoaded = generationBatchSelectors.isCurrentGenerationTopicLoaded(state);
+      expect(isLoaded).toBe(false);
+    });
+  });
+});

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

@@ -0,0 +1,36 @@
+import { GenerationBatch } from '@/types/generation';
+
+import { ImageStoreState } from '../../initialState';
+import { generationTopicSelectors } from '../generationTopic/selectors';
+
+// ====== topic batch selectors ====== //
+
+const getGenerationBatchesByTopicId = (topicId: string) => (s: ImageStoreState) => {
+  return s.generationBatchesMap[topicId] || [];
+};
+
+const currentGenerationBatches = (s: ImageStoreState): GenerationBatch[] => {
+  const activeTopicId = generationTopicSelectors.activeGenerationTopicId(s);
+  if (!activeTopicId) return [];
+  return getGenerationBatchesByTopicId(activeTopicId)(s);
+};
+
+const getGenerationBatchByBatchId = (batchId: string) => (s: ImageStoreState) => {
+  const batches = currentGenerationBatches(s);
+  return batches.find((batch) => batch.id === batchId);
+};
+
+const isCurrentGenerationTopicLoaded = (s: ImageStoreState): boolean => {
+  const activeTopicId = generationTopicSelectors.activeGenerationTopicId(s);
+  if (!activeTopicId) return false;
+  return Array.isArray(s.generationBatchesMap[activeTopicId]);
+};
+
+// ====== aggregate selectors ====== //
+
+export const generationBatchSelectors = {
+  getGenerationBatchesByTopicId,
+  currentGenerationBatches,
+  getGenerationBatchByBatchId,
+  isCurrentGenerationTopicLoaded,
+};

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

@@ -0,0 +1,351 @@
+import { act, renderHook } from '@testing-library/react';
+import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
+
+import { fluxSchnellParamsSchema } from '@/config/paramsSchemas/fal/flux-schnell';
+import {
+  ModelParamsSchema,
+  RuntimeImageGenParams,
+  extractDefaultValues,
+} from '@/libs/standard-parameters/meta-schema';
+import { useImageStore } from '@/store/image';
+import { AIImageModelCard } from '@/types/aiModel';
+
+// Mock external dependencies
+vi.mock('@/store/aiInfra', () => ({
+  aiProviderSelectors: {
+    enabledImageModelList: vi.fn(() => [
+      {
+        id: 'fal',
+        name: 'Fal',
+        children: [
+          {
+            id: 'flux/schnell',
+            displayName: 'FLUX.1 Schnell',
+            type: 'image',
+            parameters: fluxSchnellParamsSchema,
+            releasedAt: '2024-08-01',
+          } as AIImageModelCard,
+        ],
+      },
+      {
+        id: 'custom-provider',
+        name: 'Custom Provider',
+        children: [
+          {
+            id: 'custom-model',
+            displayName: 'Custom Model',
+            type: 'image',
+            parameters: {
+              prompt: { default: '' },
+              width: { default: 1024, min: 256, max: 2048, step: 64 },
+              height: { default: 1024, min: 256, max: 2048, step: 64 },
+              steps: { default: 20, min: 1, max: 50 },
+            } as ModelParamsSchema,
+            releasedAt: '2024-01-01',
+          } as AIImageModelCard,
+        ],
+      },
+    ]),
+  },
+  getAiInfraStoreState: vi.fn(() => ({})),
+}));
+
+const fluxSchnellDefaultValues = extractDefaultValues(fluxSchnellParamsSchema);
+
+const customModelSchema: ModelParamsSchema = {
+  prompt: { default: '' },
+  width: { default: 1024, min: 256, max: 2048, step: 64 },
+  height: { default: 1024, min: 256, max: 2048, step: 64 },
+  steps: { default: 20, min: 1, max: 50 },
+};
+
+beforeEach(() => {
+  vi.clearAllMocks();
+
+  // Reset store state
+  useImageStore.setState({
+    model: 'initial-model',
+    provider: 'initial-provider',
+    imageNum: 1,
+    parameters: {
+      prompt: 'initial prompt',
+      width: 512,
+      height: 512,
+    } as RuntimeImageGenParams,
+    parametersSchema: {
+      prompt: { default: '' },
+      width: { default: 512, min: 256, max: 1024 },
+      height: { default: 512, min: 256, max: 1024 },
+    },
+  });
+});
+
+afterEach(() => {
+  vi.restoreAllMocks();
+});
+
+describe('GenerationConfigAction', () => {
+  describe('setParamOnInput', () => {
+    it('should update a single parameter in the parameters object', async () => {
+      const { result } = renderHook(() => useImageStore());
+
+      act(() => {
+        result.current.setParamOnInput('prompt', 'new test prompt');
+      });
+
+      expect(result.current.parameters).toMatchObject({
+        prompt: 'new test prompt',
+        width: 512,
+        height: 512,
+      });
+    });
+
+    it('should update numeric parameters correctly', async () => {
+      const { result } = renderHook(() => useImageStore());
+
+      act(() => {
+        result.current.setParamOnInput('width', 2048);
+      });
+
+      expect(result.current.parameters).toMatchObject({
+        prompt: 'initial prompt',
+        width: 2048,
+        height: 512,
+      });
+    });
+
+    it('should handle null values correctly', async () => {
+      const { result } = renderHook(() => useImageStore());
+
+      act(() => {
+        result.current.setParamOnInput('seed', null);
+      });
+
+      expect(result.current.parameters?.seed).toBeNull();
+    });
+
+    it('should handle array values correctly', async () => {
+      const { result } = renderHook(() => useImageStore());
+
+      act(() => {
+        result.current.setParamOnInput('imageUrls', ['test1.jpg', 'test2.jpg']);
+      });
+
+      expect(result.current.parameters?.imageUrls).toEqual(['test1.jpg', 'test2.jpg']);
+    });
+  });
+
+  describe('setModelAndProviderOnSelect', () => {
+    it('should set model, provider, parameters and parametersSchema for flux/schnell', async () => {
+      const { result } = renderHook(() => useImageStore());
+
+      act(() => {
+        result.current.setModelAndProviderOnSelect('flux/schnell', 'fal');
+      });
+
+      expect(result.current.model).toBe('flux/schnell');
+      expect(result.current.provider).toBe('fal');
+      expect(result.current.parameters).toEqual(fluxSchnellDefaultValues);
+      expect(result.current.parametersSchema).toEqual(fluxSchnellParamsSchema);
+    });
+
+    it('should handle model selection with custom parameters', async () => {
+      const { result } = renderHook(() => useImageStore());
+
+      act(() => {
+        result.current.setModelAndProviderOnSelect('custom-model', 'custom-provider');
+      });
+
+      const expectedParams = extractDefaultValues(customModelSchema);
+
+      expect(result.current.model).toBe('custom-model');
+      expect(result.current.provider).toBe('custom-provider');
+      expect(result.current.parameters).toEqual(expectedParams);
+      expect(result.current.parametersSchema).toEqual(customModelSchema);
+    });
+
+    it('should replace all previous parameters with new model defaults', async () => {
+      const { result } = renderHook(() => useImageStore());
+
+      // First set some custom parameters
+      act(() => {
+        result.current.setParamOnInput('prompt', 'custom prompt');
+        result.current.setParamOnInput('steps', 50);
+      });
+
+      // Then switch model
+      act(() => {
+        result.current.setModelAndProviderOnSelect('flux/schnell', 'fal');
+      });
+
+      // Should completely replace parameters with model defaults
+      expect(result.current.parameters).toEqual(fluxSchnellDefaultValues);
+      expect(result.current.parameters?.prompt).toBe(''); // Default value, not 'custom prompt'
+    });
+  });
+
+  describe('setImageNum', () => {
+    it('should update the imageNum value', async () => {
+      const { result } = renderHook(() => useImageStore());
+
+      act(() => {
+        result.current.setImageNum(4);
+      });
+
+      expect(result.current.imageNum).toBe(4);
+    });
+
+    it('should handle different imageNum values', async () => {
+      const { result } = renderHook(() => useImageStore());
+
+      act(() => {
+        result.current.setImageNum(8);
+      });
+
+      expect(result.current.imageNum).toBe(8);
+
+      act(() => {
+        result.current.setImageNum(1);
+      });
+
+      expect(result.current.imageNum).toBe(1);
+    });
+
+    it('should handle edge case values', async () => {
+      const { result } = renderHook(() => useImageStore());
+
+      act(() => {
+        result.current.setImageNum(0);
+      });
+
+      expect(result.current.imageNum).toBe(0);
+    });
+  });
+
+  describe('reuseSettings', () => {
+    it('should set model, provider and merge settings with default values', async () => {
+      const { result } = renderHook(() => useImageStore());
+      const customSettings: Partial<RuntimeImageGenParams> = {
+        prompt: 'custom prompt',
+        steps: 8,
+        seed: 54321,
+      };
+
+      act(() => {
+        result.current.reuseSettings('flux/schnell', 'fal', customSettings);
+      });
+
+      expect(result.current.model).toBe('flux/schnell');
+      expect(result.current.provider).toBe('fal');
+      expect(result.current.parameters).toEqual({
+        ...fluxSchnellDefaultValues,
+        ...customSettings,
+      });
+      expect(result.current.parametersSchema).toEqual(fluxSchnellParamsSchema);
+    });
+
+    it('should override default values with provided settings', async () => {
+      const { result } = renderHook(() => useImageStore());
+      const customSettings: Partial<RuntimeImageGenParams> = {
+        width: 1536,
+        height: 1536,
+      };
+
+      act(() => {
+        result.current.reuseSettings('flux/schnell', 'fal', customSettings);
+      });
+
+      expect(result.current.parameters).toEqual({
+        ...fluxSchnellDefaultValues,
+        width: 1536,
+        height: 1536,
+      });
+    });
+
+    it('should handle empty settings object', async () => {
+      const { result } = renderHook(() => useImageStore());
+
+      act(() => {
+        result.current.reuseSettings('flux/schnell', 'fal', {});
+      });
+
+      expect(result.current.parameters).toEqual(fluxSchnellDefaultValues);
+    });
+
+    it('should handle partial settings with null values', async () => {
+      const { result } = renderHook(() => useImageStore());
+      const customSettings: Partial<RuntimeImageGenParams> = {
+        seed: null,
+        imageUrl: null,
+      };
+
+      act(() => {
+        result.current.reuseSettings('flux/schnell', 'fal', customSettings);
+      });
+
+      expect(result.current.parameters?.seed).toBeNull();
+      expect(result.current.parameters?.imageUrl).toBeNull();
+    });
+  });
+
+  describe('reuseSeed', () => {
+    it('should update only the seed parameter', async () => {
+      const { result } = renderHook(() => useImageStore());
+      const newSeed = 98765;
+
+      act(() => {
+        result.current.reuseSeed(newSeed);
+      });
+
+      expect(result.current.parameters).toMatchObject({
+        prompt: 'initial prompt',
+        width: 512,
+        height: 512,
+        seed: newSeed,
+      });
+    });
+
+    it('should preserve other parameters when updating seed', async () => {
+      const { result } = renderHook(() => useImageStore());
+
+      // First set some parameters
+      act(() => {
+        result.current.setParamOnInput('prompt', 'test prompt');
+        result.current.setParamOnInput('width', 1024);
+      });
+
+      const newSeed = 11111;
+      act(() => {
+        result.current.reuseSeed(newSeed);
+      });
+
+      expect(result.current.parameters).toMatchObject({
+        prompt: 'test prompt',
+        width: 1024,
+        height: 512,
+        seed: newSeed,
+      });
+    });
+
+    it('should handle seed value of 0', async () => {
+      const { result } = renderHook(() => useImageStore());
+
+      act(() => {
+        result.current.reuseSeed(0);
+      });
+
+      expect(result.current.parameters?.seed).toBe(0);
+    });
+
+    it('should handle large seed values within range', async () => {
+      const { result } = renderHook(() => useImageStore());
+      const largeSeed = 2147483647; // MAX_SEED
+
+      act(() => {
+        result.current.reuseSeed(largeSeed);
+      });
+
+      expect(result.current.parameters?.seed).toBe(largeSeed);
+    });
+  });
+});