From d5c82f68cf5a235b4f1d1d6a0227f02192491525 Mon Sep 17 00:00:00 2001
From: Arvin Xu <arvinx@foxmail.com>
Date: Mon, 8 Apr 2024 09:13:40 +0000
Subject: [PATCH] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20refactor:=20refactor=20the?=
 =?UTF-8?q?=20key?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/const/settings/index.ts                           |  6 +++---
 .../slices/settings/selectors/modelConfig.test.ts     | 11 +++++++++--
 src/types/settings/modelProvider.ts                   |  6 +++---
 3 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/src/const/settings/index.ts b/src/const/settings/index.ts
index ff8c7a40b543..acf3b274c9e7 100644
--- 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
index 5112db9243d1..f487ea3041ca 100644
--- a/src/store/global/slices/settings/selectors/modelConfig.test.ts
+++ b/src/store/global/slices/settings/selectors/modelConfig.test.ts
@@ -16,7 +16,7 @@ describe('modelConfigSelectors', () => {
         settings: {
           languageModel: {
             ollama: {
-              models: ['llava'],
+              enabledModels: ['llava'],
             },
           },
         },
@@ -24,7 +24,14 @@ describe('modelConfigSelectors', () => {
 
       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
index 99d75271f7a5..ed174d8332b2 100644
--- 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 {