feat: add QfRerankerCredential and QfBgeReranker classes for document reranking

Author: wxg0103

apps/models_provider/impl/aliyun_bai_lian_model_provider/model/tts.py

@@ -26,7 +26,7 @@ def is_cache_model():
 
     @staticmethod
     def new_instance(model_type, model_name, model_credential: Dict[str, object], **model_kwargs):
-        optional_params = {'params': {'voice': 'longxiaochun', 'speech_rate': 1.0}}
+        optional_params = {'params': {}}
         for key, value in model_kwargs.items():
             if key not in ['model_id', 'use_local', 'streaming']:
                 optional_params['params'][key] = value

apps/models_provider/impl/wenxin_model_provider/credential/reranker.py

@@ -0,0 +1,50 @@
+from typing import Dict
+
+from langchain_core.documents import Document
+
+from common import forms
+from common.exception.app_exception import AppApiException
+from common.forms import BaseForm
+from models_provider.base_model_provider import BaseModelCredential, ValidCode
+from django.utils.translation import gettext_lazy as _
+from common.utils.logger import maxkb_logger
+from models_provider.impl.wenxin_model_provider.model.reranker import QfBgeReranker
+
+
+class QfRerankerCredential(BaseForm, BaseModelCredential):
+    api_url = forms.TextInputField('API URL', required=True)
+    api_key = forms.PasswordInputField('API Key', required=True)
+
+    def is_valid(self, model_type: str, model_name, model_credential: Dict[str, object], model_params, 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,
+                                  _('{model_type} Model type is not supported').format(model_type=model_type))
+
+        for key in ['api_url', 'api_key']:
+            if key not in model_credential:
+                if raise_exception:
+                    raise AppApiException(ValidCode.valid_error.value, _('{key}  is required').format(key=key))
+                else:
+                    return False
+        try:
+            model: QfBgeReranker = provider.get_model(model_type, model_name, model_credential)
+            test_text = str(_('Hello'))
+            model.compress_documents([Document(page_content=test_text)], test_text)
+        except Exception as e:
+            maxkb_logger.error(f'Exception: {e}', exc_info=True)
+            if isinstance(e, AppApiException):
+                raise e
+            if raise_exception:
+                raise AppApiException(
+                    ValidCode.valid_error.value,
+                    _('Verification failed, please check whether the parameters are correct: {error}').format(
+                        error=str(e))
+                )
+            return False
+
+        return True
+
+    def encryption_dict(self, model_info: Dict[str, object]):
+        return {**model_info, 'api_key': super().encryption(model_info.get('api_key', ''))}

apps/models_provider/impl/wenxin_model_provider/model/reranker.py

@@ -0,0 +1,75 @@
+import json
+from typing import Sequence, Optional, Dict, Any
+
+import requests
+from langchain_core.callbacks import Callbacks
+from langchain_core.documents import BaseDocumentCompressor, Document
+
+from models_provider.base_model_provider import MaxKBBaseModel
+
+
+class QfBgeReranker(MaxKBBaseModel, BaseDocumentCompressor):
+    api_key: str
+    api_url: str
+    model: str
+    params: dict
+    top_n: int = 3
+
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        self.api_key = kwargs.get('api_key')
+        self.model = kwargs.get('model')
+        self.params = kwargs.get('params', {})
+        self.api_url = kwargs.get('api_url')
+        self.top_n = self.params.get('top_n', 3)
+
+    @staticmethod
+    def is_cache_model():
+        return False
+
+    @staticmethod
+    def new_instance(model_type, model_name, model_credential: Dict[str, object], **model_kwargs):
+        return QfBgeReranker(
+            model=model_name,
+            api_key=model_credential.get('api_key'),
+            api_url=model_credential.get('api_url'),
+            params=model_kwargs,
+        )
+
+    def compress_documents(
+            self,
+            documents: Sequence[Document],
+            query: str,
+            callbacks: Optional[Callbacks] = None
+    ) -> Sequence[Document]:
+        if not documents:
+            return []
+
+        texts = [doc.page_content for doc in documents]
+
+        headers = {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json"
+        }
+        top_n = min(self.top_n, len(texts))
+        payload = {
+            "model": self.model,
+            "query": query,
+            "documents": texts,
+            "top_n": top_n
+        }
+
+        response = requests.post(f"{self.api_url}/rerank", json=payload, headers=headers)
+
+        if response.status_code != 200:
+            raise RuntimeError(f"千帆 API 请求失败：{response.text}")
+
+        res = response.json()
+
+        return [
+            Document(
+                page_content=item.get('document', ''),
+                metadata={'relevance_score': item.get('relevance_score')}
+            )
+            for item in res.get('results', [])
+        ]

apps/models_provider/impl/wenxin_model_provider/wenxin_model_provider.py

@@ -13,13 +13,17 @@
     ModelInfoManage
 from models_provider.impl.wenxin_model_provider.credential.embedding import QianfanEmbeddingCredential
 from models_provider.impl.wenxin_model_provider.credential.llm import WenxinLLMModelCredential
+from models_provider.impl.wenxin_model_provider.credential.reranker import QfRerankerCredential
 from models_provider.impl.wenxin_model_provider.model.embedding import QianfanEmbeddings
 from models_provider.impl.wenxin_model_provider.model.llm import QianfanChatModel
 from maxkb.conf import PROJECT_DIR
 from django.utils.translation import gettext as _
 
+from models_provider.impl.wenxin_model_provider.model.reranker import QfBgeReranker
+
 win_xin_llm_model_credential = WenxinLLMModelCredential()
 qianfan_embedding_credential = QianfanEmbeddingCredential()
+qf_reranker_credential = QfRerankerCredential()
 model_info_list = [ModelInfo('ERNIE-Bot-4',
                              _('ERNIE-Bot-4 is a large language model independently developed by Baidu. It covers massive Chinese data and has stronger capabilities in dialogue Q&A, content creation and generation.'),
                              ModelTypeConst.LLM, win_xin_llm_model_credential, QianfanChatModel),
@@ -46,13 +50,19 @@
                              ModelInfo('bge-large-zh', '', ModelTypeConst.EMBEDDING, qianfan_embedding_credential,
                                        QianfanEmbeddings)
                              ]
-model_info_manage = ModelInfoManage.builder().append_model_info_list(model_info_list).append_default_model_info(
+rerank_model_info_list = [ModelInfo('bce-reranker-base',
+                                    _(''),
+                                    ModelTypeConst.RERANKER, qf_reranker_credential, QfBgeReranker),
+                          ]
+model_info_manage = (ModelInfoManage.builder().append_model_info_list(model_info_list).append_default_model_info(
     ModelInfo('ERNIE-Bot-4',
               _('ERNIE-Bot-4 is a large language model independently developed by Baidu. It covers massive Chinese data and has stronger capabilities in dialogue Q&A, content creation and generation.'),
               ModelTypeConst.LLM,
               win_xin_llm_model_credential,
               QianfanChatModel)).append_model_info_list(embedding_model_info_list).append_default_model_info(
-    embedding_model_info_list[0]).build()
+    embedding_model_info_list[0]).
+                     append_model_info_list(rerank_model_info_list).append_default_model_info(
+    rerank_model_info_list[0]).build())
 
 
 class WenxinModelProvider(IModelProvider):