@lobehub/chat 1.55.0 → 1.55.1

package/src/server/routers/lambda/aiModel.test.ts

@@ -0,0 +1,240 @@
+import { describe, expect, it, vi } from 'vitest';
+
+import { AiInfraRepos } from '@/database/repositories/aiInfra';
+import { AiModelModel } from '@/database/server/models/aiModel';
+import { UserModel } from '@/database/server/models/user';
+
+import { aiModelRouter } from './aiModel';
+
+vi.mock('@/database/server/models/aiModel');
+vi.mock('@/database/server/models/user');
+vi.mock('@/database/repositories/aiInfra');
+vi.mock('@/server/globalConfig', () => ({
+  getServerGlobalConfig: vi.fn().mockReturnValue({
+    aiProvider: {},
+  }),
+}));
+vi.mock('@/server/modules/KeyVaultsEncrypt', () => ({
+  KeyVaultsGateKeeper: {
+    initWithEnvKey: vi.fn().mockResolvedValue({
+      encrypt: vi.fn(),
+      decrypt: vi.fn(),
+    }),
+  },
+}));
+
+describe('aiModelRouter', () => {
+  const mockCtx = {
+    userId: 'test-user',
+  };
+
+  it('should create ai model', async () => {
+    const mockCreate = vi.fn().mockResolvedValue({ id: 'model-1' });
+    vi.mocked(AiModelModel).mockImplementation(
+      () =>
+        ({
+          create: mockCreate,
+        }) as any,
+    );
+
+    const caller = aiModelRouter.createCaller(mockCtx);
+
+    const result = await caller.createAiModel({
+      id: 'test-model',
+      providerId: 'test-provider',
+    });
+
+    expect(result).toBe('model-1');
+    expect(mockCreate).toHaveBeenCalledWith({
+      id: 'test-model',
+      providerId: 'test-provider',
+    });
+  });
+
+  it('should get ai model by id', async () => {
+    const mockModel = {
+      id: 'model-1',
+      name: 'Test Model',
+    };
+    const mockFindById = vi.fn().mockResolvedValue(mockModel);
+    vi.mocked(AiModelModel).mockImplementation(
+      () =>
+        ({
+          findById: mockFindById,
+        }) as any,
+    );
+
+    const caller = aiModelRouter.createCaller(mockCtx);
+
+    const result = await caller.getAiModelById({ id: 'model-1' });
+
+    expect(result).toEqual(mockModel);
+    expect(mockFindById).toHaveBeenCalledWith('model-1');
+  });
+
+  it('should get ai provider model list', async () => {
+    const mockModelList = [
+      { id: 'model-1', name: 'Model 1' },
+      { id: 'model-2', name: 'Model 2' },
+    ];
+    const mockGetList = vi.fn().mockResolvedValue(mockModelList);
+    vi.mocked(AiInfraRepos).mockImplementation(
+      () =>
+        ({
+          getAiProviderModelList: mockGetList,
+        }) as any,
+    );
+
+    const caller = aiModelRouter.createCaller(mockCtx);
+
+    const result = await caller.getAiProviderModelList({ id: 'provider-1' });
+
+    expect(result).toEqual(mockModelList);
+    expect(mockGetList).toHaveBeenCalledWith('provider-1');
+  });
+
+  it('should remove ai model', async () => {
+    const mockDelete = vi.fn().mockResolvedValue(true);
+    vi.mocked(AiModelModel).mockImplementation(
+      () =>
+        ({
+          delete: mockDelete,
+        }) as any,
+    );
+
+    const caller = aiModelRouter.createCaller(mockCtx);
+
+    await caller.removeAiModel({
+      id: 'model-1',
+      providerId: 'provider-1',
+    });
+
+    expect(mockDelete).toHaveBeenCalledWith('model-1', 'provider-1');
+  });
+
+  it('should update ai model', async () => {
+    const mockUpdate = vi.fn().mockResolvedValue(true);
+    vi.mocked(AiModelModel).mockImplementation(
+      () =>
+        ({
+          update: mockUpdate,
+        }) as any,
+    );
+
+    const caller = aiModelRouter.createCaller(mockCtx);
+
+    await caller.updateAiModel({
+      id: 'model-1',
+      providerId: 'provider-1',
+      value: {
+        displayName: 'Updated Model',
+      },
+    });
+
+    expect(mockUpdate).toHaveBeenCalledWith('model-1', 'provider-1', {
+      displayName: 'Updated Model',
+    });
+  });
+
+  it('should toggle model enabled status', async () => {
+    const mockToggle = vi.fn().mockResolvedValue(true);
+    vi.mocked(AiModelModel).mockImplementation(
+      () =>
+        ({
+          toggleModelEnabled: mockToggle,
+        }) as any,
+    );
+
+    const caller = aiModelRouter.createCaller(mockCtx);
+
+    await caller.toggleModelEnabled({
+      id: 'model-1',
+      providerId: 'provider-1',
+      enabled: true,
+    });
+
+    expect(mockToggle).toHaveBeenCalledWith({
+      id: 'model-1',
+      providerId: 'provider-1',
+      enabled: true,
+    });
+  });
+
+  it('should batch toggle ai models', async () => {
+    const mockBatchToggle = vi.fn().mockResolvedValue(true);
+    vi.mocked(AiModelModel).mockImplementation(
+      () =>
+        ({
+          batchToggleAiModels: mockBatchToggle,
+        }) as any,
+    );
+
+    const caller = aiModelRouter.createCaller(mockCtx);
+
+    await caller.batchToggleAiModels({
+      id: 'provider-1',
+      models: ['model-1', 'model-2'],
+      enabled: true,
+    });
+
+    expect(mockBatchToggle).toHaveBeenCalledWith('provider-1', ['model-1', 'model-2'], true);
+  });
+
+  it('should batch update ai models', async () => {
+    const mockBatchUpdate = vi.fn().mockResolvedValue([]);
+    vi.mocked(AiModelModel).mockImplementation(
+      () =>
+        ({
+          batchUpdateAiModels: mockBatchUpdate,
+        }) as any,
+    );
+
+    const caller = aiModelRouter.createCaller(mockCtx);
+
+    await caller.batchUpdateAiModels({
+      id: 'provider-1',
+      models: [{ id: 'model-1' }, { id: 'model-2' }],
+    });
+
+    expect(mockBatchUpdate).toHaveBeenCalledWith('provider-1', [
+      { id: 'model-1' },
+      { id: 'model-2' },
+    ]);
+  });
+
+  it('should clear models by provider', async () => {
+    const mockClear = vi.fn().mockResolvedValue(true);
+    vi.mocked(AiModelModel).mockImplementation(
+      () =>
+        ({
+          clearModelsByProvider: mockClear,
+        }) as any,
+    );
+
+    const caller = aiModelRouter.createCaller(mockCtx);
+
+    await caller.clearModelsByProvider({
+      providerId: 'provider-1',
+    });
+
+    expect(mockClear).toHaveBeenCalledWith('provider-1');
+  });
+
+  it('should clear remote models', async () => {
+    const mockClearRemote = vi.fn().mockResolvedValue(true);
+    vi.mocked(AiModelModel).mockImplementation(
+      () =>
+        ({
+          clearRemoteModels: mockClearRemote,
+        }) as any,
+    );
+
+    const caller = aiModelRouter.createCaller(mockCtx);
+
+    await caller.clearRemoteModels({
+      providerId: 'provider-1',
+    });
+
+    expect(mockClearRemote).toHaveBeenCalledWith('provider-1');
+  });
+});

