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.huggingface_tei.rerank.rerank import (8    HuggingfaceTeiRerankModel,9)10from core.model_runtime.model_providers.huggingface_tei.text_embedding.text_embedding import TeiHelper11from tests.integration_tests.model_runtime.__mock.huggingface_tei import MockTEIClass12 13MOCK = os.getenv("MOCK_SWITCH", "false").lower() == "true"14 15 16@pytest.fixture17def setup_tei_mock(request, monkeypatch: pytest.MonkeyPatch):18    if MOCK:19        monkeypatch.setattr(TeiHelper, "get_tei_extra_parameter", MockTEIClass.get_tei_extra_parameter)20        monkeypatch.setattr(TeiHelper, "invoke_tokenize", MockTEIClass.invoke_tokenize)21        monkeypatch.setattr(TeiHelper, "invoke_embeddings", MockTEIClass.invoke_embeddings)22        monkeypatch.setattr(TeiHelper, "invoke_rerank", MockTEIClass.invoke_rerank)23    yield24 25    if MOCK:26        monkeypatch.undo()27 28 29@pytest.mark.parametrize("setup_tei_mock", [["none"]], indirect=True)30def test_validate_credentials(setup_tei_mock):31    model = HuggingfaceTeiRerankModel()32    # model name is only used in mock33    model_name = "reranker"34 35    if MOCK:36        # TEI Provider will check model type by API endpoint, at real server, the model type is correct.37        # So we dont need to check model type here. Only check in mock38        with pytest.raises(CredentialsValidateFailedError):39            model.validate_credentials(40                model="embedding",41                credentials={42                    "server_url": os.environ.get("TEI_RERANK_SERVER_URL"),43                },44            )45 46    model.validate_credentials(47        model=model_name,48        credentials={49            "server_url": os.environ.get("TEI_RERANK_SERVER_URL"),50        },51    )52 53 54@pytest.mark.parametrize("setup_tei_mock", [["none"]], indirect=True)55def test_invoke_model(setup_tei_mock):56    model = HuggingfaceTeiRerankModel()57    # model name is only used in mock58    model_name = "reranker"59 60    result = model.invoke(61        model=model_name,62        credentials={63            "server_url": os.environ.get("TEI_RERANK_SERVER_URL"),64        },65        query="Who is Kasumi?",66        docs=[67            'Kasumi is a girl\'s name of Japanese origin meaning "mist".',68            "Her music is a kawaii bass, a mix of future bass, pop, and kawaii music ",69            "and she leads a team named PopiParty.",70        ],71        score_threshold=0.8,72    )73 74    assert isinstance(result, RerankResult)75    assert len(result.docs) == 176    assert result.docs[0].index == 077    assert result.docs[0].score >= 0.878