Underground-Digital/Workflow-Engine

1from typing import Optional2 3import httpx4 5from core.model_runtime.entities.common_entities import I18nObject6from core.model_runtime.entities.model_entities import AIModelEntity, FetchFrom, ModelPropertyKey, ModelType7from core.model_runtime.entities.rerank_entities import RerankDocument, RerankResult8from core.model_runtime.errors.invoke import (9    InvokeAuthorizationError,10    InvokeBadRequestError,11    InvokeConnectionError,12    InvokeError,13    InvokeRateLimitError,14    InvokeServerUnavailableError,15)16from core.model_runtime.errors.validate import CredentialsValidateFailedError17from core.model_runtime.model_providers.__base.rerank_model import RerankModel18 19 20class MixedBreadRerankModel(RerankModel):21    """22    Model class for MixedBread rerank model.23    """24 25    def _invoke(26        self,27        model: str,28        credentials: dict,29        query: str,30        docs: list[str],31        score_threshold: Optional[float] = None,32        top_n: Optional[int] = None,33        user: Optional[str] = None,34    ) -> RerankResult:35        """36        Invoke rerank model37 38        :param model: model name39        :param credentials: model credentials40        :param query: search query41        :param docs: docs for reranking42        :param score_threshold: score threshold43        :param top_n: top n documents to return44        :param user: unique user id45        :return: rerank result46        """47        if len(docs) == 0:48            return RerankResult(model=model, docs=[])49 50        base_url = credentials.get("base_url", "https://api.mixedbread.ai/v1")51        base_url = base_url.removesuffix("/")52 53        try:54            response = httpx.post(55                base_url + "/reranking",56                json={"model": model, "query": query, "input": docs, "top_k": top_n, "return_input": True},57                headers={"Authorization": f"Bearer {credentials.get('api_key')}", "Content-Type": "application/json"},58            )59            response.raise_for_status()60            results = response.json()61 62            rerank_documents = []63            for result in results["data"]:64                rerank_document = RerankDocument(65                    index=result["index"],66                    text=result["input"],67                    score=result["score"],68                )69                if score_threshold is None or result["score"] >= score_threshold:70                    rerank_documents.append(rerank_document)71 72            return RerankResult(model=model, docs=rerank_documents)73        except httpx.HTTPStatusError as e:74            raise InvokeServerUnavailableError(str(e))75 76    def validate_credentials(self, model: str, credentials: dict) -> None:77        """78        Validate model credentials79 80        :param model: model name81        :param credentials: model credentials82        :return:83        """84        try:85            self._invoke(86                model=model,87                credentials=credentials,88                query="What is the capital of the United States?",89                docs=[90                    "Carson City is the capital city of the American state of Nevada. At the 2010 United States "91                    "Census, Carson City had a population of 55,274.",92                    "The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean that "93                    "are a political division controlled by the United States. Its capital is Saipan.",94                ],95                score_threshold=0.8,96            )97        except Exception as ex:98            raise CredentialsValidateFailedError(str(ex))99 100    @property101    def _invoke_error_mapping(self) -> dict[type[InvokeError], list[type[Exception]]]:102        """103        Map model invoke error to unified error104        """105        return {106            InvokeConnectionError: [httpx.ConnectError],107            InvokeServerUnavailableError: [httpx.RemoteProtocolError],108            InvokeRateLimitError: [],109            InvokeAuthorizationError: [httpx.HTTPStatusError],110            InvokeBadRequestError: [httpx.RequestError],111        }112 113    def get_customizable_model_schema(self, model: str, credentials: dict) -> AIModelEntity:114        """115        generate custom model entities from credentials116        """117        entity = AIModelEntity(118            model=model,119            label=I18nObject(en_US=model),120            model_type=ModelType.RERANK,121            fetch_from=FetchFrom.CUSTOMIZABLE_MODEL,122            model_properties={ModelPropertyKey.CONTEXT_SIZE: int(credentials.get("context_size", "512"))},123        )124 125        return entity126