Forráskód Böngészése

fix: Show rerank in system for localai (#4652)

Giovanny Gutiérrez 11 hónapja
szülő
commit
2deb23e00e

+ 16 - 0
api/core/model_runtime/model_providers/localai/rerank/rerank.py

@@ -5,6 +5,8 @@ import httpx
 from requests import post
 from yarl import URL
 
+from core.model_runtime.entities.common_entities import I18nObject
+from core.model_runtime.entities.model_entities import AIModelEntity, FetchFrom, ModelType
 from core.model_runtime.entities.rerank_entities import RerankDocument, RerankResult
 from core.model_runtime.errors.invoke import (
     InvokeAuthorizationError,
@@ -118,3 +120,17 @@ class LocalaiRerankModel(RerankModel):
             InvokeAuthorizationError: [httpx.HTTPStatusError],  
             InvokeBadRequestError: [httpx.RequestError]
         }
+    
+    def get_customizable_model_schema(self, model: str, credentials: dict) -> AIModelEntity:
+        """
+            generate custom model entities from credentials
+        """
+        entity = AIModelEntity(
+            model=model,
+            label=I18nObject(en_US=model),
+            model_type=ModelType.RERANK,
+            fetch_from=FetchFrom.CUSTOMIZABLE_MODEL,
+            model_properties={}
+        )
+
+        return entity