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.sagemaker.rerank.rerank import SageMakerRerankModel8 9 10def test_validate_credentials():11    model = SageMakerRerankModel()12 13    with pytest.raises(CredentialsValidateFailedError):14        model.validate_credentials(15            model="bge-m3-rerank-v2",16            credentials={17                "aws_region": os.getenv("AWS_REGION"),18                "aws_access_key": os.getenv("AWS_ACCESS_KEY"),19                "aws_secret_access_key": os.getenv("AWS_SECRET_ACCESS_KEY"),20            },21            query="What is the capital of the United States?",22            docs=[23                "Carson City is the capital city of the American state of Nevada. At the 2010 United States "24                "Census, Carson City had a population of 55,274.",25                "The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean that "26                "are a political division controlled by the United States. Its capital is Saipan.",27            ],28            score_threshold=0.8,29        )30 31 32def test_invoke_model():33    model = SageMakerRerankModel()34 35    result = model.invoke(36        model="bge-m3-rerank-v2",37        credentials={38            "aws_region": os.getenv("AWS_REGION"),39            "aws_access_key": os.getenv("AWS_ACCESS_KEY"),40            "aws_secret_access_key": os.getenv("AWS_SECRET_ACCESS_KEY"),41        },42        query="What is the capital of the United States?",43        docs=[44            "Carson City is the capital city of the American state of Nevada. At the 2010 United States "45            "Census, Carson City had a population of 55,274.",46            "The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific Ocean that "47            "are a political division controlled by the United States. Its capital is Saipan.",48        ],49        score_threshold=0.8,50    )51 52    assert isinstance(result, RerankResult)53    assert len(result.docs) == 154    assert result.docs[0].index == 155    assert result.docs[0].score >= 0.856