2024-08-16 20:21:41 +08:00
|
|
|
import os
|
|
|
|
|
|
|
|
import pytest
|
|
|
|
|
|
|
|
from core.model_runtime.entities.rerank_entities import RerankResult
|
|
|
|
from core.model_runtime.errors.validate import CredentialsValidateFailedError
|
|
|
|
from core.model_runtime.model_providers.siliconflow.rerank.rerank import SiliconflowRerankModel
|
|
|
|
|
|
|
|
|
|
|
|
def test_validate_credentials():
|
|
|
|
model = SiliconflowRerankModel()
|
|
|
|
|
|
|
|
with pytest.raises(CredentialsValidateFailedError):
|
|
|
|
model.validate_credentials(
|
|
|
|
model="BAAI/bge-reranker-v2-m3",
|
2024-08-23 23:52:25 +08:00
|
|
|
credentials={"api_key": "invalid_key"},
|
2024-08-16 20:21:41 +08:00
|
|
|
)
|
|
|
|
|
|
|
|
model.validate_credentials(
|
|
|
|
model="BAAI/bge-reranker-v2-m3",
|
|
|
|
credentials={
|
|
|
|
"api_key": os.environ.get("API_KEY"),
|
|
|
|
},
|
|
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
def test_invoke_model():
|
|
|
|
model = SiliconflowRerankModel()
|
|
|
|
|
|
|
|
result = model.invoke(
|
2024-08-23 23:52:25 +08:00
|
|
|
model="BAAI/bge-reranker-v2-m3",
|
2024-08-16 20:21:41 +08:00
|
|
|
credentials={
|
|
|
|
"api_key": os.environ.get("API_KEY"),
|
|
|
|
},
|
|
|
|
query="Who is Kasumi?",
|
|
|
|
docs=[
|
2024-08-23 23:52:25 +08:00
|
|
|
'Kasumi is a girl\'s name of Japanese origin meaning "mist".',
|
2024-08-16 20:21:41 +08:00
|
|
|
"Her music is a kawaii bass, a mix of future bass, pop, and kawaii music ",
|
2024-08-23 23:52:25 +08:00
|
|
|
"and she leads a team named PopiParty.",
|
2024-08-16 20:21:41 +08:00
|
|
|
],
|
2024-08-23 23:52:25 +08:00
|
|
|
score_threshold=0.8,
|
2024-08-16 20:21:41 +08:00
|
|
|
)
|
|
|
|
|
|
|
|
assert isinstance(result, RerankResult)
|
|
|
|
assert len(result.docs) == 1
|
|
|
|
assert result.docs[0].index == 0
|
|
|
|
assert result.docs[0].score >= 0.8
|