tjdify/api/tests/integration_tests/model_runtime/cohere/test_rerank.py

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.cohere.rerank.rerank import CohereRerankModel


def test_validate_credentials():
    model = CohereRerankModel()

    with pytest.raises(CredentialsValidateFailedError):
        model.validate_credentials(model="rerank-english-v2.0", credentials={"api_key": "invalid_key"})

    model.validate_credentials(model="rerank-english-v2.0", credentials={"api_key": os.environ.get("COHERE_API_KEY")})


def test_invoke_model():
    model = CohereRerankModel()

    result = model.invoke(
        model="rerank-english-v2.0",
        credentials={"api_key": os.environ.get("COHERE_API_KEY")},
        query="What is the capital of the United States?",
        docs=[
            "Carson City is the capital city of the American state of Nevada. At the 2010 United States "
            "Census, Carson City had a population of 55,274.",
            "Washington, D.C. (also known as simply Washington or D.C., and officially as the District of Columbia) "
            "is the capital of the United States. It is a federal district. The President of the USA and many major "
            "national government offices are in the territory. This makes it the political center of the United "
            "States of America.",
        ],
        score_threshold=0.8,
    )

    assert isinstance(result, RerankResult)
    assert len(result.docs) == 1
    assert result.docs[0].index == 1
    assert result.docs[0].score >= 0.8