mirror of
https://github.com/langgenius/dify-plugin-daemon.git
synced 2026-07-22 01:35:24 -04:00
Yeuoly
105 lines
3.3 KiB
Python
105 lines
3.3 KiB
Python
from typing import Optional
|
|
|
|
import httpx
|
|
|
|
from dify_plugin import RerankModel
|
|
from dify_plugin.entities import I18nObject
|
|
from dify_plugin.entities.model import AIModelEntity, FetchFrom, ModelPropertyKey, ModelType
|
|
from dify_plugin.errors.model import (
|
|
CredentialsValidateFailedError,
|
|
InvokeAuthorizationError,
|
|
InvokeBadRequestError,
|
|
InvokeConnectionError,
|
|
InvokeError,
|
|
InvokeRateLimitError,
|
|
InvokeServerUnavailableError,
|
|
)
|
|
from dify_plugin.entities.model.rerank import (
|
|
RerankDocument,
|
|
RerankResult,
|
|
)
|
|
|
|
class {{ .PluginName | SnakeToCamel }}RerankModel(RerankModel):
|
|
"""
|
|
Model class for {{ .PluginName | SnakeToCamel }} rerank model.
|
|
"""
|
|
|
|
def _invoke(
|
|
self,
|
|
model: str,
|
|
credentials: dict,
|
|
query: str,
|
|
docs: list[str],
|
|
score_threshold: Optional[float] = None,
|
|
top_n: Optional[int] = None,
|
|
user: Optional[str] = None,
|
|
) -> RerankResult:
|
|
"""
|
|
Invoke rerank model
|
|
|
|
:param model: model name
|
|
:param credentials: model credentials
|
|
:param query: search query
|
|
:param docs: docs for reranking
|
|
:param score_threshold: score threshold
|
|
:param top_n: top n documents to return
|
|
:param user: unique user id
|
|
:return: rerank result
|
|
"""
|
|
pass
|
|
|
|
def validate_credentials(self, model: str, credentials: dict) -> None:
|
|
"""
|
|
Validate model credentials
|
|
|
|
:param model: model name
|
|
:param credentials: model credentials
|
|
:return:
|
|
"""
|
|
try:
|
|
self._invoke(
|
|
model=model,
|
|
credentials=credentials,
|
|
query="What is the capital of the United States?",
|
|
docs=[
|
|
"Carson City is the capital city of the American state of Nevada. At the 2010 United States "
|
|
"Census, Carson City had a population of 55,274.",
|
|
"The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean that "
|
|
"are a political division controlled by the United States. Its capital is Saipan.",
|
|
],
|
|
score_threshold=0.8,
|
|
)
|
|
except Exception as ex:
|
|
raise CredentialsValidateFailedError(str(ex))
|
|
|
|
@property
|
|
def _invoke_error_mapping(self) -> dict[type[InvokeError], list[type[Exception]]]:
|
|
"""
|
|
Map model invoke error to unified error
|
|
"""
|
|
return {
|
|
InvokeConnectionError: [httpx.ConnectError],
|
|
InvokeServerUnavailableError: [httpx.RemoteProtocolError],
|
|
InvokeRateLimitError: [],
|
|
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={
|
|
ModelPropertyKey.CONTEXT_SIZE: int(credentials.get("context_size") or 0)
|
|
},
|
|
)
|
|
|
|
return entity
|