Bug fix for volcengine (#909)

### 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: 海贼宅 <[email protected]>
infiniflow · May 24, 2024 · c0d71ad · c0d71ad
1 parent 735bdf0
commit c0d71ad
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 6 deletions.
diff --git 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
@@ -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) => () => {