chzphoenix 211f416806
feat:add wenxin rerank (#9431)
Co-authored-by: cuihz <cuihz@knowbox.cn>
Co-authored-by: crazywoola <427733928@qq.com>
2024-10-17 19:18:32 +08:00

22 lines
671 B
Python

import os
from time import sleep
from core.model_runtime.entities.rerank_entities import RerankResult
from core.model_runtime.model_providers.wenxin.rerank.rerank import WenxinRerankModel
def test_invoke_bce_reranker_base_v1():
sleep(3)
model = WenxinRerankModel()
response = model.invoke(
model="bce-reranker-base_v1",
credentials={"api_key": os.environ.get("WENXIN_API_KEY"), "secret_key": os.environ.get("WENXIN_SECRET_KEY")},
query="What is Deep Learning?",
docs=["Deep Learning is ...", "My Book is ..."],
user="abc-123",
)
assert isinstance(response, RerankResult)
assert len(response.docs) == 2