diff --git a/api/tests/integration_tests/model_runtime/localai/test_rerank.py b/api/tests/integration_tests/model_runtime/localai/test_rerank.py index a75439337e..99847bc852 100644 --- a/api/tests/integration_tests/model_runtime/localai/test_rerank.py +++ b/api/tests/integration_tests/model_runtime/localai/test_rerank.py @@ -1,64 +1,8 @@ import os import pytest -from api.core.model_runtime.entities.rerank_entities import RerankResult - -from core.model_runtime.errors.validate import CredentialsValidateFailedError -from core.model_runtime.model_providers.localai.rerank.rerank import LocalaiRerankModel - - -def test_validate_credentials_for_chat_model(): - model = LocalaiRerankModel() - - with pytest.raises(CredentialsValidateFailedError): - model.validate_credentials( - model='bge-reranker-v2-m3', - credentials={ - 'server_url': 'hahahaha', - 'completion_type': 'completion', - } - ) - - model.validate_credentials( - model='bge-reranker-base', - credentials={ - 'server_url': os.environ.get('LOCALAI_SERVER_URL'), - 'completion_type': 'completion', - } - ) - -def test_invoke_rerank_model(): - model = LocalaiRerankModel() - - response = model.invoke( - model='bge-reranker-base', - credentials={ - 'server_url': os.environ.get('LOCALAI_SERVER_URL') - }, - query='Organic skincare products for sensitive skin', - docs=[ - "Eco-friendly kitchenware for modern homes", - "Biodegradable cleaning supplies for eco-conscious consumers", - "Organic cotton baby clothes for sensitive skin", - "Natural organic skincare range for sensitive skin", - "Tech gadgets for smart homes: 2024 edition", - "Sustainable gardening tools and compost solutions", - "Sensitive skin-friendly facial cleansers and toners", - "Organic food wraps and storage solutions", - "Yoga mats made from recycled materials" - ], - top_n=3, - score_threshold=0.75, - user="abc-123" - ) - - assert isinstance(response, RerankResult) - assert len(response.docs) == 3 -import os - -import pytest -from api.core.model_runtime.entities.rerank_entities import RerankDocument, RerankResult +from core.model_runtime.entities.rerank_entities import RerankDocument, RerankResult from core.model_runtime.errors.validate import CredentialsValidateFailedError from core.model_runtime.model_providers.localai.rerank.rerank import LocalaiRerankModel