@lobehub/chat 0.146.2 → 0.147.0

package/src/store/global/slices/settings/selectors/modelConfig.ts

@@ -0,0 +1,179 @@
+import { uniqBy } from 'lodash-es';
+
+import { ChatModelCard, ModelProviderCard } from '@/types/llm';
+import { GeneralModelProviderConfig, GlobalLLMProviderKey } from '@/types/settings';
+
+import { GlobalStore } from '../../../store';
+import { modelProviderSelectors } from './modelProvider';
+import { currentSettings } from './settings';
+
+const modelProvider = (s: GlobalStore) => currentSettings(s).languageModel;
+
+const providerConfig = (provider: string) => (s: GlobalStore) =>
+  currentSettings(s).languageModel[provider as GlobalLLMProviderKey] as
+    | GeneralModelProviderConfig
+    | undefined;
+
+const providerEnabled = (provider: GlobalLLMProviderKey) => (s: GlobalStore) =>
+  currentSettings(s).languageModel[provider]?.enabled || false;
+
+const providerEnableModels = (provider: string) => (s: GlobalStore) => {
+  if (!providerConfig(provider)(s)?.enabledModels) return;
+
+  return providerConfig(provider)(s)?.enabledModels?.filter(Boolean);
+};
+
+const openAIConfig = (s: GlobalStore) => modelProvider(s).openai;
+
+const openAIAPIKey = (s: GlobalStore) => modelProvider(s).openai.apiKey;
+const openAIProxyUrl = (s: GlobalStore) => modelProvider(s).openai.endpoint;
+
+const zhipuAPIKey = (s: GlobalStore) => modelProvider(s).zhipu.apiKey;
+
+const bedrockConfig = (s: GlobalStore) => modelProvider(s).bedrock;
+
+const googleAPIKey = (s: GlobalStore) => modelProvider(s).google.apiKey;
+
+const enableAzure = (s: GlobalStore) => modelProvider(s).azure.enabled;
+const azureConfig = (s: GlobalStore) => modelProvider(s).azure;
+
+const mistralAPIKey = (s: GlobalStore) => modelProvider(s).mistral.apiKey;
+
+const moonshotAPIKey = (s: GlobalStore) => modelProvider(s).moonshot.apiKey;
+
+const ollamaProxyUrl = (s: GlobalStore) => modelProvider(s).ollama.endpoint;
+
+const perplexityAPIKey = (s: GlobalStore) => modelProvider(s).perplexity.apiKey;
+
+const anthropicAPIKey = (s: GlobalStore) => modelProvider(s).anthropic.apiKey;
+const anthropicProxyUrl = (s: GlobalStore) => modelProvider(s).anthropic.endpoint;
+
+const groqAPIKey = (s: GlobalStore) => modelProvider(s).groq.apiKey;
+
+const openrouterAPIKey = (s: GlobalStore) => modelProvider(s).openrouter.apiKey;
+
+const togetheraiAPIKey = (s: GlobalStore) => modelProvider(s).togetherai.apiKey;
+
+const zerooneAPIKey = (s: GlobalStore) => modelProvider(s).zeroone.apiKey;
+
+const modelSelectList = (s: GlobalStore): ModelProviderCard[] => {
+  return modelProviderSelectors.providerModelList(s).map((list) => ({
+    ...list,
+    chatModels: list.chatModels.map((model) => {
+      const models = providerEnableModels(list.id)(s);
+
+      if (!models) return model;
+
+      return {
+        ...model,
+        enabled: models?.some((m) => m === model.id),
+      };
+    }),
+    enabled: providerEnabled(list.id as any)(s),
+  }));
+};
+
+const enabledModelProviderList = (s: GlobalStore): ModelProviderCard[] =>
+  modelSelectList(s)
+    .filter((s) => s.enabled)
+    .map((provider) => ({
+      ...provider,
+      chatModels: provider.chatModels.filter((model) => model.enabled),
+    }));
+
+const providerCard = (provider: string) => (s: GlobalStore) =>
+  modelProviderSelectors.providerModelList(s).find((s) => s.id === provider);
+
+const providerModelCards =
+  (provider: string) =>
+  (s: GlobalStore): ChatModelCard[] => {
+    const builtinCards = providerCard(provider)(s)?.chatModels || [];
+
+    const userCards = (providerConfig(provider)(s)?.customModelCards || []).map((model) => ({
+      ...model,
+      isCustom: true,
+    }));
+
+    return uniqBy([...userCards, ...builtinCards], 'id');
+  };
+
+const getCustomModelCardById =
+  ({ id, provider }: { id?: string; provider?: string }) =>
+  (s: GlobalStore) => {
+    if (!provider) return;
+
+    const config = providerConfig(provider)(s);
+
+    return config?.customModelCards?.find((m) => m.id === id);
+  };
+
+const currentEditingCustomModelCard = (s: GlobalStore) => {
+  if (!s.editingCustomCardModel) return;
+  const { id, provider } = s.editingCustomCardModel;
+
+  return getCustomModelCardById({ id, provider })(s);
+};
+
+const enabledAutoFetchModels =
+  (provider: GlobalLLMProviderKey) =>
+  (s: GlobalStore): boolean => {
+    return providerConfig(provider)(s)?.autoFetchModelLists || false;
+  };
+
+/* eslint-disable sort-keys-fix/sort-keys-fix,  */
+export const modelConfigSelectors = {
+  providerEnabled,
+  providerEnableModels,
+  providerConfig,
+  providerModelCards,
+  currentEditingCustomModelCard,
+  getCustomModelCardById,
+
+  modelSelectList,
+  enabledModelProviderList,
+  enabledAutoFetchModels,
+
+  // OpenAI
+  openAIConfig,
+  openAIAPIKey,
+  openAIProxyUrl,
+
+  // Azure OpenAI
+  enableAzure,
+  azureConfig,
+  // Zhipu
+  zhipuAPIKey,
+  // Google
+  googleAPIKey,
+
+  // Bedrock
+  bedrockConfig,
+
+  // Moonshot
+  moonshotAPIKey,
+
+  // Ollama
+  ollamaProxyUrl,
+
+  // Perplexity
+  perplexityAPIKey,
+
+  // Anthropic
+  anthropicAPIKey,
+  anthropicProxyUrl,
+
+  // Mistral
+  mistralAPIKey,
+
+  // Groq
+  groqAPIKey,
+
+  // OpenRouter
+  openrouterAPIKey,
+
+  // ZeroOne 零一万物
+  zerooneAPIKey,
+
+  // TogetherAI
+  togetheraiAPIKey,
+};

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

