From e894dca586a5d2de0501d97235c3dbbbb5c5d8df Mon Sep 17 00:00:00 2001 From: CaptainB Date: Thu, 19 Dec 2024 18:39:48 +0800 Subject: [PATCH] feat: Support gemini embedding model --- .../credential/embedding.py | 45 +++++++++++++++++++ .../gemini_model_provider.py | 12 +++++ .../gemini_model_provider/model/embedding.py | 22 +++++++++ 3 files changed, 79 insertions(+) create mode 100644 apps/setting/models_provider/impl/gemini_model_provider/credential/embedding.py create mode 100644 apps/setting/models_provider/impl/gemini_model_provider/model/embedding.py diff --git a/apps/setting/models_provider/impl/gemini_model_provider/credential/embedding.py b/apps/setting/models_provider/impl/gemini_model_provider/credential/embedding.py new file mode 100644 index 00000000000..5fc0a9c880f --- /dev/null +++ b/apps/setting/models_provider/impl/gemini_model_provider/credential/embedding.py @@ -0,0 +1,45 @@ +# coding=utf-8 +""" + @project: MaxKB + @Author：虎 + @file： embedding.py + @date：2024/7/12 16:45 + @desc: +""" +from typing import Dict + +from common import forms +from common.exception.app_exception import AppApiException +from common.forms import BaseForm +from setting.models_provider.base_model_provider import BaseModelCredential, ValidCode + + +class GeminiEmbeddingCredential(BaseForm, BaseModelCredential): + def is_valid(self, model_type: str, model_name, model_credential: Dict[str, object], provider, + raise_exception=True): + model_type_list = provider.get_model_type_list() + if not any(list(filter(lambda mt: mt.get('value') == model_type, model_type_list))): + raise AppApiException(ValidCode.valid_error.value, f'{model_type} 模型类型不支持') + + for key in ['api_key']: + if key not in model_credential: + if raise_exception: + raise AppApiException(ValidCode.valid_error.value, f'{key} 字段为必填字段') + else: + return False + try: + model = provider.get_model(model_type, model_name, model_credential) + model.embed_query('你好') + except Exception as e: + if isinstance(e, AppApiException): + raise e + if raise_exception: + raise AppApiException(ValidCode.valid_error.value, f'校验失败,请检查参数是否正确: {str(e)}') + else: + return False + return True + + def encryption_dict(self, model: Dict[str, object]): + return {**model, 'api_key': super().encryption(model.get('api_key', ''))} + + api_key = forms.PasswordInputField('API Key', required=True) diff --git a/apps/setting/models_provider/impl/gemini_model_provider/gemini_model_provider.py b/apps/setting/models_provider/impl/gemini_model_provider/gemini_model_provider.py index 2556328d798..e8adb808211 100644 --- a/apps/setting/models_provider/impl/gemini_model_provider/gemini_model_provider.py +++ b/apps/setting/models_provider/impl/gemini_model_provider/gemini_model_provider.py @@ -11,9 +11,11 @@ from common.util.file_util import get_file_content from setting.models_provider.base_model_provider import IModelProvider, ModelProvideInfo, ModelInfo, ModelTypeConst, \ ModelInfoManage +from setting.models_provider.impl.gemini_model_provider.credential.embedding import GeminiEmbeddingCredential from setting.models_provider.impl.gemini_model_provider.credential.image import GeminiImageModelCredential from setting.models_provider.impl.gemini_model_provider.credential.llm import GeminiLLMModelCredential from setting.models_provider.impl.gemini_model_provider.credential.stt import GeminiSTTModelCredential +from setting.models_provider.impl.gemini_model_provider.model.embedding import GeminiEmbeddingModel from setting.models_provider.impl.gemini_model_provider.model.image import GeminiImage from setting.models_provider.impl.gemini_model_provider.model.llm import GeminiChatModel from setting.models_provider.impl.gemini_model_provider.model.stt import GeminiSpeechToText @@ -22,6 +24,7 @@ gemini_llm_model_credential = GeminiLLMModelCredential() gemini_image_model_credential = GeminiImageModelCredential() gemini_stt_model_credential = GeminiSTTModelCredential() +gemini_embedding_model_credential = GeminiEmbeddingCredential() model_info_list = [ ModelInfo('gemini-1.0-pro', '最新的Gemini 1.0 Pro模型，随Google更新而更新', @@ -56,14 +59,23 @@ GeminiSpeechToText), ] +model_embedding_info_list = [ + ModelInfo('models/embedding-001', '', + ModelTypeConst.EMBEDDING, + gemini_embedding_model_credential, + GeminiEmbeddingModel), +] + model_info_manage = ( ModelInfoManage.builder() .append_model_info_list(model_info_list) .append_model_info_list(model_image_info_list) .append_model_info_list(model_stt_info_list) + .append_model_info_list(model_embedding_info_list) .append_default_model_info(model_info_list[0]) .append_default_model_info(model_image_info_list[0]) .append_default_model_info(model_stt_info_list[0]) + .append_default_model_info(model_embedding_info_list[0]) .build() ) diff --git a/apps/setting/models_provider/impl/gemini_model_provider/model/embedding.py b/apps/setting/models_provider/impl/gemini_model_provider/model/embedding.py new file mode 100644 index 00000000000..7474023377e --- /dev/null +++ b/apps/setting/models_provider/impl/gemini_model_provider/model/embedding.py @@ -0,0 +1,22 @@ +# coding=utf-8 +""" + @project: MaxKB + @Author：虎 + @file： embedding.py + @date：2024/7/12 17:44 + @desc: +""" +from typing import Dict + +from langchain_google_genai import GoogleGenerativeAIEmbeddings + +from setting.models_provider.base_model_provider import MaxKBBaseModel + + +class GeminiEmbeddingModel(MaxKBBaseModel, GoogleGenerativeAIEmbeddings): + @staticmethod + def new_instance(model_type, model_name, model_credential: Dict[str, object], **model_kwargs): + return GoogleGenerativeAIEmbeddings( + google_api_key=model_credential.get('api_key'), + model=model_name, + )