Underground-Digital/Workflow-Engine

1import os2 3import pytest4 5from core.model_runtime.entities.rerank_entities import RerankDocument, RerankResult6from core.model_runtime.errors.validate import CredentialsValidateFailedError7from core.model_runtime.model_providers.gpustack.rerank.rerank import (8    GPUStackRerankModel,9)10 11 12def test_validate_credentials_for_rerank_model():13    model = GPUStackRerankModel()14 15    with pytest.raises(CredentialsValidateFailedError):16        model.validate_credentials(17            model="bge-reranker-v2-m3",18            credentials={19                "endpoint_url": "invalid_url",20                "api_key": "invalid_api_key",21            },22        )23 24    model.validate_credentials(25        model="bge-reranker-v2-m3",26        credentials={27            "endpoint_url": os.environ.get("GPUSTACK_SERVER_URL"),28            "api_key": os.environ.get("GPUSTACK_API_KEY"),29        },30    )31 32 33def test_invoke_rerank_model():34    model = GPUStackRerankModel()35 36    response = model.invoke(37        model="bge-reranker-v2-m3",38        credentials={39            "endpoint_url": os.environ.get("GPUSTACK_SERVER_URL"),40            "api_key": os.environ.get("GPUSTACK_API_KEY"),41        },42        query="Organic skincare products for sensitive skin",43        docs=[44            "Eco-friendly kitchenware for modern homes",45            "Biodegradable cleaning supplies for eco-conscious consumers",46            "Organic cotton baby clothes for sensitive skin",47            "Natural organic skincare range for sensitive skin",48            "Tech gadgets for smart homes: 2024 edition",49            "Sustainable gardening tools and compost solutions",50            "Sensitive skin-friendly facial cleansers and toners",51            "Organic food wraps and storage solutions",52            "Yoga mats made from recycled materials",53        ],54        top_n=3,55        score_threshold=-0.75,56        user="abc-123",57    )58 59    assert isinstance(response, RerankResult)60    assert len(response.docs) == 361 62 63def test__invoke():64    model = GPUStackRerankModel()65 66    # Test case 1: Empty docs67    result = model._invoke(68        model="bge-reranker-v2-m3",69        credentials={70            "endpoint_url": os.environ.get("GPUSTACK_SERVER_URL"),71            "api_key": os.environ.get("GPUSTACK_API_KEY"),72        },73        query="Organic skincare products for sensitive skin",74        docs=[],75        top_n=3,76        score_threshold=0.75,77        user="abc-123",78    )79    assert isinstance(result, RerankResult)80    assert len(result.docs) == 081 82    # Test case 2: Expected docs83    result = model._invoke(84        model="bge-reranker-v2-m3",85        credentials={86            "endpoint_url": os.environ.get("GPUSTACK_SERVER_URL"),87            "api_key": os.environ.get("GPUSTACK_API_KEY"),88        },89        query="Organic skincare products for sensitive skin",90        docs=[91            "Eco-friendly kitchenware for modern homes",92            "Biodegradable cleaning supplies for eco-conscious consumers",93            "Organic cotton baby clothes for sensitive skin",94            "Natural organic skincare range for sensitive skin",95            "Tech gadgets for smart homes: 2024 edition",96            "Sustainable gardening tools and compost solutions",97            "Sensitive skin-friendly facial cleansers and toners",98            "Organic food wraps and storage solutions",99            "Yoga mats made from recycled materials",100        ],101        top_n=3,102        score_threshold=-0.75,103        user="abc-123",104    )105    assert isinstance(result, RerankResult)106    assert len(result.docs) == 3107    assert all(isinstance(doc, RerankDocument) for doc in result.docs)108