Underground-Digital/Workflow-Engine

1import os2from unittest.mock import Mock, patch3 4import pytest5 6from core.model_runtime.entities.text_embedding_entities import TextEmbeddingResult7from core.model_runtime.errors.validate import CredentialsValidateFailedError8from core.model_runtime.model_providers.mixedbread.text_embedding.text_embedding import MixedBreadTextEmbeddingModel9 10 11def test_validate_credentials():12    model = MixedBreadTextEmbeddingModel()13 14    with pytest.raises(CredentialsValidateFailedError):15        model.validate_credentials(model="mxbai-embed-large-v1", credentials={"api_key": "invalid_key"})16    with patch("requests.post") as mock_post:17        mock_response = Mock()18        mock_response.json.return_value = {19            "usage": {"prompt_tokens": 3, "total_tokens": 3},20            "model": "mixedbread-ai/mxbai-embed-large-v1",21            "data": [{"embedding": [0.23333 for _ in range(1024)], "index": 0, "object": "embedding"}],22            "object": "list",23            "normalized": "true",24            "encoding_format": "float",25            "dimensions": 1024,26        }27        mock_response.status_code = 20028        mock_post.return_value = mock_response29        model.validate_credentials(30            model="mxbai-embed-large-v1", credentials={"api_key": os.environ.get("MIXEDBREAD_API_KEY")}31        )32 33 34def test_invoke_model():35    model = MixedBreadTextEmbeddingModel()36 37    with patch("requests.post") as mock_post:38        mock_response = Mock()39        mock_response.json.return_value = {40            "usage": {"prompt_tokens": 6, "total_tokens": 6},41            "model": "mixedbread-ai/mxbai-embed-large-v1",42            "data": [43                {"embedding": [0.23333 for _ in range(1024)], "index": 0, "object": "embedding"},44                {"embedding": [0.23333 for _ in range(1024)], "index": 1, "object": "embedding"},45            ],46            "object": "list",47            "normalized": "true",48            "encoding_format": "float",49            "dimensions": 1024,50        }51        mock_response.status_code = 20052        mock_post.return_value = mock_response53        result = model.invoke(54            model="mxbai-embed-large-v1",55            credentials={56                "api_key": os.environ.get("MIXEDBREAD_API_KEY"),57            },58            texts=["hello", "world"],59            user="abc-123",60        )61 62        assert isinstance(result, TextEmbeddingResult)63        assert len(result.embeddings) == 264        assert result.usage.total_tokens == 665 66 67def test_get_num_tokens():68    model = MixedBreadTextEmbeddingModel()69 70    num_tokens = model.get_num_tokens(71        model="mxbai-embed-large-v1",72        credentials={73            "api_key": os.environ.get("MIXEDBREAD_API_KEY"),74        },75        texts=["ping"],76    )77 78    assert num_tokens == 179