diff --git a/backend/apps/ai_model/model_factory.py b/backend/apps/ai_model/model_factory.py index 6391501c..f7e5817d 100644 --- a/backend/apps/ai_model/model_factory.py +++ b/backend/apps/ai_model/model_factory.py @@ -11,7 +11,7 @@ from apps.ai_model.openai.llm import BaseChatOpenAI 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 @@ class BaseLLM(ABC): """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, diff --git a/backend/apps/system/api/aimodel.py b/backend/apps/system/api/aimodel.py index 9e949af6..6eea3549 100644 --- a/backend/apps/system/api/aimodel.py +++ b/backend/apps/system/api/aimodel.py @@ -20,7 +20,7 @@ async def check_llm(info: AiModelCreator, trans: Trans): 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}%")) diff --git a/backend/apps/system/schemas/ai_model_schema.py b/backend/apps/system/schemas/ai_model_schema.py index 5bfd206b..8b5946bf 100644 --- a/backend/apps/system/schemas/ai_model_schema.py +++ b/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): diff --git a/frontend/src/assets/model/icon_vllm_colorful.png b/frontend/src/assets/model/icon_vllm_colorful.png new file mode 100644 index 00000000..aefd0525 Binary files /dev/null and b/frontend/src/assets/model/icon_vllm_colorful.png differ diff --git a/frontend/src/entity/supplier.ts b/frontend/src/entity/supplier.ts index 2c8a271f..11932f85 100644 --- a/frontend/src/entity/supplier.ts +++ b/frontend/src/entity/supplier.ts @@ -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) => { diff --git a/frontend/src/views/system/model/ModelForm.vue b/frontend/src/views/system/model/ModelForm.vue index 501865d0..3a211d6c 100644 --- a/frontend/src/views/system/model/ModelForm.vue +++ b/frontend/src/views/system/model/ModelForm.vue @@ -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) => { 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)