mirror of
https://github.com/langgenius/dify-plugin-sdks.git
synced 2026-07-22 18:35:29 -04:00
Yeuoly
25 lines
752 B
Python
25 lines
752 B
Python
from typing import Generator
|
|
from dify_plugin.entities.model.rerank import RerankModelConfig
|
|
from dify_plugin.entities.tool import ToolInvokeMessage
|
|
from dify_plugin.interfaces.tool import Tool
|
|
|
|
|
|
class Rerank(Tool):
|
|
def _invoke(
|
|
self, tool_parameters: dict
|
|
) -> Generator[ToolInvokeMessage, None, None]:
|
|
response = self.session.model.rerank.invoke(
|
|
model_config=RerankModelConfig(
|
|
provider="jina",
|
|
model="jina-embeddings-v2-base-v1.0",
|
|
score_threshold=0.5,
|
|
top_n=10,
|
|
),
|
|
docs=["Kasumi", "Utae", "Arisa"],
|
|
query="Utae",
|
|
)
|
|
|
|
yield self.create_json_message({
|
|
"data": response.docs,
|
|
})
|