From c0d71adaa2c6c8d5b54d2ec5e4fe51ad0ea07c01 Mon Sep 17 00:00:00 2001
From: yungongzi <yun_gong_zi@yeah.net>
Date: Fri, 24 May 2024 11:34:39 +0800
Subject: [PATCH] Bug fix for volcengine (#909)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

### What problem does this PR solve?
Bug fixes for the VolcEngine

- Bug fix for front-end configuration code of VolcEngine

- Bug fix for tokens counting logic of VolcEngine


### Type of change

- [x] Bug Fix (non-breaking change which fixes an issue)

---------

Co-authored-by: 海贼宅 <stu_xyx@163.com>
---
 rag/llm/chat_model.py                              | 8 ++++----
 web/src/pages/user-setting/setting-model/index.tsx | 6 ++++--
 2 files changed, 8 insertions(+), 6 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
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) => () => {