Underground-Digital/Workflow-Engine

1import os2from unittest.mock import Mock, patch3 4import pytest5 6from core.model_runtime.entities.rerank_entities import RerankResult7from core.model_runtime.errors.validate import CredentialsValidateFailedError8from core.model_runtime.model_providers.mixedbread.rerank.rerank import MixedBreadRerankModel9 10 11def test_validate_credentials():12    model = MixedBreadRerankModel()13 14    with pytest.raises(CredentialsValidateFailedError):15        model.validate_credentials(16            model="mxbai-rerank-large-v1",17            credentials={"api_key": "invalid_key"},18        )19    with patch("httpx.post") as mock_post:20        mock_response = Mock()21        mock_response.json.return_value = {22            "usage": {"prompt_tokens": 86, "total_tokens": 86},23            "model": "mixedbread-ai/mxbai-rerank-large-v1",24            "data": [25                {26                    "index": 0,27                    "score": 0.06762695,28                    "input": "Carson City is the capital city of the American state of Nevada. At the 2010 United "29                    "States Census, Carson City had a population of 55,274.",30                    "object": "text_document",31                },32                {33                    "index": 1,34                    "score": 0.057403564,35                    "input": "The Commonwealth of the Northern Mariana Islands is a group of islands in the Pacific "36                    "Ocean that are a political division controlled by the United States. Its capital is "37                    "Saipan.",38                    "object": "text_document",39                },40            ],41            "object": "list",42            "top_k": 2,43            "return_input": True,44        }45        mock_response.status_code = 20046        mock_post.return_value = mock_response47        model.validate_credentials(48            model="mxbai-rerank-large-v1",49            credentials={50                "api_key": os.environ.get("MIXEDBREAD_API_KEY"),51            },52        )53 54 55def test_invoke_model():56    model = MixedBreadRerankModel()57    with patch("httpx.post") as mock_post:58        mock_response = Mock()59        mock_response.json.return_value = {60            "usage": {"prompt_tokens": 56, "total_tokens": 56},61            "model": "mixedbread-ai/mxbai-rerank-large-v1",62            "data": [63                {64                    "index": 0,65                    "score": 0.6044922,66                    "input": "Kasumi is a girl name of Japanese origin meaning mist.",67                    "object": "text_document",68                },69                {70                    "index": 1,71                    "score": 0.0703125,72                    "input": "Her music is a kawaii bass, a mix of future bass, pop, and kawaii music and she leads a "73                    "team named PopiParty.",74                    "object": "text_document",75                },76            ],77            "object": "list",78            "top_k": 2,79            "return_input": "true",80        }81        mock_response.status_code = 20082        mock_post.return_value = mock_response83        result = model.invoke(84            model="mxbai-rerank-large-v1",85            credentials={86                "api_key": os.environ.get("MIXEDBREAD_API_KEY"),87            },88            query="Who is Kasumi?",89            docs=[90                "Kasumi is a girl name of Japanese origin meaning mist.",91                "Her music is a kawaii bass, a mix of future bass, pop, and kawaii music and she leads a team named "92                "PopiParty.",93            ],94            score_threshold=0.5,95        )96 97        assert isinstance(result, RerankResult)98        assert len(result.docs) == 199        assert result.docs[0].index == 0100        assert result.docs[0].score >= 0.5101