Add support for VolcEngine - the current version supports SDK2 (#885)

- The main idea is to assemble **ak**, **sk**, and **ep_id** into a
dictionary and store it in the database **api_key** field
- I don’t know much about the front-end, so I learned from Ollama, which
may be redundant.

### Configuration method

- model name

- Format requirements: {"VolcEngine model name":"endpoint_id"}
- For example: {"Skylark-pro-32K":"ep-xxxxxxxxx"}

- Volcano ACCESS_KEY
- Format requirements: VOLC_ACCESSKEY of the volcano engine
corresponding to the model

- Volcano SECRET_KEY
- Format requirements: VOLC_SECRETKEY of the volcano engine
corresponding to the model

### What problem does this PR solve?

_Briefly describe what this PR aims to solve. Include background context
that will help reviewers understand the purpose of the PR._

### Type of change

- [x] New Feature (non-breaking change which adds functionality)

api/apps/llm_app.py

@@ -96,16 +96,29 @@ def set_api_key():
 @validate_request("llm_factory", "llm_name", "model_type")
 def add_llm():
     req = request.json
+    factory = req["llm_factory"]
+    # For VolcEngine, due to its special authentication method
+    # Assemble volc_ak, volc_sk, endpoint_id into api_key
+    if factory == "VolcEngine":
+        temp = list(eval(req["llm_name"]).items())[0]
+        llm_name = temp[0]
+        endpoint_id = temp[1]
+        api_key = '{' + f'"volc_ak": "{req.get("volc_ak", "")}", ' \
+                        f'"volc_sk": "{req.get("volc_sk", "")}", ' \
+                        f'"ep_id": "{endpoint_id}", ' + '}'
+    else:
+        llm_name = req["llm_name"]
+        api_key = "xxxxxxxxxxxxxxx"
+
     llm = {
         "tenant_id": current_user.id,
-        "llm_factory": req["llm_factory"],
+        "llm_factory": factory,
         "model_type": req["model_type"],
-        "llm_name": req["llm_name"],
+        "llm_name": llm_name,
         "api_base": req.get("api_base", ""),
-        "api_key": "xxxxxxxxxxxxxxx"
+        "api_key": api_key
     }
 
-    factory = req["llm_factory"]
     msg = ""
     if llm["model_type"] == LLMType.EMBEDDING.value:
         mdl = EmbeddingModel[factory](
@@ -118,7 +131,10 @@ def add_llm():
             msg += f"\nFail to access embedding model({llm['llm_name']})." + str(e)
     elif llm["model_type"] == LLMType.CHAT.value:
         mdl = ChatModel[factory](
-            key=None, model_name=llm["llm_name"], base_url=llm["api_base"])
+            key=llm['api_key'] if factory == "VolcEngine" else None,
+            model_name=llm["llm_name"],
+            base_url=llm["api_base"]
+        )
         try:
             m, tc = mdl.chat(None, [{"role": "user", "content": "Hello! How are you doing!"}], {
                              "temperature": 0.9})
@@ -134,7 +150,6 @@ def add_llm():
     if msg:
         return get_data_error_result(retmsg=msg)
 
-
     if not TenantLLMService.filter_update(
             [TenantLLM.tenant_id == current_user.id, TenantLLM.llm_factory == factory, TenantLLM.llm_name == llm["llm_name"]], llm):
         TenantLLMService.save(**llm)

api/db/init_data.py

@@ -132,7 +132,12 @@ factory_infos = [{
     "logo": "",
     "tags": "LLM",
     "status": "1",
-},
+},{
+    "name": "VolcEngine",
+    "logo": "",
+    "tags": "LLM, TEXT EMBEDDING",
+    "status": "1",
+}
     # {
     #     "name": "文心一言",
     #     "logo": "",
@@ -372,6 +377,21 @@ def init_llm_factory():
             "max_tokens": 16385,
             "model_type": LLMType.CHAT.value
         },
+        # ------------------------ VolcEngine -----------------------
+        {
+            "fid": factory_infos[9]["name"],
+            "llm_name": "Skylark2-pro-32k",
+            "tags": "LLM,CHAT,32k",
+            "max_tokens": 32768,
+            "model_type": LLMType.CHAT.value
+        },
+        {
+            "fid": factory_infos[9]["name"],
+            "llm_name": "Skylark2-pro-4k",
+            "tags": "LLM,CHAT,4k",
+            "max_tokens": 4096,
+            "model_type": LLMType.CHAT.value
+        },
     ]
     for info in factory_infos:
         try:

