🐛 fix: support openrouter custom models env (lobehub#1647)

* ✨ feat: add openrouter custom models env * ✅ add OPENROUTER_CUSTOM_MODELS env var test case
Maxc100 · Mar 25, 2024 · 78baa16 · 78baa16
1 parent 601d6cd
commit 78baa16
Show file tree

Hide file tree

Showing 9 changed files with 197 additions and 4 deletions.
diff --git a/.env.example b/.env.example
@@ -98,6 +98,7 @@ OPENAI_API_KEY=sk-xxxxxxxxx
 ########################################
 
 #OPENROUTER_API_KEY=xxxxxxxxxxxxxxxxxxxxxxxxxxxxx
+#OPENROUTER_CUSTOM_MODELS=model1,model2,model3
 
 ########################################
 ######### 01.AI Service ##########

diff --git a/Dockerfile b/Dockerfile
@@ -94,6 +94,7 @@ ENV MISTRAL_API_KEY ""
 
 # OpenRouter
 ENV OPENROUTER_API_KEY ""
+ENV OPENROUTER_CUSTOM_MODELS ""
 
 # 01.AI
 ENV ZEROONE_API_KEY ""

diff --git a/docs/self-hosting/environment-variables/model-provider.mdx b/docs/self-hosting/environment-variables/model-provider.mdx
@@ -188,5 +188,12 @@ When using the `turn` mode, the API Keys will be retrieved in a round-robin mann
 - Default: -
 - Example: `sk-or-v1-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx=`
 
+### `OPENROUTER_CUSTOM_MODELS`
+
+- Type: Optional
+- Description: Used to control the model list, use `+` to add a model, use `-` to hide a model, use `model_name=display_name` to customize the display name of a model, separated by commas.
+- Default: `-`
+- Example: `-all,+01-ai/yi-34b-chat,+huggingfaceh4/zephyr-7b-beta`
+
 
 [azure-api-verion-url]: https://docs.microsoft.com/zh-cn/azure/developer/javascript/api-reference/es-modules/azure-sdk/ai-translation/translationconfiguration?view=azure-node-latest#api-version
diff --git a/docs/self-hosting/environment-variables/model-provider.zh-CN.mdx b/docs/self-hosting/environment-variables/model-provider.zh-CN.mdx
@@ -186,6 +186,13 @@ LobeChat 在部署时提供了丰富的模型服务商相关的环境变量，
 - 默认值：-
 - 示例：`sk-or-v1-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx=`
 
+### `OPENROUTER_CUSTOM_MODELS`
+
+- 类型：可选
+- 描述：用来控制模型列表，使用 `+` 增加一个模型，使用 `-` 来隐藏一个模型，使用 `模型名=展示名` 来自定义模型的展示名，用英文逗号隔开。
+- 默认值：`-`
+- 示例：`-all,+01-ai/yi-34b-chat,+huggingfaceh4/zephyr-7b-beta`
+
 ## 01 AI
 
 ### `ZEROONE_API_KEY`

diff --git a/src/app/api/config/route.ts b/src/app/api/config/route.ts
@@ -26,6 +26,7 @@ export const GET = async () => {
     ENABLED_ZEROONE,
     DEFAULT_AGENT_CONFIG,
     OLLAMA_CUSTOM_MODELS,
+    OPENROUTER_CUSTOM_MODELS,
   } = getServerConfig();
 
   const config: GlobalServerConfig = {
@@ -43,7 +44,7 @@ export const GET = async () => {
       mistral: { enabled: ENABLED_MISTRAL },
       moonshot: { enabled: ENABLED_MOONSHOT },
       ollama: { customModelName: OLLAMA_CUSTOM_MODELS, enabled: ENABLE_OLLAMA },
-      openrouter: { enabled: ENABLED_OPENROUTER },
+      openrouter: { customModelName: OPENROUTER_CUSTOM_MODELS, enabled: ENABLED_OPENROUTER },
       perplexity: { enabled: ENABLED_PERPLEXITY },
       zeroone: { enabled: ENABLED_ZEROONE },
       zhipu: { enabled: ENABLED_ZHIPU },

diff --git a/src/config/server/provider.ts b/src/config/server/provider.ts
@@ -45,6 +45,7 @@ declare global {
 
       // OpenRouter Provider
       OPENROUTER_API_KEY?: string;
+      OPENROUTER_CUSTOM_MODELS?: string;
 
       // ZeroOne Provider
       ZEROONE_API_KEY?: string;
@@ -118,6 +119,7 @@ export const getProviderConfig = () => {
 
     ENABLED_OPENROUTER: !!OPENROUTER_API_KEY,
     OPENROUTER_API_KEY,
+    OPENROUTER_CUSTOM_MODELS: process.env.OPENROUTER_CUSTOM_MODELS,
 
     ENABLED_MOONSHOT: !!MOONSHOT_API_KEY,
     MOONSHOT_API_KEY,

diff --git a/src/store/global/slices/settings/selectors/__snapshots__/modelProvider.test.ts.snap b/src/store/global/slices/settings/selectors/__snapshots__/modelProvider.test.ts.snap
@@ -86,3 +86,145 @@ exports[`modelProviderSelectors > CUSTOM_MODELS > should work correct with gpt-4
   },
 ]
 `;
+
+exports[`modelProviderSelectors > OPENROUTER_CUSTOM_MODELS > custom deletion, addition, and renaming of models 1`] = `
+[
+  {
+    "chatModels": [
+      {
+        "description": "GPT 3.5 Turbo，适用于各种文本生成和理解任务",
+        "displayName": "GPT-3.5 Turbo",
+        "functionCall": true,
+        "id": "gpt-3.5-turbo",
+        "tokens": 16385,
+      },
+      {
+        "displayName": "GPT-3.5 Turbo (0125)",
+        "functionCall": true,
+        "hidden": true,
+        "id": "gpt-3.5-turbo-0125",
+        "tokens": 16385,
+      },
+      {
+        "displayName": "GPT-3.5 Turbo (1106)",
+        "functionCall": true,
+        "hidden": true,
+        "id": "gpt-3.5-turbo-1106",
+        "tokens": 16385,
+      },
+      {
+        "displayName": "GPT-3.5 Turbo Instruct",
+        "hidden": true,
+        "id": "gpt-3.5-turbo-instruct",
+        "tokens": 4096,
+      },
+      {
+        "displayName": "GPT-3.5 Turbo 16K",
+        "hidden": true,
+        "id": "gpt-3.5-turbo-16k",
+        "tokens": 16385,
+      },
+      {
+        "displayName": "GPT-3.5 Turbo (0613)",
+        "hidden": true,
+        "id": "gpt-3.5-turbo-0613",
+        "legacy": true,
+        "tokens": 4096,
+      },
+      {
+        "displayName": "GPT-3.5 Turbo 16K (0613)",
+        "hidden": true,
+        "id": "gpt-3.5-turbo-16k-0613",
+        "legacy": true,
+        "tokens": 4096,
+      },
+      {
+        "displayName": "GPT-4 Turbo Preview",
+        "functionCall": true,
+        "id": "gpt-4-turbo-preview",
+        "tokens": 128000,
+      },
+      {
+        "displayName": "GPT-4 Turbo Preview (0125)",
+        "functionCall": true,
+        "hidden": true,
+        "id": "gpt-4-0125-preview",
+        "tokens": 128000,
+      },
+      {
+        "description": "GPT-4 视觉预览版，支持视觉任务",
+        "displayName": "GPT-4 Turbo Vision Preview",
+        "id": "gpt-4-vision-preview",
+        "tokens": 128000,
+        "vision": true,
+      },
+      {
+        "displayName": "GPT-4 Turbo Preview (1106)",
+        "functionCall": true,
+        "hidden": true,
+        "id": "gpt-4-1106-preview",
+        "tokens": 128000,
+      },
+      {
+        "displayName": "GPT-4",
+        "functionCall": true,
+        "hidden": true,
+        "id": "gpt-4",
+        "tokens": 8192,
+      },
+      {
+        "displayName": "GPT-4 (0613)",
+        "functionCall": true,
+        "hidden": true,
+        "id": "gpt-4-0613",
+        "tokens": 8192,
+      },
+      {
+        "displayName": "GPT-4 32K",
+        "functionCall": true,
+        "hidden": true,
+        "id": "gpt-4-32k",
+        "tokens": 32768,
+      },
+      {
+        "displayName": "GPT-4 32K (0613)",
+        "functionCall": true,
+        "hidden": true,
+        "id": "gpt-4-32k-0613",
+        "tokens": 32768,
+      },
+      {
+        "displayName": "GPT-4 ALL",
+        "files": true,
+        "functionCall": true,
+        "hidden": true,
+        "id": "gpt-4-all",
+        "tokens": 32768,
+        "vision": true,
+      },
+    ],
+    "enabled": true,
+    "id": "openai",
+  },
+  {
+    "chatModels": [
+      {
+        "displayName": "google/gemma-7b-it",
+        "functionCall": true,
+        "id": "google/gemma-7b-it",
+        "isCustom": true,
+        "vision": true,
+      },
+      {
+        "displayName": "Mistral-7B-Instruct",
+        "functionCall": true,
+        "id": "mistralai/mistral-7b-instruct",
+        "isCustom": true,
+        "vision": true,
+      },
+    ],
+    "enabled": true,
+    "id": "openrouter",
+  },
+]
+`;
diff --git a/src/store/global/slices/settings/selectors/modelProvider.test.ts b/src/store/global/slices/settings/selectors/modelProvider.test.ts
@@ -131,6 +131,32 @@ describe('modelProviderSelectors', () => {
     });
   });
 
+  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();
+    });
+  });
+
   describe('modelEnabledVision', () => {
     it('should return true if the model has vision ability', () => {
       const hasAbility = modelProviderSelectors.modelEnabledVision('gpt-4-vision-preview')(

diff --git a/src/store/global/slices/settings/selectors/modelProvider.ts b/src/store/global/slices/settings/selectors/modelProvider.ts
@@ -149,9 +149,15 @@ const modelSelectList = (s: GlobalStore): ModelProviderCard[] => {
 
   const ollamaChatModels = processChatModels(ollamaModelConfig, OllamaProvider.chatModels);
 
-  const openrouterModelConfig = parseModelString(
-    currentSettings(s).languageModel.openrouter.customModelName,
-  )
+  const openrouterModelString = [
+    s.serverConfig.languageModel?.openrouter?.customModelName,
+    currentSettings(s).languageModel.openrouter.customModelName
+  ]
+    .filter(Boolean)
+    .join(',');
+
+  const openrouterModelConfig = parseModelString(openrouterModelString);
+
   const openrouterChatModels = processChatModels(openrouterModelConfig, OpenRouterProvider.chatModels);
 
   return [