Underground-Digital/Workflow-Engine

1from typing import Optional2 3from core.model_manager import ModelInstance, ModelManager4from core.model_runtime.entities.model_entities import ModelType5from core.model_runtime.errors.invoke import InvokeAuthorizationError6from core.rag.data_post_processor.reorder import ReorderRunner7from core.rag.models.document import Document8from core.rag.rerank.entity.weight import KeywordSetting, VectorSetting, Weights9from core.rag.rerank.rerank_base import BaseRerankRunner10from core.rag.rerank.rerank_factory import RerankRunnerFactory11from core.rag.rerank.rerank_type import RerankMode12 13 14class DataPostProcessor:15    """Interface for data post-processing document."""16 17    def __init__(18        self,19        tenant_id: str,20        reranking_mode: str,21        reranking_model: Optional[dict] = None,22        weights: Optional[dict] = None,23        reorder_enabled: bool = False,24    ):25        self.rerank_runner = self._get_rerank_runner(reranking_mode, tenant_id, reranking_model, weights)26        self.reorder_runner = self._get_reorder_runner(reorder_enabled)27 28    def invoke(29        self,30        query: str,31        documents: list[Document],32        score_threshold: Optional[float] = None,33        top_n: Optional[int] = None,34        user: Optional[str] = None,35    ) -> list[Document]:36        if self.rerank_runner:37            documents = self.rerank_runner.run(query, documents, score_threshold, top_n, user)38 39        if self.reorder_runner:40            documents = self.reorder_runner.run(documents)41 42        return documents43 44    def _get_rerank_runner(45        self,46        reranking_mode: str,47        tenant_id: str,48        reranking_model: Optional[dict] = None,49        weights: Optional[dict] = None,50    ) -> Optional[BaseRerankRunner]:51        if reranking_mode == RerankMode.WEIGHTED_SCORE.value and weights:52            runner = RerankRunnerFactory.create_rerank_runner(53                runner_type=reranking_mode,54                tenant_id=tenant_id,55                weights=Weights(56                    vector_setting=VectorSetting(57                        vector_weight=weights["vector_setting"]["vector_weight"],58                        embedding_provider_name=weights["vector_setting"]["embedding_provider_name"],59                        embedding_model_name=weights["vector_setting"]["embedding_model_name"],60                    ),61                    keyword_setting=KeywordSetting(62                        keyword_weight=weights["keyword_setting"]["keyword_weight"],63                    ),64                ),65            )66            return runner67        elif reranking_mode == RerankMode.RERANKING_MODEL.value:68            rerank_model_instance = self._get_rerank_model_instance(tenant_id, reranking_model)69            if rerank_model_instance is None:70                return None71            runner = RerankRunnerFactory.create_rerank_runner(72                runner_type=reranking_mode, rerank_model_instance=rerank_model_instance73            )74            return runner75        return None76 77    def _get_reorder_runner(self, reorder_enabled) -> Optional[ReorderRunner]:78        if reorder_enabled:79            return ReorderRunner()80        return None81 82    def _get_rerank_model_instance(self, tenant_id: str, reranking_model: Optional[dict]) -> ModelInstance | None:83        if reranking_model:84            try:85                model_manager = ModelManager()86                rerank_model_instance = model_manager.get_model_instance(87                    tenant_id=tenant_id,88                    provider=reranking_model["reranking_provider_name"],89                    model_type=ModelType.RERANK,90                    model=reranking_model["reranking_model_name"],91                )92                return rerank_model_instance93            except InvokeAuthorizationError:94                return None95        return None96