Underground-Digital/Workflow-Engine

1import os2 3import pytest4 5from core.model_runtime.entities.text_embedding_entities import TextEmbeddingResult6from core.model_runtime.errors.validate import CredentialsValidateFailedError7from core.model_runtime.model_providers.huggingface_tei.text_embedding.text_embedding import (8    HuggingfaceTeiTextEmbeddingModel,9    TeiHelper,10)11from 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 = HuggingfaceTeiTextEmbeddingModel()32    # model name is only used in mock33    model_name = "embedding"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="reranker",41                credentials={42                    "server_url": os.environ.get("TEI_EMBEDDING_SERVER_URL", ""),43                },44            )45 46    model.validate_credentials(47        model=model_name,48        credentials={49            "server_url": os.environ.get("TEI_EMBEDDING_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 = HuggingfaceTeiTextEmbeddingModel()57    model_name = "embedding"58 59    result = model.invoke(60        model=model_name,61        credentials={62            "server_url": os.environ.get("TEI_EMBEDDING_SERVER_URL", ""),63        },64        texts=["hello", "world"],65        user="abc-123",66    )67 68    assert isinstance(result, TextEmbeddingResult)69    assert len(result.embeddings) == 270    assert result.usage.total_tokens > 071