@@ -9,154 +9,36 @@ import { initialSettingsState } from '../initialState';
 import { modelProviderSelectors } from './modelProvider';
 
 describe('modelProviderSelectors', () => {
-  describe('CUSTOM_MODELS', () => {
-    it('custom deletion, addition, and renaming of models', () => {
-      const s = merge(initialSettingsState, {
-        serverConfig: {
-          customModelName:
-            '-all,+llama,+claude-2，-gpt-3.5-turbo,gpt-4-0125-preview=gpt-4-turbo,gpt-4-0125-preview=gpt-4-32k',
-        },
-      }) as unknown as GlobalStore;
-
-      const result = modelProviderSelectors.modelSelectList(s).filter((r) => r.enabled);
-
-      expect(result).toMatchSnapshot();
-    });
-
-    it('should work correct with gpt-4', () => {
-      const s = merge(initialSettingsState, {
-        serverConfig: {
-          customModelName:
-            '-all,+gpt-3.5-turbo-1106,+gpt-3.5-turbo,+gpt-3.5-turbo-16k,+gpt-4,+gpt-4-32k,+gpt-4-1106-preview,+gpt-4-vision-preview',
-        },
-      }) as unknown as GlobalStore;
+  describe('providerCard', () => {
+    it('should return the correct ModelProviderCard when provider ID matches', () => {
+      const s = merge(initialSettingsState, {}) as unknown as GlobalStore;
 
-      const result = modelProviderSelectors.modelSelectList(s).filter((r) => r.enabled);
-
-      expect(result[0].chatModels).toMatchSnapshot();
+      const result = modelProviderSelectors.providerCard('openai')(s);
+      expect(result).not.toBeUndefined();
     });
-    it('duplicate naming model', () => {
-      const s = merge(initialSettingsState, {
-        serverConfig: {},
-        settings: {
-          languageModel: {
-            openAI: {
-              customModelName: 'gpt-4-0125-preview=gpt-4-turbo，gpt-4-0125-preview=gpt-4-32k',
-            },
-          },
-        },
-      }) as unknown as GlobalStore;
-
-      const result = modelProviderSelectors.modelSelectList(s).filter((r) => r.enabled);
-
-      expect(result[0].chatModels.find((s) => s.id === 'gpt-4-0125-preview')?.displayName).toEqual(
-        'gpt-4-32k',
-      );
-    });
-
-    it('should delete model', () => {
-      const s = merge(initialSettingsState, {
-        serverConfig: { customModelName: '-gpt-4' },
-      }) as unknown as GlobalStore;
-
-      const result = modelProviderSelectors.modelSelectList(s).filter((r) => r.enabled);
 
-      expect(result.find((r) => r.id === 'gpt-4')).toBeUndefined();
+    it('should return undefined when provider ID does not exist', () => {
+      const s = merge(initialSettingsState, {}) as unknown as GlobalStore;
+      const result = modelProviderSelectors.providerCard('nonExistingProvider')(s);
+      expect(result).toBeUndefined();
     });
+  });
 
-    it('show the hidden model', () => {
-      const s = merge(initialSettingsState, {
-        serverConfig: {},
-        settings: {
-          languageModel: {
-            openAI: {
-              customModelName: '+gpt-4-1106-preview',
-            },
-          },
-        },
-      }) as unknown as GlobalStore;
-
-      const result = modelProviderSelectors.modelSelectList(s).filter((r) => r.enabled);
-
-      expect(result[0].chatModels.find((o) => o.id === 'gpt-4-1106-preview')).toEqual({
-        displayName: 'GPT-4 Turbo Preview (1106)',
-        functionCall: true,
-        id: 'gpt-4-1106-preview',
-        tokens: 128000,
-      });
-    });
+  describe('defaultEnabledProviderModels', () => {
+    it('should return enabled models for a given provider', () => {
+      const s = merge(initialSettingsState, {}) as unknown as GlobalStore;
 
-    it('only add the model', () => {
-      const s = merge(initialSettingsState, {
-        serverConfig: {},
-        settings: {
-          languageModel: {
-            openAI: {
-              customModelName: 'model1,model2,model3，model4',
-            },
-          },
-        },
-      }) as unknown as GlobalStore;
-
-      const result = modelProviderSelectors.modelSelectList(s).filter((r) => r.enabled);
-
-      expect(result[0].chatModels).toContainEqual({
-        displayName: 'model1',
-        functionCall: true,
-        id: 'model1',
-        isCustom: true,
-        vision: true,
-      });
-      expect(result[0].chatModels).toContainEqual({
-        displayName: 'model2',
-        functionCall: true,
-        id: 'model2',
-        isCustom: true,
-        vision: true,
-      });
-      expect(result[0].chatModels).toContainEqual({
-        displayName: 'model3',
-        functionCall: true,
-        id: 'model3',
-        isCustom: true,
-        vision: true,
-      });
-      expect(result[0].chatModels).toContainEqual({
-        displayName: 'model4',
-        functionCall: true,
-        id: 'model4',
-        isCustom: true,
-        vision: true,
-      });
+      const result = modelProviderSelectors.defaultEnabledProviderModels('openai')(s);
+      expect(result).toEqual(['gpt-3.5-turbo', 'gpt-4-turbo-preview', 'gpt-4-vision-preview']);
     });
-  });
 
-  describe('OPENROUTER_CUSTOM_MODELS', () => {
-    it('custom deletion, addition, and renaming of models', () => {
-      const s = merge(initialSettingsState, {
-        settings: {
-          languageModel: {
-            openrouter: {
-              enabled: true,
-            },
-          },
-        },
-        serverConfig: {
-          languageModel: {
-            openrouter: {
-              apiKey: 'test-openrouter-api-key',
-              customModelName:
-                '-all,+google/gemma-7b-it,+mistralai/mistral-7b-instruct=Mistral-7B-Instruct',
-            },
-          },
-        },
-      }) as unknown as GlobalStore;
-
-      const result = modelProviderSelectors.modelSelectList(s).filter((r) => r.enabled);
-      expect(result).toMatchSnapshot();
+    it('should return undefined for a non-existing provider', () => {
+      const s = merge(initialSettingsState, {}) as unknown as GlobalStore;
+
+      const result = modelProviderSelectors.defaultEnabledProviderModels('nonExistingProvider')(s);
+      expect(result).toBeUndefined();
     });
   });
-
   describe('modelEnabledVision', () => {
     it('should return true if the model has vision ability', () => {
       const hasAbility = modelProviderSelectors.modelEnabledVision('gpt-4-vision-preview')(
@@ -213,4 +95,31 @@ describe('modelProviderSelectors', () => {
       expect(show).toBe(false);
     });
   });
+
+  describe('modelMaxToken', () => {
+    it('should return the correct token count for a model with specified tokens', () => {
+      const model1Tokens = modelProviderSelectors.modelMaxToken('gpt-3.5-turbo')(
+        useGlobalStore.getState(),
+      );
+
+      expect(model1Tokens).toEqual(16385);
+    });
+
+    it('should return 0 for a model without a specified token count', () => {
+      // 测试未指定tokens属性的模型的tokens值，期望为0
+      const tokens = modelProviderSelectors.modelMaxToken('chat-bison-001')(
+        useGlobalStore.getState(),
+      );
+      expect(tokens).toEqual(0);
+    });
+
+    it('should return 0 for a non-existing model', () => {
+      // 测试一个不存在的模型的tokens值，期望为0
+      const tokens = modelProviderSelectors.modelMaxToken('nonExistingModel')(
+        useGlobalStore.getState(),
+      );
+
+      expect(tokens).toEqual(0);
+    });
+  });
 });