From cdd74de666bc2e9d737dc75c6e1fcd92338fd70e Mon Sep 17 00:00:00 2001
From: rxliuli <rxliuli@gmail.com>
Date: Mon, 26 Aug 2024 02:06:32 +0800
Subject: [PATCH] fix: close 5315

---
 app/utils/model.ts | 29 +++++++++++++++++++++++++++--
 1 file changed, 27 insertions(+), 2 deletions(-)

diff --git a/app/utils/model.ts b/app/utils/model.ts
index 0b62b53be..fa9888042 100644
--- 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,
           );