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

This commit is contained in:
Giovanny Gutiérrez 2024-05-26 23:09:51 -05:00 committed by GitHub
parent 8152bc6fbf
commit 2deb23e00e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -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