package/src/store/aiInfra/slices/aiModel/selectors.test.ts

@@ -0,0 +1,228 @@
+import { describe, expect, it } from 'vitest';
+
+import { AIProviderStoreState } from '@/store/aiInfra/initialState';
+import { AiModelSourceEnum } from '@/types/aiModel';
+
+import { aiModelSelectors } from './selectors';
+
+describe('aiModelSelectors', () => {
+  const mockState: AIProviderStoreState = {
+    aiProviderModelList: [
+      {
+        id: 'model1',
+        type: 'chat',
+        enabled: true,
+        displayName: 'Model One',
+      },
+      {
+        id: 'model2',
+        type: 'chat',
+        enabled: false,
+        displayName: 'Model Two',
+      },
+      {
+        id: 'model3',
+        type: 'embedding',
+        enabled: true,
+        displayName: 'Model Three',
+      },
+      {
+        id: 'model4',
+        type: 'chat',
+        enabled: true,
+        source: AiModelSourceEnum.Remote,
+        displayName: 'Remote Model',
+      },
+    ],
+    modelSearchKeyword: '',
+    aiModelLoadingIds: ['model2'],
+    enabledAiModels: [
+      {
+        id: 'model1',
+        providerId: 'provider1',
+        abilities: {
+          functionCall: true,
+          vision: true,
+          reasoning: true,
+        },
+        contextWindowTokens: 4000,
+        type: 'chat',
+      },
+      {
+        id: 'model4',
+        providerId: 'provider2',
+        abilities: {
+          functionCall: false,
+          vision: false,
+          reasoning: false,
+        },
+        type: 'chat',
+      },
+    ],
+    activeProviderModelList: [],
+    aiProviderConfigUpdatingIds: [],
+    aiProviderList: [],
+    aiProviderLoadingIds: [],
+    providerSearchKeyword: '',
+    aiProviderRuntimeConfig: {},
+    initAiProviderList: false,
+  };
+
+  describe('aiProviderChatModelListIds', () => {
+    it('should return ids of chat type models', () => {
+      const result = aiModelSelectors.aiProviderChatModelListIds(mockState);
+      expect(result).toEqual(['model1', 'model2', 'model4']);
+    });
+  });
+
+  describe('enabledAiProviderModelList', () => {
+    it('should return enabled models', () => {
+      const result = aiModelSelectors.enabledAiProviderModelList(mockState);
+      expect(result).toHaveLength(3);
+      expect(result.map((m) => m.id)).toEqual(['model1', 'model3', 'model4']);
+    });
+  });
+
+  describe('disabledAiProviderModelList', () => {
+    it('should return disabled models', () => {
+      const result = aiModelSelectors.disabledAiProviderModelList(mockState);
+      expect(result).toHaveLength(1);
+      expect(result[0].id).toBe('model2');
+    });
+  });
+
+  describe('filteredAiProviderModelList', () => {
+    it('should filter models by id', () => {
+      const state = { ...mockState, modelSearchKeyword: 'model1' };
+      const result = aiModelSelectors.filteredAiProviderModelList(state);
+      expect(result).toHaveLength(1);
+      expect(result[0].id).toBe('model1');
+    });
+
+    it('should filter models by display name', () => {
+      const state = { ...mockState, modelSearchKeyword: 'remote' };
+      const result = aiModelSelectors.filteredAiProviderModelList(state);
+      expect(result).toHaveLength(1);
+      expect(result[0].id).toBe('model4');
+    });
+
+    it('should handle empty keyword', () => {
+      const result = aiModelSelectors.filteredAiProviderModelList(mockState);
+      expect(result).toHaveLength(mockState.aiProviderModelList.length);
+    });
+  });
+
+  describe('totalAiProviderModelList', () => {
+    it('should return total number of models', () => {
+      const result = aiModelSelectors.totalAiProviderModelList(mockState);
+      expect(result).toBe(4);
+    });
+  });
+
+  describe('isEmptyAiProviderModelList', () => {
+    it('should return true when list is empty', () => {
+      const state = { ...mockState, aiProviderModelList: [] };
+      const result = aiModelSelectors.isEmptyAiProviderModelList(state);
+      expect(result).toBe(true);
+    });
+
+    it('should return false when list is not empty', () => {
+      const result = aiModelSelectors.isEmptyAiProviderModelList(mockState);
+      expect(result).toBe(false);
+    });
+  });
+
+  describe('hasRemoteModels', () => {
+    it('should return true when remote models exist', () => {
+      const result = aiModelSelectors.hasRemoteModels(mockState);
+      expect(result).toBe(true);
+    });
+
+    it('should return false when no remote models exist', () => {
+      const state = {
+        ...mockState,
+        aiProviderModelList: mockState.aiProviderModelList.filter(
+          (m) => !('source' in m) || m.source !== AiModelSourceEnum.Remote,
+        ),
+      };
+      const result = aiModelSelectors.hasRemoteModels(state);
+      expect(result).toBe(false);
+    });
+  });
+
+  describe('isModelEnabled', () => {
+    it('should return true for enabled model', () => {
+      const result = aiModelSelectors.isModelEnabled('model1')(mockState);
+      expect(result).toBe(true);
+    });
+
+    it('should return false for disabled model', () => {
+      const result = aiModelSelectors.isModelEnabled('model2')(mockState);
+      expect(result).toBe(false);
+    });
+  });
+
+  describe('isModelLoading', () => {
+    it('should return true for loading model', () => {
+      const result = aiModelSelectors.isModelLoading('model2')(mockState);
+      expect(result).toBe(true);
+    });
+
+    it('should return false for non-loading model', () => {
+      const result = aiModelSelectors.isModelLoading('model1')(mockState);
+      expect(result).toBe(false);
+    });
+  });
+
+  describe('getAiModelById', () => {
+    it('should return model by id', () => {
+      const result = aiModelSelectors.getAiModelById('model1')(mockState);
+      expect(result).toBeDefined();
+      expect(result?.id).toBe('model1');
+    });
+
+    it('should return undefined for non-existent model', () => {
+      const result = aiModelSelectors.getAiModelById('nonexistent')(mockState);
+      expect(result).toBeUndefined();
+    });
+  });
+
+  describe('model capability checks', () => {
+    it('should check tool use support', () => {
+      expect(aiModelSelectors.isModelSupportToolUse('model1', 'provider1')(mockState)).toBe(true);
+      expect(aiModelSelectors.isModelSupportToolUse('model4', 'provider2')(mockState)).toBe(false);
+    });
+
+    it('should check vision support', () => {
+      expect(aiModelSelectors.isModelSupportVision('model1', 'provider1')(mockState)).toBe(true);
+      expect(aiModelSelectors.isModelSupportVision('model4', 'provider2')(mockState)).toBe(false);
+    });
+
+    it('should check reasoning support', () => {
+      expect(aiModelSelectors.isModelSupportReasoning('model1', 'provider1')(mockState)).toBe(true);
+      expect(aiModelSelectors.isModelSupportReasoning('model4', 'provider2')(mockState)).toBe(
+        false,
+      );
+    });
+  });
+
+  describe('context window checks', () => {
+    it('should check if model has context window tokens', () => {
+      expect(aiModelSelectors.isModelHasContextWindowToken('model1', 'provider1')(mockState)).toBe(
+        true,
+      );
+      expect(aiModelSelectors.isModelHasContextWindowToken('model4', 'provider2')(mockState)).toBe(
+        false,
+      );
+    });
+
+    it('should get model context window tokens', () => {
+      expect(aiModelSelectors.modelContextWindowTokens('model1', 'provider1')(mockState)).toBe(
+        4000,
+      );
+      expect(
+        aiModelSelectors.modelContextWindowTokens('model4', 'provider2')(mockState),
+      ).toBeUndefined();
+    });
+  });
+});