Underground-Digital/Workflow-Engine

1import os2 3import pytest4 5from core.model_runtime.entities.rerank_entities import RerankResult6from core.model_runtime.errors.validate import CredentialsValidateFailedError7from core.model_runtime.model_providers.siliconflow.rerank.rerank import SiliconflowRerankModel8 9 10def test_validate_credentials():11    model = SiliconflowRerankModel()12 13    with pytest.raises(CredentialsValidateFailedError):14        model.validate_credentials(15            model="BAAI/bge-reranker-v2-m3",16            credentials={"api_key": "invalid_key"},17        )18 19    model.validate_credentials(20        model="BAAI/bge-reranker-v2-m3",21        credentials={22            "api_key": os.environ.get("API_KEY"),23        },24    )25 26 27def test_invoke_model():28    model = SiliconflowRerankModel()29 30    result = model.invoke(31        model="BAAI/bge-reranker-v2-m3",32        credentials={33            "api_key": os.environ.get("API_KEY"),34        },35        query="Who is Kasumi?",36        docs=[37            'Kasumi is a girl\'s name of Japanese origin meaning "mist".',38            "Her music is a kawaii bass, a mix of future bass, pop, and kawaii music ",39            "and she leads a team named PopiParty.",40        ],41        score_threshold=0.8,42    )43 44    assert isinstance(result, RerankResult)45    assert len(result.docs) == 146    assert result.docs[0].index == 047    assert result.docs[0].score >= 0.848