Underground-Digital/Workflow-Engine

1import os
2
3import pytest
4
5from core.model_runtime.entities.rerank_entities import RerankResult
6from core.model_runtime.errors.validate import CredentialsValidateFailedError
7from core.model_runtime.model_providers.gitee_ai.rerank.rerank import GiteeAIRerankModel
8
9
10def test_validate_credentials():
11    model = GiteeAIRerankModel()
12
13    with pytest.raises(CredentialsValidateFailedError):
14        model.validate_credentials(
15            model="bge-reranker-v2-m3",
16            credentials={"api_key": "invalid_key"},
17        )
18
19    model.validate_credentials(
20        model="bge-reranker-v2-m3",
21        credentials={
22            "api_key": os.environ.get("GITEE_AI_API_KEY"),
23        },
24    )
25
26
27def test_invoke_model():
28    model = GiteeAIRerankModel()
29    result = model.invoke(
30        model="bge-reranker-v2-m3",
31        credentials={
32            "api_key": os.environ.get("GITEE_AI_API_KEY"),
33        },
34        query="What is the capital of the United States?",
35        docs=[
36            "Carson City is the capital city of the American state of Nevada. At the 2010 United States "
37            "Census, Carson City had a population of 55,274.",
38            "The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean that "
39            "are a political division controlled by the United States. Its capital is Saipan.",
40        ],
41        top_n=1,
42        score_threshold=0.01,
43    )
44
45    assert isinstance(result, RerankResult)
46    assert len(result.docs) == 1
47    assert result.docs[0].score >= 0.01
48