From c9f07f0eb5dbe7e9f852fa211af87f1e4af74d65 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=B5=B7=E8=B4=BC=E5=AE=85?= <stu_xyx@163.com>
Date: Thu, 23 May 2024 23:27:10 +0800
Subject: [PATCH 1/2] BUG fix - Bug fix for tokens counting logic of VolcEngine

---
 rag/llm/chat_model.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/rag/llm/chat_model.py b/rag/llm/chat_model.py
index 9bf4e96afb..705bfb6110 100644
--- a/rag/llm/chat_model.py
+++ b/rag/llm/chat_model.py
@@ -359,6 +359,7 @@ def chat_streamly(self, system, history, gen_conf):
         if system:
             history.insert(0, {"role": "system", "content": system})
         ans = ""
+        tk_count = 0
         try:
             req = {
                 "parameters": {
@@ -376,10 +377,9 @@ def chat_streamly(self, system, history, gen_conf):
                 if not resp.choices[0].message.content:
                     continue
                 ans += resp.choices[0].message.content
-                yield ans
                 if resp.choices[0].finish_reason == "stop":
-                    yield resp.usage.total_tokens
-
+                    tk_count = resp.usage.total_tokens
+                yield ans
         except Exception as e:
             yield ans + "\n**ERROR**: " + str(e)
-        yield 0
+        yield tk_count

From 084247c59f1a782ae4139fce740aed84ec1ec319 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E6=B5=B7=E8=B4=BC=E5=AE=85?= <stu_xyx@163.com>
Date: Thu, 23 May 2024 23:28:22 +0800
Subject: [PATCH 2/2] BUG fix - Bug fix for front-end configuration code of
 VolcEngine

---
 web/src/pages/user-setting/setting-model/index.tsx | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/web/src/pages/user-setting/setting-model/index.tsx b/web/src/pages/user-setting/setting-model/index.tsx
index 69a770108e..071d9ec2ad 100644
--- a/web/src/pages/user-setting/setting-model/index.tsx
+++ b/web/src/pages/user-setting/setting-model/index.tsx
@@ -102,7 +102,7 @@ const ModelCard = ({ item, clickApiKey }: IModelCardProps) => {
           <Col span={12} className={styles.factoryOperationWrapper}>
             <Space size={'middle'}>
               <Button onClick={handleApiKeyClick}>
-                {isLocalLlmFactory(item.name) ? t('addTheModel') : 'API-Key'}
+                {isLocalLlmFactory(item.name) || item.name === 'VolcEngine' ? t('addTheModel') : 'API-Key'}
                 <SettingOutlined />
               </Button>
               <Button onClick={handleShowMoreClick}>
@@ -181,11 +181,13 @@ const UserSettingModel = () => {
     (llmFactory: string) => {
       if (isLocalLlmFactory(llmFactory)) {
         showLlmAddingModal(llmFactory);
+      } else if (llmFactory === 'VolcEngine') {
+        showVolcAddingModal('VolcEngine');
       } else {
         showApiKeyModal({ llm_factory: llmFactory });
       }
     },
-    [showApiKeyModal, showLlmAddingModal],
+    [showApiKeyModal, showLlmAddingModal, showVolcAddingModal],
   );
 
   const handleAddModel = (llmFactory: string) => () => {