fix: close 5315

ChatGPTNextWeb · Aug 25, 2024 · cdd74de · cdd74de
1 parent 718782f
commit cdd74de
Showing 1 changed file with 27 additions and 2 deletions.
diff --git a/app/utils/model.ts b/app/utils/model.ts
@@ -62,6 +62,26 @@ export function collectModelTable(
     };
   });
 
+  function parseModelName(name: string): {
+    customModelName: string;
+    customProviderName: string;
+  } {
+    let customModelName, customProviderName;
+    if (name.startsWith("'") || name.startsWith('"')) {
+      const match = name.match(/^(['"])(.*?)\1(@.*)?$/);
+      if (match) {
+        customModelName = match[2];
+        customProviderName = match[3]?.slice(1) || customModelName;
+      }
+    } else {
+      [customModelName, customProviderName] = name.split("@");
+    }
+    return { customModelName, customProviderName } as {
+      customModelName: string;
+      customProviderName: string;
+    };
+  }
+
   // server custom models
   customModels
     .split(",")
@@ -79,7 +99,12 @@ export function collectModelTable(
         );
       } else {
         // 1. find model by name, and set available value
-        const [customModelName, customProviderName] = name.split("@");
+        // modelName@azure => 'modelName', 'azure'
+        // 'modelName@azure' => 'modelName@azure', ''
+        // modelName@azure=deploymentName => 'modelName', 'azure=deploymentName'
+        // 'modelName@azure'=deploymentName => 'modelName@azure', '=deploymentName
+        // 'modelName@azure'@azure=deploymentName => 'modelName@azure', 'azure=deploymentName'
+        let { customModelName, customProviderName } = parseModelName(name);
         let count = 0;
         for (const fullName in modelTable) {
           const [modelName, providerName] = fullName.split("@");
@@ -102,7 +127,7 @@ export function collectModelTable(
         }
         // 2. if model not exists, create new model with available value
         if (count === 0) {
-          let [customModelName, customProviderName] = name.split("@");
+          let { customModelName, customProviderName } = parseModelName(name);
           const provider = customProvider(
             customProviderName || customModelName,
           );