|
@@ -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
|
|
|
|