Feat: Add model provider Text Embedding Inference for embedding and rerank (#7132)

api/tests/integration_tests/model_runtime/__mock/huggingface_tei.py

@@ -0,0 +1,94 @@
+
+from api.core.model_runtime.model_providers.huggingface_tei.tei_helper import TeiModelExtraParameter
+
+
+class MockTEIClass:
+    @staticmethod
+    def get_tei_extra_parameter(server_url: str, model_name: str) -> TeiModelExtraParameter:
+        # During mock, we don't have a real server to query, so we just return a dummy value
+        if 'rerank' in model_name:
+            model_type = 'reranker'
+        else:
+            model_type = 'embedding'
+
+        return TeiModelExtraParameter(model_type=model_type, max_input_length=512, max_client_batch_size=1)
+    
+    @staticmethod
+    def invoke_tokenize(server_url: str, texts: list[str]) -> list[list[dict]]:
+        # Use space as token separator, and split the text into tokens
+        tokenized_texts = []
+        for text in texts:
+            tokens = text.split(' ')
+            current_index = 0
+            tokenized_text = []
+            for idx, token in enumerate(tokens):
+                s_token = {
+                    'id': idx,
+                    'text': token,
+                    'special': False,
+                    'start': current_index,
+                    'stop': current_index + len(token),
+                }
+                current_index += len(token) + 1
+                tokenized_text.append(s_token)
+            tokenized_texts.append(tokenized_text)
+        return tokenized_texts
+
+    @staticmethod
+    def invoke_embeddings(server_url: str, texts: list[str]) -> dict:
+        # {
+        #     "object": "list",
+        #     "data": [
+        #         {
+        #             "object": "embedding",
+        #             "embedding": [...],
+        #             "index": 0
+        #         }
+        #     ],
+        #     "model": "MODEL_NAME",
+        #     "usage": {
+        #         "prompt_tokens": 3,
+        #         "total_tokens": 3
+        #     }
+        # }
+        embeddings = []
+        for idx, text in enumerate(texts):
+            embedding = [0.1] * 768
+            embeddings.append(
+                {
+                    'object': 'embedding',
+                    'embedding': embedding,
+                    'index': idx,
+                }
+            )
+        return {
+            'object': 'list',
+            'data': embeddings,
+            'model': 'MODEL_NAME',
+            'usage': {
+                'prompt_tokens': sum(len(text.split(' ')) for text in texts),
+                'total_tokens': sum(len(text.split(' ')) for text in texts),
+            },
+        }
+
+    def invoke_rerank(server_url: str, query: str, texts: list[str]) -> list[dict]:
+        #         Example response:
+        # [
+        #     {
+        #         "index": 0,
+        #         "text": "Deep Learning is ...",
+        #         "score": 0.9950755
+        #     }
+        # ]
+        reranked_docs = []
+        for idx, text in enumerate(texts):
+            reranked_docs.append(
+                {
+                    'index': idx,
+                    'text': text,
+                    'score': 0.9,
+                }
+            )
+            # For mock, only return the first document
+            break
+        return reranked_docs

api/tests/integration_tests/model_runtime/huggingface_tei/test_embeddings.py

@@ -0,0 +1,72 @@
+import os
+
+import pytest
+from api.core.model_runtime.model_providers.huggingface_tei.text_embedding.text_embedding import TeiHelper
+
+from core.model_runtime.entities.text_embedding_entities import TextEmbeddingResult
+from core.model_runtime.errors.validate import CredentialsValidateFailedError
+from core.model_runtime.model_providers.huggingface_tei.text_embedding.text_embedding import (
+    HuggingfaceTeiTextEmbeddingModel,
+)
+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 = HuggingfaceTeiTextEmbeddingModel()
+    # model name is only used in mock
+    model_name = 'embedding'
+
+    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='reranker',
+                credentials={
+                    'server_url': os.environ.get('TEI_EMBEDDING_SERVER_URL', ""),
+                }
+            )
+
+    model.validate_credentials(
+        model=model_name,
+        credentials={
+            'server_url': os.environ.get('TEI_EMBEDDING_SERVER_URL', ""),
+        }
+    )
+
+@pytest.mark.parametrize('setup_tei_mock', [['none']], indirect=True)
+def test_invoke_model(setup_tei_mock):
+    model = HuggingfaceTeiTextEmbeddingModel()
+    model_name = 'embedding'
+
+    result = model.invoke(
+        model=model_name,
+        credentials={
+            'server_url': os.environ.get('TEI_EMBEDDING_SERVER_URL', ""),
+        },
+        texts=[
+            "hello",
+            "world"
+        ],
+        user="abc-123"
+    )
+
+    assert isinstance(result, TextEmbeddingResult)
+    assert len(result.embeddings) == 2
+    assert result.usage.total_tokens > 0

api/tests/integration_tests/model_runtime/huggingface_tei/test_rerank.py

@@ -0,0 +1,76 @@
+import os
+
+import pytest
+
+from core.model_runtime.entities.rerank_entities import RerankDocument, RerankResult
+from core.model_runtime.entities.text_embedding_entities import TextEmbeddingResult
+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