dify/api/tests/integration_tests/model_runtime/huggingface_tei/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.huggingface_tei.rerank.rerank import (
    HuggingfaceTeiRerankModel,
)
from core.model_runtime.model_providers.huggingface_tei.text_embedding.text_embedding import TeiHelper
from tests.integration_tests.model_runtime.__mock.huggingface_tei import MockTEIClass

MOCK = os.getenv("MOCK_SWITCH", "false").lower() == "true"


@pytest.fixture
def setup_tei_mock(request, monkeypatch: pytest.MonkeyPatch):
    if MOCK:
        monkeypatch.setattr(TeiHelper, "get_tei_extra_parameter", MockTEIClass.get_tei_extra_parameter)
        monkeypatch.setattr(TeiHelper, "invoke_tokenize", MockTEIClass.invoke_tokenize)
        monkeypatch.setattr(TeiHelper, "invoke_embeddings", MockTEIClass.invoke_embeddings)
        monkeypatch.setattr(TeiHelper, "invoke_rerank", MockTEIClass.invoke_rerank)
    yield

    if MOCK:
        monkeypatch.undo()


@pytest.mark.parametrize("setup_tei_mock", [["none"]], indirect=True)
def test_validate_credentials(setup_tei_mock):
    model = HuggingfaceTeiRerankModel()
    # model name is only used in mock
    model_name = "reranker"

    if MOCK:
        # TEI Provider will check model type by API endpoint, at real server, the model type is correct.
        # So we dont need to check model type here. Only check in mock
        with pytest.raises(CredentialsValidateFailedError):
            model.validate_credentials(
                model="embedding",
                credentials={
                    "server_url": os.environ.get("TEI_RERANK_SERVER_URL"),
                },
            )

    model.validate_credentials(
        model=model_name,
        credentials={
            "server_url": os.environ.get("TEI_RERANK_SERVER_URL"),
        },
    )


@pytest.mark.parametrize("setup_tei_mock", [["none"]], indirect=True)
def test_invoke_model(setup_tei_mock):
    model = HuggingfaceTeiRerankModel()
    # model name is only used in mock
    model_name = "reranker"

    result = model.invoke(
        model=model_name,
        credentials={
            "server_url": os.environ.get("TEI_RERANK_SERVER_URL"),
        },
        query="Who is Kasumi?",
        docs=[
            'Kasumi is a girl\'s name of Japanese origin meaning "mist".',
            "Her music is a kawaii bass, a mix of future bass, pop, and kawaii music ",
            "and she leads a team named PopiParty.",
        ],
        score_threshold=0.8,
    )

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