Underground-Digital/Workflow-Engine

1from typing import Optional2 3from core.model_manager import ModelInstance4from core.rag.models.document import Document5from core.rag.rerank.rerank_base import BaseRerankRunner6 7 8class RerankModelRunner(BaseRerankRunner):9    def __init__(self, rerank_model_instance: ModelInstance) -> None:10        self.rerank_model_instance = rerank_model_instance11 12    def run(13        self,14        query: str,15        documents: list[Document],16        score_threshold: Optional[float] = None,17        top_n: Optional[int] = None,18        user: Optional[str] = None,19    ) -> list[Document]:20        """21        Run rerank model22        :param query: search query23        :param documents: documents for reranking24        :param score_threshold: score threshold25        :param top_n: top n26        :param user: unique user id if needed27        :return:28        """29        docs = []30        doc_id = set()31        unique_documents = []32        for document in documents:33            if document.provider == "dify" and document.metadata["doc_id"] not in doc_id:34                doc_id.add(document.metadata["doc_id"])35                docs.append(document.page_content)36                unique_documents.append(document)37            elif document.provider == "external":38                if document not in unique_documents:39                    docs.append(document.page_content)40                    unique_documents.append(document)41 42        documents = unique_documents43 44        rerank_result = self.rerank_model_instance.invoke_rerank(45            query=query, docs=docs, score_threshold=score_threshold, top_n=top_n, user=user46        )47 48        rerank_documents = []49 50        for result in rerank_result.docs:51            # format document52            rerank_document = Document(53                page_content=result.text,54                metadata=documents[result.index].metadata,55                provider=documents[result.index].provider,56            )57            rerank_document.metadata["score"] = result.score58            rerank_documents.append(rerank_document)59 60        return rerank_documents61