mirror of
https://github.com/langgenius/dify.git
synced 2026-10-04 13:47:55 +08:00
test: use real rerank model instances (#43221)
This commit is contained in:
Asuka Minato authored and 
GitHub committed
1 parent
0f611811c2
commit
b7bf3f8a29
1 file changed
+45
-30
@@ -19,7 +19,9 @@ from unittest.mock import MagicMock, Mock, patch
|
||||
import pytest
|
||||
from sqlalchemy.orm import Session
|
||||
|
||||
from core.entities.provider_configuration import ProviderModelBundle
|
||||
from core.model_manager import ModelInstance
|
||||
from core.plugin.impl.model_runtime_factory import create_plugin_model_runtime
|
||||
from core.rag.index_processor.constant.doc_type import DocType
|
||||
from core.rag.index_processor.constant.query_type import QueryType
|
||||
from core.rag.models.document import Document
|
||||
@@ -30,23 +32,35 @@ from core.rag.rerank.rerank_model import RerankModelRunner
|
||||
from core.rag.rerank.rerank_type import RerankMode
|
||||
from core.rag.rerank.weight_rerank import WeightRerankRunner
|
||||
from extensions.storage.storage_type import StorageType
|
||||
from graphon.model_runtime.entities.model_entities import ModelFeature
|
||||
from graphon.model_runtime.entities.model_entities import ModelFeature, ModelType
|
||||
from graphon.model_runtime.entities.rerank_entities import RerankDocument, RerankResult
|
||||
from graphon.model_runtime.model_providers.base.rerank_model import RerankModel
|
||||
from models.enums import CreatorUserRole
|
||||
from models.model import UploadFile
|
||||
from tests.unit_tests.core.model_fixtures import make_model_config
|
||||
|
||||
|
||||
def create_mock_model_instance() -> ModelInstance:
|
||||
"""Create a properly configured mock ModelInstance for reranking tests."""
|
||||
mock_instance = Mock(spec=ModelInstance)
|
||||
# Setup provider_model_bundle chain for check_model_support_vision
|
||||
mock_instance.provider_model_bundle = Mock()
|
||||
mock_instance.provider_model_bundle.configuration = Mock()
|
||||
mock_instance.provider_model_bundle.configuration.tenant_id = "test-tenant-id"
|
||||
mock_instance.provider = "test-provider"
|
||||
mock_instance.model_name = "test-model"
|
||||
mock_instance.get_model_schema.return_value = Mock(features=[])
|
||||
return mock_instance
|
||||
def create_model_instance() -> ModelInstance:
|
||||
"""Build a rerank model with isolated schema and plugin invocation boundaries."""
|
||||
config = make_model_config(provider="test-provider", model="test-model", mode="chat")
|
||||
configuration = config.provider_model_bundle.configuration
|
||||
configuration.provider.supported_model_types = [ModelType.RERANK]
|
||||
schema = config.model_schema.model_copy(update={"model_type": ModelType.RERANK, "features": []})
|
||||
instance = ModelInstance(
|
||||
provider_model_bundle=ProviderModelBundle(
|
||||
configuration=configuration,
|
||||
model_type_instance=RerankModel(
|
||||
provider_schema=configuration.provider,
|
||||
model_runtime=create_plugin_model_runtime(tenant_id=configuration.tenant_id),
|
||||
),
|
||||
),
|
||||
model=config.model,
|
||||
credentials={},
|
||||
)
|
||||
instance.get_model_schema = Mock(return_value=schema)
|
||||
instance.invoke_rerank = Mock()
|
||||
instance.invoke_multimodal_rerank = Mock()
|
||||
return instance
|
||||
|
||||
|
||||
class _UsesSQLiteSession:
|
||||
@@ -70,8 +84,8 @@ class TestRerankModelRunner(_UsesSQLiteSession):
|
||||
|
||||
@pytest.fixture
|
||||
def mock_model_instance(self):
|
||||
"""Create a mock ModelInstance for reranking."""
|
||||
return create_mock_model_instance()
|
||||
"""Create the model with plugin I/O isolated."""
|
||||
return create_model_instance()
|
||||
|
||||
@pytest.fixture
|
||||
def rerank_runner(self, mock_model_instance):
|
||||
@@ -428,7 +442,7 @@ class TestBaseRerankRunner:
|
||||
class TestRerankModelRunnerMultimodal(_UsesSQLiteSession):
|
||||
@pytest.fixture
|
||||
def mock_model_instance(self):
|
||||
return create_mock_model_instance()
|
||||
return create_model_instance()
|
||||
|
||||
@pytest.fixture
|
||||
def rerank_runner(self, mock_model_instance):
|
||||
@@ -455,7 +469,8 @@ class TestRerankModelRunnerMultimodal(_UsesSQLiteSession):
|
||||
docs=[RerankDocument(index=0, text="doc", score=0.88)],
|
||||
)
|
||||
|
||||
rerank_runner.rerank_model_instance.get_model_schema.return_value = Mock(features=[ModelFeature.VISION])
|
||||
schema = rerank_runner.rerank_model_instance.get_model_schema.return_value
|
||||
schema.features = [ModelFeature.VISION]
|
||||
with patch.object(
|
||||
rerank_runner,
|
||||
"fetch_multimodal_rerank",
|
||||
@@ -1092,7 +1107,7 @@ class TestRerankRunnerFactory(_UsesSQLiteSession):
|
||||
- Parameters are forwarded to runner constructor
|
||||
"""
|
||||
# Arrange: Mock model instance
|
||||
mock_model_instance = create_mock_model_instance()
|
||||
mock_model_instance = create_model_instance()
|
||||
|
||||
# Act: Create runner via factory
|
||||
runner = RerankRunnerFactory.create_rerank_runner(
|
||||
@@ -1155,7 +1170,7 @@ class TestRerankRunnerFactory(_UsesSQLiteSession):
|
||||
- String values are properly matched
|
||||
"""
|
||||
# Arrange: Mock model instance
|
||||
mock_model_instance = create_mock_model_instance()
|
||||
mock_model_instance = create_model_instance()
|
||||
|
||||
# Act: Create runner using enum value
|
||||
runner = RerankRunnerFactory.create_rerank_runner(
|
||||
@@ -1186,7 +1201,7 @@ class TestRerankIntegration(_UsesSQLiteSession):
|
||||
- Top results are returned correctly
|
||||
"""
|
||||
# Arrange: Create mock model and documents
|
||||
mock_model_instance = create_mock_model_instance()
|
||||
mock_model_instance = create_model_instance()
|
||||
mock_rerank_result = RerankResult(
|
||||
model="bge-reranker-base",
|
||||
docs=[
|
||||
@@ -1243,7 +1258,7 @@ class TestRerankIntegration(_UsesSQLiteSession):
|
||||
- Normalization is consistent
|
||||
"""
|
||||
# Arrange: Create mock model with various scores
|
||||
mock_model_instance = create_mock_model_instance()
|
||||
mock_model_instance = create_model_instance()
|
||||
mock_rerank_result = RerankResult(
|
||||
model="bge-reranker-base",
|
||||
docs=[
|
||||
@@ -1292,7 +1307,7 @@ class TestRerankEdgeCases(_UsesSQLiteSession):
|
||||
- Empty metadata documents are processed correctly
|
||||
"""
|
||||
# Arrange: Create documents with empty metadata
|
||||
mock_model_instance = create_mock_model_instance()
|
||||
mock_model_instance = create_model_instance()
|
||||
mock_rerank_result = RerankResult(
|
||||
model="bge-reranker-base",
|
||||
docs=[
|
||||
@@ -1338,7 +1353,7 @@ class TestRerankEdgeCases(_UsesSQLiteSession):
|
||||
- Score comparison logic works at boundary
|
||||
"""
|
||||
# Arrange: Create mock with various scores including negatives
|
||||
mock_model_instance = create_mock_model_instance()
|
||||
mock_model_instance = create_model_instance()
|
||||
mock_rerank_result = RerankResult(
|
||||
model="bge-reranker-base",
|
||||
docs=[
|
||||
@@ -1374,7 +1389,7 @@ class TestRerankEdgeCases(_UsesSQLiteSession):
|
||||
- No overflow or precision issues
|
||||
"""
|
||||
# Arrange: All documents with perfect scores
|
||||
mock_model_instance = create_mock_model_instance()
|
||||
mock_model_instance = create_model_instance()
|
||||
mock_rerank_result = RerankResult(
|
||||
model="bge-reranker-base",
|
||||
docs=[
|
||||
@@ -1409,7 +1424,7 @@ class TestRerankEdgeCases(_UsesSQLiteSession):
|
||||
- Content encoding is preserved
|
||||
"""
|
||||
# Arrange: Documents with special characters
|
||||
mock_model_instance = create_mock_model_instance()
|
||||
mock_model_instance = create_model_instance()
|
||||
mock_rerank_result = RerankResult(
|
||||
model="bge-reranker-base",
|
||||
docs=[
|
||||
@@ -1451,7 +1466,7 @@ class TestRerankEdgeCases(_UsesSQLiteSession):
|
||||
- Content is not truncated unexpectedly
|
||||
"""
|
||||
# Arrange: Documents with very long content
|
||||
mock_model_instance = create_mock_model_instance()
|
||||
mock_model_instance = create_model_instance()
|
||||
long_content = "This is a very long document. " * 1000 # ~30,000 characters
|
||||
|
||||
mock_rerank_result = RerankResult(
|
||||
@@ -1488,7 +1503,7 @@ class TestRerankEdgeCases(_UsesSQLiteSession):
|
||||
- All documents are processed correctly
|
||||
"""
|
||||
# Arrange: Create 100 documents
|
||||
mock_model_instance = create_mock_model_instance()
|
||||
mock_model_instance = create_model_instance()
|
||||
num_docs = 100
|
||||
|
||||
# Create rerank results for all documents
|
||||
@@ -1579,7 +1594,7 @@ class TestRerankEdgeCases(_UsesSQLiteSession):
|
||||
- Documents can still be ranked
|
||||
"""
|
||||
# Arrange: Empty query
|
||||
mock_model_instance = create_mock_model_instance()
|
||||
mock_model_instance = create_model_instance()
|
||||
mock_rerank_result = RerankResult(
|
||||
model="bge-reranker-base",
|
||||
docs=[
|
||||
@@ -1626,7 +1641,7 @@ class TestRerankPerformance(_UsesSQLiteSession):
|
||||
- Efficient batch processing
|
||||
"""
|
||||
# Arrange: Multiple documents
|
||||
mock_model_instance = create_mock_model_instance()
|
||||
mock_model_instance = create_model_instance()
|
||||
mock_rerank_result = RerankResult(
|
||||
model="bge-reranker-base",
|
||||
docs=[RerankDocument(index=i, text=f"Doc {i}", score=0.9 - i * 0.1) for i in range(5)],
|
||||
@@ -1736,7 +1751,7 @@ class TestRerankErrorHandling(_UsesSQLiteSession):
|
||||
- Error context is preserved
|
||||
"""
|
||||
# Arrange: Mock model that raises exception
|
||||
mock_model_instance = create_mock_model_instance()
|
||||
mock_model_instance = create_model_instance()
|
||||
mock_model_instance.invoke_rerank.side_effect = RuntimeError("Model invocation failed")
|
||||
|
||||
documents = [
|
||||
@@ -1762,7 +1777,7 @@ class TestRerankErrorHandling(_UsesSQLiteSession):
|
||||
- Invalid results don't corrupt output
|
||||
"""
|
||||
# Arrange: Rerank result with invalid index
|
||||
mock_model_instance = create_mock_model_instance()
|
||||
mock_model_instance = create_model_instance()
|
||||
mock_rerank_result = RerankResult(
|
||||
model="bge-reranker-base",
|
||||
docs=[
|
||||
|
||||
Reference in new issue
Block a user