From 1caa61f4c0e8d35bfff2dd670925f8c1ceb8267a Mon Sep 17 00:00:00 2001
From: Dogtiti <499960698@qq.com>
Date: Sat, 6 Jul 2024 22:59:20 +0800
Subject: [PATCH] feat: swap name and displayName for bytedance in custom
 models

---
 app/client/api.ts    |  2 ++
 app/config/server.ts |  6 +++---
 app/constant.ts      |  9 ++++++++-
 app/utils/model.ts   | 12 ++++++++++--
 4 files changed, 23 insertions(+), 6 deletions(-)

diff --git a/app/client/api.ts b/app/client/api.ts
index d2eeca46a..f2e83c391 100644
--- a/app/client/api.ts
+++ b/app/client/api.ts
@@ -225,6 +225,8 @@ export function getClientApi(provider: ServiceProvider): ClientApi {
       return new ClientApi(ModelProvider.GeminiPro);
     case ServiceProvider.Anthropic:
       return new ClientApi(ModelProvider.Claude);
+    case ServiceProvider.ByteDance:
+      return new ClientApi(ModelProvider.Doubao);
     default:
       return new ClientApi(ModelProvider.GPT);
   }
diff --git a/app/config/server.ts b/app/config/server.ts
index d50dbf1a1..0f57d2d6d 100644
--- a/app/config/server.ts
+++ b/app/config/server.ts
@@ -32,13 +32,13 @@ declare global {
       GOOGLE_API_KEY?: string;
       GOOGLE_URL?: string;
 
+      // google tag manager
+      GTM_ID?: string;
+
       // bytedance only
       BYTEDANCE_URL?: string;
       BYTEDANCE_API_KEY?: string;
 
-      // google tag manager
-      GTM_ID?: string;
-
       // custom template for preprocessing user input
       DEFAULT_INPUT_TEMPLATE?: string;
     }
diff --git a/app/constant.ts b/app/constant.ts
index 1ed292d21..5b52073bb 100644
--- a/app/constant.ts
+++ b/app/constant.ts
@@ -183,7 +183,14 @@ const anthropicModels = [
   "claude-3-5-sonnet-20240620",
 ];
 
-const bytedanceModels = ["ep-20240520082937-424bw=Doubao-lite-4k"];
+const bytedanceModels = [
+  "Doubao-lite-4k",
+  "Doubao-lite-32k",
+  "Doubao-lite-128k",
+  "Doubao-pro-4k",
+  "Doubao-pro-32k",
+  "Doubao-pro-128k",
+];
 
 export const DEFAULT_MODELS = [
   ...openaiModels.map((name) => ({
diff --git a/app/utils/model.ts b/app/utils/model.ts
index 249987726..62ecc55b3 100644
--- a/app/utils/model.ts
+++ b/app/utils/model.ts
@@ -39,7 +39,7 @@ export function collectModelTable(
       const available = !m.startsWith("-");
       const nameConfig =
         m.startsWith("+") || m.startsWith("-") ? m.slice(1) : m;
-      const [name, displayName] = nameConfig.split("=");
+      let [name, displayName] = nameConfig.split("=");
 
       // enable or disable all models
       if (name === "all") {
@@ -50,9 +50,17 @@ export function collectModelTable(
         // 1. find model by name(), and set available value
         let count = 0;
         for (const fullName in modelTable) {
-          if (fullName.split("@").shift() == name) {
+          const [modelName, providerName] = fullName.split("@");
+          if (modelName === name) {
             count += 1;
             modelTable[fullName]["available"] = available;
+            // swap name and displayName for bytedance
+            if (providerName === "bytedance") {
+              const tempName = name;
+              name = displayName;
+              displayName = tempName;
+              modelTable[fullName]["name"] = name;
+            }
             if (displayName) {
               modelTable[fullName]["displayName"] = displayName;
             }