Feat/add retriever rerank (#1560)

Co-authored-by: jyong <jyong@dify.ai>

api/core/model_providers/model_factory.py

@@ -9,6 +9,7 @@ from core.model_providers.models.embedding.base import BaseEmbedding
 from core.model_providers.models.entity.model_params import ModelKwargs, ModelType
 from core.model_providers.models.llm.base import BaseLLM
 from core.model_providers.models.moderation.base import BaseModeration
+from core.model_providers.models.reranking.base import BaseReranking
 from core.model_providers.models.speech2text.base import BaseSpeech2Text
 from extensions.ext_database import db
 from models.provider import TenantDefaultModel
@@ -140,6 +141,44 @@ class ModelFactory:
             name=model_name
         )
 
+
+    @classmethod
+    def get_reranking_model(cls,
+                            tenant_id: str,
+                            model_provider_name: Optional[str] = None,
+                            model_name: Optional[str] = None) -> Optional[BaseReranking]:
+        """
+        get reranking model.
+
+        :param tenant_id: a string representing the ID of the tenant.
+        :param model_provider_name:
+        :param model_name:
+        :return:
+        """
+        if model_provider_name is None and model_name is None:
+            default_model = cls.get_default_model(tenant_id, ModelType.RERANKING)
+
+            if not default_model:
+                raise LLMBadRequestError(f"Default model is not available. "
+                                         f"Please configure a Default Reranking Model "
+                                         f"in the Settings -> Model Provider.")
+
+            model_provider_name = default_model.provider_name
+            model_name = default_model.model_name
+
+        # get model provider
+        model_provider = ModelProviderFactory.get_preferred_model_provider(tenant_id, model_provider_name)
+
+        if not model_provider:
+            raise ProviderTokenNotInitError(f"Model {model_name} provider credentials is not initialized.")
+
+        # init reranking model
+        model_class = model_provider.get_model_class(model_type=ModelType.RERANKING)
+        return model_class(
+            model_provider=model_provider,
+            name=model_name
+        )
+
     @classmethod
     def get_speech2text_model(cls,
                               tenant_id: str,

api/core/model_providers/model_provider_factory.py

@@ -72,6 +72,9 @@ class ModelProviderFactory:
         elif provider_name == 'localai':
             from core.model_providers.providers.localai_provider import LocalAIProvider
             return LocalAIProvider
+        elif provider_name == 'cohere':
+            from core.model_providers.providers.cohere_provider import CohereProvider
+            return CohereProvider
         else:
             raise NotImplementedError
 

api/core/model_providers/models/entity/model_params.py

@@ -17,7 +17,7 @@ class ModelType(enum.Enum):
     IMAGE = 'image'
     VIDEO = 'video'
     MODERATION = 'moderation'
-
+    RERANKING = 'reranking'
     @staticmethod
     def value_of(value):
         for member in ModelType:

api/core/model_providers/models/reranking/base.py

@@ -0,0 +1,36 @@
+from abc import abstractmethod
+from typing import Any, Optional, List
+from langchain.schema import Document
+
+from core.model_providers.models.base import BaseProviderModel
+from core.model_providers.models.entity.model_params import ModelType
+from core.model_providers.providers.base import BaseModelProvider
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class BaseReranking(BaseProviderModel):
+    name: str
+    type: ModelType = ModelType.RERANKING
+
+    def __init__(self, model_provider: BaseModelProvider, client: Any, name: str):
+        super().__init__(model_provider, client)
+        self.name = name
+
+    @property
+    def base_model_name(self) -> str:
+        """
+        get base model name
+        
+        :return: str
+        """
+        return self.name
+
+    @abstractmethod
+    def rerank(self, query: str, documents: List[Document], score_threshold: Optional[float], top_k: Optional[int]) -> Optional[List[Document]]:
+        raise NotImplementedError
+
+    @abstractmethod
+    def handle_exceptions(self, ex: Exception) -> Exception:
+        raise NotImplementedError

api/core/model_providers/models/reranking/cohere_reranking.py

@@ -0,0 +1,73 @@
+import logging
+from typing import Optional, List
+
+import cohere
+import openai
+from langchain.schema import Document
+
+from core.model_providers.error import LLMBadRequestError, LLMAPIConnectionError, LLMAPIUnavailableError, \
+    LLMRateLimitError, LLMAuthorizationError
+from core.model_providers.models.reranking.base import BaseReranking
+from core.model_providers.providers.base import BaseModelProvider
+
+
+class CohereReranking(BaseReranking):
+
+    def __init__(self, model_provider: BaseModelProvider, name: str):
+        self.credentials = model_provider.get_model_credentials(
+            model_name=name,
+            model_type=self.type
+        )
+
+        client = cohere.Client(self.credentials.get('api_key'))
+
+        super().__init__(model_provider, client, name)
+
+    def rerank(self, query: str, documents: List[Document], score_threshold: Optional[float], top_k: Optional[int]) -> Optional[List[Document]]:
+        docs = []
+        doc_id = []
+        for document in documents:
+            if document.metadata['doc_id'] not in doc_id:
+                doc_id.append(document.metadata['doc_id'])
+                docs.append(document.page_content)
+        results = self.client.rerank(query=query, documents=docs, model=self.name, top_n=top_k)
+        rerank_documents = []
+
+        for idx, result in enumerate(results):
+            # format document
+            rerank_document = Document(
+                page_content=result.document['text'],
+                metadata={
+                    "doc_id": documents[result.index].metadata['doc_id'],
+                    "doc_hash": documents[result.index].metadata['doc_hash'],
+                    "document_id": documents[result.index].metadata['document_id'],
+                    "dataset_id": documents[result.index].metadata['dataset_id'],
+                    'score': result.relevance_score
+                }
+            )
+            # score threshold check
+            if score_threshold is not None:
+                if result.relevance_score >= score_threshold:
+                    rerank_documents.append(rerank_document)
+            else:
+                rerank_documents.append(rerank_document)
+        return rerank_documents
+
+    def handle_exceptions(self, ex: Exception) -> Exception:
+        if isinstance(ex, openai.error.InvalidRequestError):
+            logging.warning("Invalid request to OpenAI API.")
+            return LLMBadRequestError(str(ex))
+        elif isinstance(ex, openai.error.APIConnectionError):
+            logging.warning("Failed to connect to OpenAI API.")
+            return LLMAPIConnectionError(ex.__class__.__name__ + ":" + str(ex))
+        elif isinstance(ex, (openai.error.APIError, openai.error.ServiceUnavailableError, openai.error.Timeout)):
+            logging.warning("OpenAI service unavailable.")
+            return LLMAPIUnavailableError(ex.__class__.__name__ + ":" + str(ex))
+        elif isinstance(ex, openai.error.RateLimitError):
+            return LLMRateLimitError(str(ex))
+        elif isinstance(ex, openai.error.AuthenticationError):
+            return LLMAuthorizationError(str(ex))
+        elif isinstance(ex, openai.error.OpenAIError):
+            return LLMBadRequestError(ex.__class__.__name__ + ":" + str(ex))
+        else:
+            return ex

api/core/model_providers/providers/cohere_provider.py

@@ -0,0 +1,152 @@
+import json
+from json import JSONDecodeError
+from typing import Type
+
+from langchain.schema import HumanMessage
+
+from core.helper import encrypter
+from core.model_providers.models.base import BaseProviderModel
+from core.model_providers.models.entity.model_params import ModelKwargsRules, KwargRule, ModelType, ModelMode
+from core.model_providers.models.reranking.cohere_reranking import CohereReranking
+from core.model_providers.providers.base import BaseModelProvider, CredentialsValidateFailedError
+from models.provider import ProviderType
+
+
+class CohereProvider(BaseModelProvider):
+
+    @property
+    def provider_name(self):
+        """
+        Returns the name of a provider.
+        """
+        return 'cohere'
+    
+    def _get_text_generation_model_mode(self, model_name) -> str:
+        return ModelMode.CHAT.value
+
+    def _get_fixed_model_list(self, model_type: ModelType) -> list[dict]:
+        if model_type == ModelType.RERANKING:
+            return [
+                {
+                    'id': 'rerank-english-v2.0',
+                    'name': 'rerank-english-v2.0'
+                },
+                {
+                    'id': 'rerank-multilingual-v2.0',
+                    'name': 'rerank-multilingual-v2.0'
+                }
+            ]
+        else:
+            return []
+
+    def get_model_class(self, model_type: ModelType) -> Type[BaseProviderModel]:
+        """
+        Returns the model class.
+
+        :param model_type:
+        :return:
+        """
+        if model_type == ModelType.RERANKING:
+            model_class = CohereReranking
+        else:
+            raise NotImplementedError
+
+        return model_class
+
+    def get_model_parameter_rules(self, model_name: str, model_type: ModelType) -> ModelKwargsRules:
+        """
+        get model parameter rules.
+
+        :param model_name:
+        :param model_type:
+        :return:
+        """
+        return ModelKwargsRules(
+            temperature=KwargRule[float](min=0, max=1, default=0.3, precision=2),
+            top_p=KwargRule[float](min=0, max=0.99, default=0.85, precision=2),
+            presence_penalty=KwargRule[float](enabled=False),
+            frequency_penalty=KwargRule[float](enabled=False),
+            max_tokens=KwargRule[int](enabled=False),
+        )
+
+    @classmethod
+    def is_provider_credentials_valid_or_raise(cls, credentials: dict):
+        """
+        Validates the given credentials.
+        """
+        if 'api_key' not in credentials:
+            raise CredentialsValidateFailedError('Cohere api_key must be provided.')
+
+        try:
+            credential_kwargs = {
+                'api_key': credentials['api_key'],
+            }
+            # todo validate
+        except Exception as ex:
+            raise CredentialsValidateFailedError(str(ex))
+
+    @classmethod
+    def encrypt_provider_credentials(cls, tenant_id: str, credentials: dict) -> dict:
+        credentials['api_key'] = encrypter.encrypt_token(tenant_id, credentials['api_key'])
+        return credentials
+
+    def get_provider_credentials(self, obfuscated: bool = False) -> dict:
+        if self.provider.provider_type == ProviderType.CUSTOM.value:
+            try:
+                credentials = json.loads(self.provider.encrypted_config)
+            except JSONDecodeError:
+                credentials = {
+                    'api_key': None,
+                }
+
+            if credentials['api_key']:
+                credentials['api_key'] = encrypter.decrypt_token(
+                    self.provider.tenant_id,
+                    credentials['api_key']
+                )
+
+                if obfuscated:
+                    credentials['api_key'] = encrypter.obfuscated_token(credentials['api_key'])
+
+            return credentials
+        else:
+            return {}
+
+    def should_deduct_quota(self):
+        return True
+
+    @classmethod
+    def is_model_credentials_valid_or_raise(cls, model_name: str, model_type: ModelType, credentials: dict):
+        """
+        check model credentials valid.
+
+        :param model_name:
+        :param model_type:
+        :param credentials:
+        """
+        return
+
+    @classmethod
+    def encrypt_model_credentials(cls, tenant_id: str, model_name: str, model_type: ModelType,
+                                  credentials: dict) -> dict:
+        """
+        encrypt model credentials for save.
+
+        :param tenant_id:
+        :param model_name:
+        :param model_type:
+        :param credentials:
+        :return:
+        """
+        return {}
+
+    def get_model_credentials(self, model_name: str, model_type: ModelType, obfuscated: bool = False) -> dict:
+        """
+        get credentials for llm use.
+
+        :param model_name:
+        :param model_type:
+        :param obfuscated:
+        :return:
+        """
+        return self.get_provider_credentials(obfuscated)

api/core/model_providers/rules/_providers.json

@@ -13,5 +13,6 @@
   "huggingface_hub",
   "xinference",
   "openllm",
-  "localai"
+  "localai",
+  "cohere"
 ]

api/core/model_providers/rules/cohere.json

@@ -0,0 +1,7 @@
+{
+    "support_provider_types": [
+        "custom"
+    ],
+    "system_config": null,
+    "model_flexibility": "fixed"
+}