Merge branch 'main' of https://github.com/dataease/SQLBot

Author: dataeaseShu

backend/apps/ai_model/model_factory.py

@@ -11,7 +11,7 @@
 from apps.system.models.system_model import AiModelDetail
 from common.core.db import engine
 from common.utils.utils import prepare_model_arg
-
+from langchain_community.llms import VLLMOpenAI
 
 # from langchain_community.llms import Tongyi, VLLM
 
@@ -60,7 +60,14 @@ def llm(self) -> BaseChatModel:
         """Return the langchain LLM instance"""
         return self._llm
 
-
+class OpenAIvLLM(BaseLLM):
+    def _init_llm(self) -> VLLMOpenAI:
+        return VLLMOpenAI(
+            openai_api_key=self.config.api_key or 'Empty',
+            openai_api_base=self.config.api_base_url,
+            model_name=self.config.model_name,
+            **self.config.additional_params,
+        )
 class OpenAILLM(BaseLLM):
     def _init_llm(self) -> BaseChatModel:
         return BaseChatOpenAI(
@@ -81,7 +88,7 @@ class LLMFactory:
     _llm_types: Dict[str, Type[BaseLLM]] = {
         "openai": OpenAILLM,
         "tongyi": OpenAILLM,
-        "vllm": OpenAILLM
+        "vllm": OpenAIvLLM
     }
 
     @classmethod
@@ -129,7 +136,7 @@ def get_default_config() -> LLMConfig:
         # 构造 LLMConfig
         return LLMConfig(
             model_id=db_model.id,
-            model_type="openai",
+            model_type="openai" if db_model.protocol == 1 else "vllm",
             model_name=db_model.base_model,
             api_key=db_model.api_key,
             api_base_url=db_model.api_domain,

backend/apps/system/api/aimodel.py

@@ -20,7 +20,7 @@ async def generate():
         try:
             additional_params = {item.key: prepare_model_arg(item.val) for item in info.config_list if item.key and item.val}
             config = LLMConfig(
-                model_type="openai",
+                model_type="openai" if info.protocol == 1 else "vllm",
                 model_name=info.base_model,
                 api_key=info.api_key,
                 api_base_url=info.api_domain,
@@ -74,7 +74,8 @@ async def query(
                        AiModelDetail.name, 
                        AiModelDetail.model_type, 
                        AiModelDetail.base_model, 
-                       AiModelDetail.supplier, 
+                       AiModelDetail.supplier,
+                       AiModelDetail.protocol, 
                        AiModelDetail.default_model)
     if keyword is not None:
         statement = statement.where(AiModelDetail.name.like(f"%{keyword}%"))

backend/apps/system/schemas/ai_model_schema.py

@@ -9,6 +9,7 @@ class AiModelItem(BaseModel):
     model_type: int
     base_model: str
     supplier: int
+    protocol: int
     default_model: bool = False
 
 class AiModelGridItem(AiModelItem, BaseCreatorDTO):

frontend/src/assets/model/icon_vllm_colorful.png

@@ -8,18 +8,25 @@ import icon_openai_colorful from '@/assets/model/icon_openai_colorful.png'
 import icon_kimi_colorful from '@/assets/model/icon_kimi_colorful.png'
 import icon_txhy_colorful from '@/assets/model/icon_txhy_colorful.png'
 import icon_hsyq_colorful from '@/assets/model/icon_hsyq_colorful.png'
+import icon_vllm_colorful from '@/assets/model/icon_vllm_colorful.png'
 
 type ModelArg = { key: string; val?: string | number; type: string; range?: string }
 type ModelOption = { name: string; api_domain?: string; args?: ModelArg[] }
 type ModelConfig = Record<
   number,
-  { api_domain: string; common_args?: ModelArg[]; model_options: ModelOption[] }
+  {
+    api_domain: string
+    common_args?: ModelArg[]
+    model_options: ModelOption[]
+  }
 >
 
 export const supplierList: Array<{
   id: number
   name: string
   icon: any
+  type?: string
+  is_private?: boolean
   model_config: ModelConfig
 }> = [
   {
@@ -259,6 +266,20 @@ export const supplierList: Array<{
       },
     },
   },
+  {
+    id: 11,
+    name: 'vLLM',
+    icon: icon_vllm_colorful,
+    type: 'vllm',
+    is_private: true,
+    model_config: {
+      0: {
+        api_domain: 'http://127.0.0.1:8000/v1',
+        common_args: [{ key: 'temperature', val: 0.6, type: 'number', range: '[0, 1]' }],
+        model_options: [],
+      },
+    },
+  },
 ]
 
 export const base_model_options = (supplier_id: number, model_type?: number) => {

frontend/src/entity/supplier.ts

@@ -38,6 +38,7 @@ const modelForm = reactive({
   api_key: '',
   api_domain: '',
   config_list: [],
+  protocol: 1,
 })
 const isCreate = ref(false)
 const modelRef = ref()
@@ -90,7 +91,7 @@ const handleCurrentChange = (val: any) => {
   currentPage.value = val
 }
 
-const rules = {
+const rules = computed(() => ({
   model_type: [
     {
       required: true,
@@ -109,12 +110,12 @@ const rules = {
   name: [{ required: true, message: t('model.the_basic_model'), trigger: 'blur' }],
   api_key: [
     {
-      required: true,
+      required: !currentSupplier.value?.is_private,
       message: t('datasource.please_enter') + t('common.empty') + 'API Key',
       trigger: 'blur',
     },
   ],
-}
+}))
 
 onMounted(() => {
   setTimeout(() => {
@@ -169,6 +170,7 @@ const supplierChang = (supplier: any) => {
   const config = supplier.model_config[modelForm.model_type || 0]
   modelForm.api_domain = config.api_domain
   modelForm.base_model = ''
+  modelForm.protocol = supplier.type === 'vllm' ? 2 : 1
   advancedSetting.value = []
 }
 let curId = +new Date()
@@ -201,7 +203,7 @@ const formatAdvancedSetting = (list: Array<any>) => {
   advancedSetting.value = setting_list
 }
 const baseModelChange = (val: string) => {
-  if (!val || !modelForm.supplier || !modelList.value?.length) {
+  if (!val || !modelForm.supplier) {
     return
   }
   const current_model = modelList.value?.find((model: any) => model.name == val)