From f7ce3b8686356852b66b08989581e38fbf7453a3 Mon Sep 17 00:00:00 2001
From: kangfenmao <kangfenmao@qq.com>
Date: Mon, 8 Jul 2024 13:47:04 +0800
Subject: [PATCH] feat: add 01-yi provider

---
 src/renderer/src/config/models.ts | 50 +++++++++++++++++++++++++++++++
 src/renderer/src/store/index.ts   |  6 ++--
 src/renderer/src/store/llm.ts     |  8 +++++
 src/renderer/src/store/migrate.ts | 28 +++++++++++++++++
 4 files changed, 90 insertions(+), 2 deletions(-)
 create mode 100644 src/renderer/src/store/migrate.ts

diff --git a/src/renderer/src/config/models.ts b/src/renderer/src/config/models.ts
index efa5f28f..12ca9232 100644
--- a/src/renderer/src/config/models.ts
+++ b/src/renderer/src/config/models.ts
@@ -201,6 +201,56 @@ export const SYSTEM_MODELS: Record<string, SystemModel[]> = {
       defaultEnabled: true
     }
   ],
+  yi: [
+    {
+      id: 'yi-large',
+      provider: 'yi',
+      name: 'yi-large',
+      group: 'Yi',
+      temperature: 0.7,
+      defaultEnabled: false
+    },
+    {
+      id: 'yi-large-turbo',
+      provider: 'yi',
+      name: 'yi-large-turbo',
+      group: 'Yi',
+      temperature: 0.7,
+      defaultEnabled: true
+    },
+    {
+      id: 'yi-large-rag',
+      provider: 'yi',
+      name: 'yi-large-rag',
+      group: 'Yi',
+      temperature: 0.7,
+      defaultEnabled: false
+    },
+    {
+      id: 'yi-medium',
+      provider: 'yi',
+      name: 'yi-medium',
+      group: 'Yi',
+      temperature: 0.7,
+      defaultEnabled: true
+    },
+    {
+      id: 'yi-medium-200k',
+      provider: 'yi',
+      name: 'yi-medium-200k',
+      group: 'Yi',
+      temperature: 0.7,
+      defaultEnabled: false
+    },
+    {
+      id: 'yi-spark',
+      provider: 'yi',
+      name: 'yi-spark',
+      group: 'Yi',
+      temperature: 0.7,
+      defaultEnabled: false
+    }
+  ],
   groq: [
     {
       id: 'llama3-8b-8192',
diff --git a/src/renderer/src/store/index.ts b/src/renderer/src/store/index.ts
index bb864f1a..c2d96ddc 100644
--- a/src/renderer/src/store/index.ts
+++ b/src/renderer/src/store/index.ts
@@ -6,6 +6,7 @@ import assistants from './assistants'
 import settings from './settings'
 import llm from './llm'
 import runtime from './runtime'
+import migrate from './migrate'
 
 const rootReducer = combineReducers({
   assistants,
@@ -18,8 +19,9 @@ const persistedReducer = persistReducer(
   {
     key: 'cherry-studio',
     storage,
-    version: 1,
-    blacklist: ['runtime']
+    version: 2,
+    blacklist: ['runtime'],
+    migrate
   },
   rootReducer
 )
diff --git a/src/renderer/src/store/llm.ts b/src/renderer/src/store/llm.ts
index 6f304638..80e10564 100644
--- a/src/renderer/src/store/llm.ts
+++ b/src/renderer/src/store/llm.ts
@@ -37,6 +37,14 @@ const initialState: LlmState = {
       isSystem: true,
       models: SYSTEM_MODELS.deepseek.filter((m) => m.defaultEnabled)
     },
+    {
+      id: 'yi',
+      name: 'Yi',
+      apiKey: '',
+      apiHost: 'https://api.lingyiwanwu.com',
+      isSystem: true,
+      models: SYSTEM_MODELS.yi.filter((m) => m.defaultEnabled)
+    },
     {
       id: 'groq',
       name: 'Groq',
diff --git a/src/renderer/src/store/migrate.ts b/src/renderer/src/store/migrate.ts
new file mode 100644
index 00000000..1eb5f601
--- /dev/null
+++ b/src/renderer/src/store/migrate.ts
@@ -0,0 +1,28 @@
+import { createMigrate } from 'redux-persist'
+import { RootState } from '.'
+import { SYSTEM_MODELS } from '@renderer/config/models'
+
+const migrate = createMigrate({
+  // @ts-ignore store type is unknown
+  '2': (state: RootState) => {
+    return {
+      ...state,
+      llm: {
+        ...state.llm,
+        providers: [
+          ...state.llm.providers,
+          {
+            id: 'yi',
+            name: 'Yi',
+            apiKey: '',
+            apiHost: 'https://api.lingyiwanwu.com',
+            isSystem: true,
+            models: SYSTEM_MODELS.yi.filter((m) => m.defaultEnabled)
+          }
+        ]
+      }
+    }
+  }
+})
+
+export default migrate