From 9203870df529146b5cc35f6d83925f1edafa59ac Mon Sep 17 00:00:00 2001
From: lloydzhou <lloydzhou@qq.com>
Date: Fri, 12 Jul 2024 20:19:36 +0800
Subject: [PATCH] hotfix: using custom models, create custom provider

---
 app/utils/model.ts | 17 ++++++++++-------
 1 file changed, 10 insertions(+), 7 deletions(-)

diff --git a/app/utils/model.ts b/app/utils/model.ts
index 8f6a1a6c7..2681882ef 100644
--- a/app/utils/model.ts
+++ b/app/utils/model.ts
@@ -1,9 +1,9 @@
 import { DEFAULT_MODELS } from "../constant";
 import { LLMModel } from "../client/api";
 
-const customProvider = (modelName: string) => ({
-  id: modelName,
-  providerName: "Custom",
+const customProvider = (providerName: string) => ({
+  id: providerName.toLowerCase(),
+  providerName: providerName,
   providerType: "custom",
 });
 
@@ -71,10 +71,13 @@ export function collectModelTable(
         }
         // 2. if model not exists, create new model with available value
         if (count === 0) {
-          const provider = customProvider(name);
-          modelTable[`${name}@${provider?.id}`] = {
-            name,
-            displayName: displayName || name,
+          const [customModelName, customProviderName] = name.split("@");
+          const provider = customProvider(
+            customProviderName || customModelName,
+          );
+          modelTable[`${customModelName}@${provider?.id}`] = {
+            name: customModelName,
+            displayName: displayName || customModelName,
             available,
             provider, // Use optional chaining
           };