rag/llm/chat_model.py

@@ -19,6 +19,7 @@ from abc import ABC
 from openai import OpenAI
 import openai
 from ollama import Client
+from volcengine.maas.v2 import MaasService
 from rag.nlp import is_english
 from rag.utils import num_tokens_from_string
 
@@ -315,3 +316,71 @@ class LocalLLM(Base):
             yield answer + "\n**ERROR**: " + str(e)
 
         yield token_count
+
+
+class VolcEngineChat(Base):
+    def __init__(self, key, model_name, base_url):
+        """
+        Since do not want to modify the original database fields, and the VolcEngine authentication method is quite special,
+        Assemble ak, sk, ep_id into api_key, store it as a dictionary type, and parse it for use
+        model_name is for display only
+        """
+        self.client = MaasService('maas-api.ml-platform-cn-beijing.volces.com', 'cn-beijing')
+        self.volc_ak = eval(key).get('volc_ak', '')
+        self.volc_sk = eval(key).get('volc_sk', '')
+        self.client.set_ak(self.volc_ak)
+        self.client.set_sk(self.volc_sk)
+        self.model_name = eval(key).get('ep_id', '')
+
+    def chat(self, system, history, gen_conf):
+        if system:
+            history.insert(0, {"role": "system", "content": system})
+        try:
+            req = {
+                "parameters": {
+                    "min_new_tokens": gen_conf.get("min_new_tokens", 1),
+                    "top_k": gen_conf.get("top_k", 0),
+                    "max_prompt_tokens": gen_conf.get("max_prompt_tokens", 30000),
+                    "temperature": gen_conf.get("temperature", 0.1),
+                    "max_new_tokens": gen_conf.get("max_tokens", 1000),
+                    "top_p": gen_conf.get("top_p", 0.3),
+                },
+                "messages": history
+            }
+            response = self.client.chat(self.model_name, req)
+            ans = response.choices[0].message.content.strip()
+            if response.choices[0].finish_reason == "length":
+                ans += "...\nFor the content length reason, it stopped, continue?" if is_english(
+                    [ans]) else "······\n由于长度的原因，回答被截断了，要继续吗？"
+            return ans, response.usage.total_tokens
+        except Exception as e:
+            return "**ERROR**: " + str(e), 0
+
+    def chat_streamly(self, system, history, gen_conf):
+        if system:
+            history.insert(0, {"role": "system", "content": system})
+        ans = ""
+        try:
+            req = {
+                "parameters": {
+                    "min_new_tokens": gen_conf.get("min_new_tokens", 1),
+                    "top_k": gen_conf.get("top_k", 0),
+                    "max_prompt_tokens": gen_conf.get("max_prompt_tokens", 30000),
+                    "temperature": gen_conf.get("temperature", 0.1),
+                    "max_new_tokens": gen_conf.get("max_tokens", 1000),
+                    "top_p": gen_conf.get("top_p", 0.3),
+                },
+                "messages": history
+            }
+            stream = self.client.stream_chat(self.model_name, req)
+            for resp in stream:
+                if not resp.choices[0].message.content:
+                    continue
+                ans += resp.choices[0].message.content
+                yield ans
+                if resp.choices[0].finish_reason == "stop":
+                    return resp.usage.total_tokens
+
+        except Exception as e:
+            yield ans + "\n**ERROR**: " + str(e)
+        yield 0

web/src/locales/en.ts

