♻️ refactor: refactor the key

bentwnghk · Apr 10, 2024 · d5c82f6 · d5c82f6
1 parent e8ed847
commit d5c82f6
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 8 deletions.
diff --git a/src/const/settings/index.ts b/src/const/settings/index.ts
@@ -103,18 +103,18 @@ export const DEFAULT_LLM_CONFIG: GlobalLLMConfig = {
   },
   ollama: {
     enabled: false,
+    enabledModels: filterEnabledModels(OllamaProvider),
     endpoint: '',
-    models: filterEnabledModels(OllamaProvider),
   },
   openAI: {
     OPENAI_API_KEY: '',
     enabled: true,
-    models: filterEnabledModels(OpenAIProvider),
+    enabledModels: filterEnabledModels(OpenAIProvider),
   },
   openrouter: {
     apiKey: '',
     enabled: false,
-    models: filterEnabledModels(OpenRouterProvider),
+    enabledModels: filterEnabledModels(OpenRouterProvider),
   },
   perplexity: {
     apiKey: '',

diff --git a/src/store/global/slices/settings/selectors/modelConfig.test.ts b/src/store/global/slices/settings/selectors/modelConfig.test.ts
@@ -16,15 +16,22 @@ describe('modelConfigSelectors', () => {
         settings: {
           languageModel: {
             ollama: {
-              models: ['llava'],
+              enabledModels: ['llava'],
             },
           },
         },
       } as GlobalSettingsState) as unknown as GlobalStore;
 
       const ollamaList = modelConfigSelectors.modelSelectList(s).find((r) => r.id === 'ollama');
 
-      expect(ollamaList?.chatModels).toEqual([]);
+      expect(ollamaList?.chatModels.find((c) => c.id === 'llava')).toEqual({
+        displayName: 'LLaVA 7B',
+        functionCall: false,
+        hidden: false,
+        id: 'llava',
+        tokens: 4000,
+        vision: true,
+      });
     });
   });
 });
diff --git a/src/types/settings/modelProvider.ts b/src/types/settings/modelProvider.ts
@@ -21,8 +21,8 @@ export interface OpenAIConfig {
    */
   customModelName?: string;
   enabled: boolean;
+  enabledModels?: string[];
   endpoint?: string;
-  models?: string[];
   useAzure?: boolean;
 }
 
@@ -45,15 +45,15 @@ export interface AWSBedrockConfig {
 export interface OllamaConfig {
   customModelName?: string;
   enabled?: boolean;
+  enabledModels: string[];
   endpoint?: string;
-  models: string[];
 }
 
 export interface OpenRouterConfig {
   apiKey?: string;
   customModelName?: string;
   enabled?: boolean;
-  models: string[];
+  enabledModels: string[];
 }
 
 export interface TogetherAIConfig {