@@ -477,6 +477,11 @@ The above is the content you need to summarize.`,
       baseUrlNameMessage: 'Please input your base url!',
       vision: 'Does it support Vision?',
       ollamaLink: 'How to integrate {{name}}',
+      volcModelNameMessage: 'Please input your model name! Format: {"ModelName":"EndpointID"}',
+      addVolcEngineAK: 'VOLC ACCESS_KEY',
+      volcAKMessage: 'Please input your VOLC_ACCESS_KEY',
+      addVolcEngineSK: 'VOLC SECRET_KEY',
+      volcSKMessage: 'Please input your SECRET_KEY',
     },
     message: {
       registered: 'Registered!',

web/src/locales/zh-traditional.ts

@@ -440,7 +440,12 @@ export default {
       modelNameMessage: '請輸入模型名稱！',
       modelTypeMessage: '請輸入模型類型！',
       baseUrlNameMessage: '請輸入基礎 Url！',
-      ollamaLink: '如何集成Ollama',
+      ollamaLink: '如何集成 {{name}}',
+      volcModelNameMessage: '請輸入模型名稱！格式：{"模型名稱":"EndpointID"}',
+      addVolcEngineAK: '火山 ACCESS_KEY',
+      volcAKMessage: '請輸入VOLC_ACCESS_KEY',
+      addVolcEngineSK: '火山 SECRET_KEY',
+      volcSKMessage: '請輸入VOLC_SECRET_KEY',
     },
     message: {
       registered: '註冊成功',

web/src/locales/zh.ts

@@ -458,6 +458,11 @@ export default {
       modelTypeMessage: '请输入模型类型！',
       baseUrlNameMessage: '请输入基础 Url！',
       ollamaLink: '如何集成 {{name}}',
+      volcModelNameMessage: '请输入模型名称！格式：{"模型名称":"EndpointID"}',
+      addVolcEngineAK: '火山 ACCESS_KEY',
+      volcAKMessage: '请输入VOLC_ACCESS_KEY',
+      addVolcEngineSK: '火山 SECRET_KEY',
+      volcSKMessage: '请输入VOLC_SECRET_KEY',
     },
     message: {
       registered: '注册成功',

web/src/pages/user-setting/setting-model/hooks.ts

@@ -166,6 +166,41 @@ export const useSubmitOllama = () => {
   };
 };
 
+export const useSubmitVolcEngine = () => {
+  const loading = useOneNamespaceEffectsLoading('settingModel', ['add_llm']);
+  const [selectedVolcFactory, setSelectedVolcFactory] = useState<string>('');
+  const addLlm = useAddLlm();
+  const {
+    visible: volcAddingVisible,
+    hideModal: hideVolcAddingModal,
+    showModal: showVolcAddingModal,
+  } = useSetModalState();
+
+  const onVolcAddingOk = useCallback(
+    async (payload: IAddLlmRequestBody) => {
+      const ret = await addLlm(payload);
+      if (ret === 0) {
+        hideVolcAddingModal();
+      }
+    },
+    [hideVolcAddingModal, addLlm],
+  );
+
+  const handleShowVolcAddingModal = (llmFactory: string) => {
+    setSelectedVolcFactory(llmFactory);
+    showVolcAddingModal();
+  };
+
+  return {
+    volcAddingLoading: loading,
+    onVolcAddingOk,
+    volcAddingVisible,
+    hideVolcAddingModal,
+    showVolcAddingModal: handleShowVolcAddingModal,
+    selectedVolcFactory,
+  };
+};
+
 export const useHandleDeleteLlm = (llmFactory: string) => {
   const deleteLlm = useDeleteLlm();
   const showDeleteConfirm = useShowDeleteConfirm();

web/src/pages/user-setting/setting-model/index.tsx

@@ -37,10 +37,12 @@ import {
   useSelectModelProvidersLoading,
   useSubmitApiKey,
   useSubmitOllama,
+  useSubmitVolcEngine,
   useSubmitSystemModelSetting,
 } from './hooks';
 import styles from './index.less';
 import OllamaModal from './ollama-modal';
+import VolcEngineModal from "./volcengine-model";
 import SystemModelSettingModal from './system-model-setting-modal';
 
 const IconMap = {
@@ -52,6 +54,7 @@ const IconMap = {
   Ollama: 'ollama',
   Xinference: 'xinference',
   DeepSeek: 'deepseek',
+  VolcEngine: 'volc_engine',
 };
 
 const LlmIcon = ({ name }: { name: string }) => {
@@ -165,6 +168,15 @@ const UserSettingModel = () => {
     selectedLlmFactory,
   } = useSubmitOllama();
 
+  const {
+    volcAddingVisible,
+    hideVolcAddingModal,
+    showVolcAddingModal,
+    onVolcAddingOk,
+    volcAddingLoading,
+    selectedVolcFactory,
+  } = useSubmitVolcEngine();
+
   const handleApiKeyClick = useCallback(
     (llmFactory: string) => {
       if (isLocalLlmFactory(llmFactory)) {
@@ -179,6 +191,8 @@ const UserSettingModel = () => {
   const handleAddModel = (llmFactory: string) => () => {
     if (isLocalLlmFactory(llmFactory)) {
       showLlmAddingModal(llmFactory);
+    } else if (llmFactory === 'VolcEngine') {
+      showVolcAddingModal('VolcEngine');
     } else {
       handleApiKeyClick(llmFactory);
     }
@@ -270,6 +284,13 @@ const UserSettingModel = () => {
         loading={llmAddingLoading}
         llmFactory={selectedLlmFactory}
       ></OllamaModal>
+      <VolcEngineModal
+        visible={volcAddingVisible}
+        hideModal={hideVolcAddingModal}
+        onOk={onVolcAddingOk}
+        loading={volcAddingLoading}
+        llmFactory={selectedVolcFactory}
+      ></VolcEngineModal>
     </section>
   );
 };

web/src/pages/user-setting/setting-model/volcengine-model/index.tsx

@@ -0,0 +1,118 @@
+import { useTranslate } from '@/hooks/commonHooks';
+import { IModalProps } from '@/interfaces/common';
+import { IAddLlmRequestBody } from '@/interfaces/request/llm';
+import { Flex, Form, Input, Modal, Select, Space, Switch } from 'antd';
+import omit from 'lodash/omit';
+
+type FieldType = IAddLlmRequestBody & { vision: boolean };
+
+const { Option } = Select;
+
+const VolcEngineModal = ({
+  visible,
+  hideModal,
+  onOk,
+  loading,
+  llmFactory
+}: IModalProps<IAddLlmRequestBody> & { llmFactory: string }) => {
+  const [form] = Form.useForm<FieldType>();
+
+  const { t } = useTranslate('setting');
+
+  const handleOk = async () => {
+    const values = await form.validateFields();
+    const modelType =
+      values.model_type === 'chat' && values.vision
+        ? 'image2text'
+        : values.model_type;
+
+    const data = {
+      ...omit(values, ['vision']),
+      model_type: modelType,
+      llm_factory: llmFactory,
+    };
+    console.info(data);
+
+    onOk?.(data);
+  };
+
+  return (
+    <Modal
+      title={t('addLlmTitle', { name: llmFactory })}
+      open={visible}
+      onOk={handleOk}
+      onCancel={hideModal}
+      okButtonProps={{ loading }}
+      footer={(originNode: React.ReactNode) => {
+        return (
+          <Flex justify={'space-between'}>
+            <a
+              href="https://www.volcengine.com/docs/82379/1095322"
+              target="_blank"
+              rel="noreferrer"
+            >
+              {t('ollamaLink', { name: llmFactory })}
+            </a>
+            <Space>{originNode}</Space>
+          </Flex>
+        );
+      }}
+    >
+      <Form
+        name="basic"
+        style={{ maxWidth: 600 }}
+        autoComplete="off"
+        layout={'vertical'}
+        form={form}
+      >
+        <Form.Item<FieldType>
+          label={t('modelType')}
+          name="model_type"
+          initialValue={'chat'}
+          rules={[{ required: true, message: t('modelTypeMessage') }]}
+        >
+          <Select placeholder={t('modelTypeMessage')}>
+            <Option value="chat">chat</Option>
+            <Option value="embedding">embedding</Option>
+          </Select>
+        </Form.Item>
+        <Form.Item<FieldType>
+          label={t('modelName')}
+          name="llm_name"
+          rules={[{ required: true, message: t('volcModelNameMessage') }]}
+        >
+          <Input placeholder={t('volcModelNameMessage')} />
+        </Form.Item>
+        <Form.Item<FieldType>
+          label={t('addVolcEngineAK')}
+          name="volc_ak"
+          rules={[{ required: true, message: t('volcAKMessage') }]}
+        >
+          <Input placeholder={t('volcAKMessage')} />
+        </Form.Item>
+        <Form.Item<FieldType>
+          label={t('addVolcEngineSK')}
+          name="volc_sk"
+          rules={[{ required: true, message: t('volcAKMessage') }]}
+        >
+          <Input placeholder={t('volcAKMessage')} />
+        </Form.Item>
+        <Form.Item noStyle dependencies={['model_type']}>
+          {({ getFieldValue }) =>
+            getFieldValue('model_type') === 'chat' && (
+              <Form.Item
+                label={t('vision')}
+                valuePropName="checked"
+                name={'vision'}
+              >
+                <Switch />
+              </Form.Item>
+            )
+          }
+        </Form.Item>
+      </Form>
+    </Modal>
+  );
+};
+
+export default VolcEngineModal;