wehub-resource-sync
949 lines
33 KiB
Python
949 lines
33 KiB
Python
# Copyright 2025-present the zvec project
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
from __future__ import annotations
|
|
|
|
from unittest.mock import patch, MagicMock
|
|
import pytest
|
|
import os
|
|
|
|
from zvec import Doc, MetricType, VectorSchema, DataType, FlatIndexParam
|
|
from zvec.extension.multi_vector_reranker import (
|
|
CallbackReRanker,
|
|
RrfReRanker,
|
|
WeightedReRanker,
|
|
)
|
|
from zvec.extension.sentence_transformer_rerank_function import (
|
|
DefaultLocalReRanker,
|
|
)
|
|
from zvec.extension.qwen_rerank_function import QwenReRanker
|
|
|
|
# Set ZVEC_RUN_INTEGRATION_TESTS=1 to run real API tests
|
|
RUN_INTEGRATION_TESTS = os.environ.get("ZVEC_RUN_INTEGRATION_TESTS", "0") == "1"
|
|
|
|
|
|
# ----------------------------
|
|
# RrfReRanker Test Case
|
|
# ----------------------------
|
|
class TestRrfReRanker:
|
|
def test_init(self):
|
|
reranker = RrfReRanker(rank_constant=100)
|
|
assert reranker.rank_constant == 100
|
|
|
|
def test_default_rank_constant(self):
|
|
reranker = RrfReRanker()
|
|
assert reranker.rank_constant == 60
|
|
|
|
def test_rerank(self):
|
|
reranker = RrfReRanker(rank_constant=60)
|
|
|
|
doc1 = Doc(id="1", score=0.8)
|
|
doc2 = Doc(id="2", score=0.7)
|
|
doc3 = Doc(id="3", score=0.9)
|
|
doc4 = Doc(id="4", score=0.6)
|
|
|
|
query_results = [[doc1, doc2, doc3], [doc3, doc1, doc4]]
|
|
|
|
results = reranker.rerank(query_results, topn=3)
|
|
|
|
assert len(results) <= 3
|
|
|
|
for doc in results:
|
|
assert hasattr(doc, "score")
|
|
|
|
scores = [doc.score for doc in results]
|
|
assert scores == sorted(scores, reverse=True)
|
|
|
|
|
|
# ----------------------------
|
|
# WeightedReRanker Test Case
|
|
# ----------------------------
|
|
class TestWeightedReRanker:
|
|
@staticmethod
|
|
def _make_fields(metrics):
|
|
return [
|
|
VectorSchema(
|
|
name=f"vector{i}",
|
|
data_type=DataType.VECTOR_FP32,
|
|
dimension=4,
|
|
index_param=FlatIndexParam(metric_type=metric),
|
|
)
|
|
for i, metric in enumerate(metrics)
|
|
]
|
|
|
|
def test_init(self):
|
|
reranker = WeightedReRanker([0.7, 0.3])
|
|
assert reranker.weights == [0.7, 0.3]
|
|
|
|
def test_rerank(self):
|
|
reranker = WeightedReRanker([0.7, 0.3])
|
|
|
|
doc1 = Doc(id="1", score=0.8)
|
|
doc2 = Doc(id="2", score=0.7)
|
|
doc3 = Doc(id="3", score=0.9)
|
|
|
|
query_results = [[doc1, doc2], [doc2, doc3]]
|
|
fields = self._make_fields([MetricType.L2, MetricType.L2])
|
|
|
|
results = reranker.rerank(query_results, topn=3, fields=fields)
|
|
|
|
assert len(results) <= 3
|
|
|
|
for doc in results:
|
|
assert hasattr(doc, "score")
|
|
|
|
|
|
# ----------------------------
|
|
# CallbackReRanker Test Case
|
|
# ----------------------------
|
|
class TestCallbackReRanker:
|
|
def test_rerank(self):
|
|
def my_callback(query_results, fields, topn):
|
|
all_docs = []
|
|
for docs in query_results:
|
|
all_docs.extend(docs)
|
|
all_docs.sort(key=lambda d: d.score, reverse=True)
|
|
return all_docs[:topn]
|
|
|
|
reranker = CallbackReRanker(my_callback)
|
|
|
|
doc1 = Doc(id="1", score=0.8)
|
|
doc2 = Doc(id="2", score=0.9)
|
|
doc3 = Doc(id="3", score=0.7)
|
|
doc4 = Doc(id="4", score=0.6)
|
|
|
|
query_results = [[doc1, doc2], [doc3, doc4]]
|
|
|
|
results = reranker.rerank(query_results, topn=3)
|
|
|
|
assert len(results) == 3
|
|
scores = [doc.score for doc in results]
|
|
assert scores == sorted(scores, reverse=True)
|
|
|
|
def test_callback_with_topn(self):
|
|
received_topn = []
|
|
|
|
def my_callback(query_results, fields, topn):
|
|
received_topn.append(topn)
|
|
return []
|
|
|
|
reranker = CallbackReRanker(my_callback)
|
|
reranker.rerank([[Doc(id="1", score=0.5)]], topn=7)
|
|
|
|
assert received_topn == [7]
|
|
|
|
|
|
# ----------------------------
|
|
# QwenReRanker Test Case
|
|
# ----------------------------
|
|
class TestQwenReRanker:
|
|
def test_init_without_query(self):
|
|
with pytest.raises(ValueError, match="Query is required for QwenReRanker"):
|
|
QwenReRanker(api_key="test_key")
|
|
|
|
def test_init_without_api_key(self):
|
|
with patch.dict(os.environ, {}, clear=True):
|
|
with pytest.raises(ValueError, match="DashScope API key is required"):
|
|
QwenReRanker(query="test")
|
|
|
|
@patch.dict(os.environ, {"DASHSCOPE_API_KEY": "test_key"})
|
|
def test_init_with_env_api_key(self):
|
|
reranker = QwenReRanker(query="test", rerank_field="content")
|
|
assert reranker.query == "test"
|
|
assert reranker._api_key == "test_key"
|
|
assert reranker.rerank_field == "content"
|
|
|
|
def test_init_with_explicit_api_key(self):
|
|
reranker = QwenReRanker(
|
|
query="test", api_key="explicit_key", rerank_field="content"
|
|
)
|
|
assert reranker.query == "test"
|
|
assert reranker._api_key == "explicit_key"
|
|
|
|
def test_model_property(self):
|
|
reranker = QwenReRanker(
|
|
query="test", api_key="test_key", rerank_field="content"
|
|
)
|
|
assert reranker.model == "gte-rerank-v2"
|
|
|
|
reranker = QwenReRanker(
|
|
query="test",
|
|
model="custom-model",
|
|
api_key="test_key",
|
|
rerank_field="content",
|
|
)
|
|
assert reranker.model == "custom-model"
|
|
|
|
def test_query_property(self):
|
|
reranker = QwenReRanker(
|
|
query="test query", api_key="test_key", rerank_field="content"
|
|
)
|
|
assert reranker.query == "test query"
|
|
|
|
def test_rerank_field_property(self):
|
|
reranker = QwenReRanker(query="test", api_key="test_key", rerank_field="title")
|
|
assert reranker.rerank_field == "title"
|
|
|
|
def test_rerank_empty_results(self):
|
|
reranker = QwenReRanker(
|
|
query="test", api_key="test_key", rerank_field="content"
|
|
)
|
|
results = reranker.rerank({})
|
|
assert results == []
|
|
|
|
def test_rerank_no_valid_documents(self):
|
|
reranker = QwenReRanker(
|
|
query="test", api_key="test_key", rerank_field="content"
|
|
)
|
|
# Document without the rerank_field
|
|
query_results = {"vector1": [Doc(id="1")]}
|
|
with pytest.raises(ValueError, match="No documents to rerank"):
|
|
reranker.rerank(query_results)
|
|
|
|
def test_rerank_skip_empty_content(self):
|
|
reranker = QwenReRanker(
|
|
query="test", api_key="test_key", rerank_field="content"
|
|
)
|
|
query_results = {
|
|
"vector1": [
|
|
Doc(id="1", fields={"content": ""}),
|
|
Doc(id="2", fields={"content": " "}),
|
|
]
|
|
}
|
|
with pytest.raises(ValueError, match="No documents to rerank"):
|
|
reranker.rerank(query_results)
|
|
|
|
@patch("zvec.extension.qwen_function.require_module")
|
|
def test_rerank_success(self, mock_require_module):
|
|
# Mock dashscope module
|
|
mock_dashscope = MagicMock()
|
|
mock_require_module.return_value = mock_dashscope
|
|
|
|
# Mock API response
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 200
|
|
mock_response.output = {
|
|
"results": [
|
|
{"index": 0, "relevance_score": 0.95},
|
|
{"index": 1, "relevance_score": 0.85},
|
|
]
|
|
}
|
|
mock_dashscope.TextReRank.call.return_value = mock_response
|
|
|
|
reranker = QwenReRanker(
|
|
query="test query", api_key="test_key", rerank_field="content"
|
|
)
|
|
|
|
query_results = {
|
|
"vector1": [
|
|
Doc(id="1", fields={"content": "Document 1"}),
|
|
Doc(id="2", fields={"content": "Document 2"}),
|
|
]
|
|
}
|
|
|
|
results = reranker.rerank(query_results, topn=2)
|
|
|
|
assert len(results) == 2
|
|
assert results[0].id == "1"
|
|
assert results[0].score == 0.95
|
|
assert results[1].id == "2"
|
|
assert results[1].score == 0.85
|
|
|
|
# Verify API call
|
|
mock_dashscope.TextReRank.call.assert_called_once_with(
|
|
model="gte-rerank-v2",
|
|
query="test query",
|
|
documents=["Document 1", "Document 2"],
|
|
top_n=2,
|
|
return_documents=False,
|
|
)
|
|
|
|
@patch("zvec.extension.qwen_function.require_module")
|
|
def test_rerank_deduplicate_documents(self, mock_require_module):
|
|
# Mock dashscope module
|
|
mock_dashscope = MagicMock()
|
|
mock_require_module.return_value = mock_dashscope
|
|
|
|
# Mock API response
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 200
|
|
mock_response.output = {
|
|
"results": [
|
|
{"index": 0, "relevance_score": 0.9},
|
|
]
|
|
}
|
|
mock_dashscope.TextReRank.call.return_value = mock_response
|
|
|
|
reranker = QwenReRanker(
|
|
query="test", api_key="test_key", rerank_field="content"
|
|
)
|
|
|
|
# Same document in multiple vector results
|
|
doc1 = Doc(id="1", fields={"content": "Document 1"})
|
|
query_results = {"vector1": [doc1], "vector2": [doc1]}
|
|
|
|
results = reranker.rerank(query_results, topn=5)
|
|
|
|
# Should only call API with document once
|
|
call_args = mock_dashscope.TextReRank.call.call_args
|
|
assert len(call_args[1]["documents"]) == 1
|
|
|
|
@patch("zvec.extension.qwen_function.require_module")
|
|
def test_rerank_api_error(self, mock_require_module):
|
|
# Mock dashscope module
|
|
mock_dashscope = MagicMock()
|
|
mock_require_module.return_value = mock_dashscope
|
|
|
|
# Mock API error response
|
|
mock_response = MagicMock()
|
|
mock_response.status_code = 400
|
|
mock_response.message = "Invalid request"
|
|
mock_response.code = "InvalidParameter"
|
|
mock_dashscope.TextReRank.call.return_value = mock_response
|
|
|
|
reranker = QwenReRanker(
|
|
query="test", api_key="test_key", rerank_field="content"
|
|
)
|
|
|
|
query_results = {"vector1": [Doc(id="1", fields={"content": "Document 1"})]}
|
|
|
|
with pytest.raises(ValueError, match="DashScope API error"):
|
|
reranker.rerank(query_results)
|
|
|
|
@patch("zvec.extension.qwen_function.require_module")
|
|
def test_rerank_runtime_error(self, mock_require_module):
|
|
# Mock dashscope module that raises exception
|
|
mock_dashscope = MagicMock()
|
|
mock_require_module.return_value = mock_dashscope
|
|
mock_dashscope.TextReRank.call.side_effect = Exception("Network error")
|
|
|
|
reranker = QwenReRanker(
|
|
query="test", api_key="test_key", rerank_field="content"
|
|
)
|
|
|
|
query_results = {"vector1": [Doc(id="1", fields={"content": "Document 1"})]}
|
|
|
|
with pytest.raises(RuntimeError, match="Failed to call DashScope API"):
|
|
reranker.rerank(query_results)
|
|
|
|
@pytest.mark.skipif(
|
|
not RUN_INTEGRATION_TESTS,
|
|
reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.",
|
|
)
|
|
def test_real_qwen_rerank(self):
|
|
"""Integration test with real DashScope TextReRank API.
|
|
|
|
To run this test, set environment variables:
|
|
export ZVEC_RUN_INTEGRATION_TESTS=1
|
|
export DASHSCOPE_API_KEY=your-api-key
|
|
"""
|
|
# Create reranker with real API
|
|
reranker = QwenReRanker(
|
|
query="What is machine learning?",
|
|
rerank_field="content",
|
|
model="gte-rerank-v2",
|
|
)
|
|
|
|
# Prepare test documents
|
|
query_results = {
|
|
"vector1": [
|
|
Doc(
|
|
id="1",
|
|
score=0.8,
|
|
fields={
|
|
"content": "Machine learning is a subset of artificial intelligence that focuses on building systems that can learn from data."
|
|
},
|
|
),
|
|
Doc(
|
|
id="2",
|
|
score=0.7,
|
|
fields={
|
|
"content": "The weather is nice today with clear skies and sunshine."
|
|
},
|
|
),
|
|
Doc(
|
|
id="3",
|
|
score=0.75,
|
|
fields={
|
|
"content": "Deep learning is a specialized branch of machine learning using neural networks with multiple layers."
|
|
},
|
|
),
|
|
],
|
|
"vector2": [
|
|
Doc(
|
|
id="4",
|
|
score=0.6,
|
|
fields={
|
|
"content": "Python is a popular programming language for data science and machine learning applications."
|
|
},
|
|
),
|
|
Doc(
|
|
id="5",
|
|
score=0.65,
|
|
fields={
|
|
"content": "A recipe for chocolate cake includes flour, sugar, eggs, and cocoa powder."
|
|
},
|
|
),
|
|
],
|
|
}
|
|
|
|
# Call real API
|
|
results = reranker.rerank(query_results, topn=3)
|
|
|
|
# Verify results
|
|
assert len(results) <= 3, "Should return at most topn documents"
|
|
assert len(results) > 0, "Should return at least one document"
|
|
|
|
# All results should have valid scores
|
|
for doc in results:
|
|
assert hasattr(doc, "score"), "Each document should have a score"
|
|
assert isinstance(doc.score, (int, float)), "Score should be numeric"
|
|
assert doc.score > 0, "Score should be positive"
|
|
|
|
# Verify scores are in descending order
|
|
scores = [doc.score for doc in results]
|
|
assert scores == sorted(scores, reverse=True), (
|
|
"Results should be sorted by score in descending order"
|
|
)
|
|
|
|
# Verify relevant documents are ranked higher
|
|
# Document 1 and 3 are about machine learning, should rank higher than weather/recipe docs
|
|
result_ids = [doc.id for doc in results]
|
|
|
|
# At least one of the ML-related documents should be in top results
|
|
ml_related_docs = {"1", "3", "4"}
|
|
assert any(doc_id in ml_related_docs for doc_id in result_ids[:2]), (
|
|
"ML-related documents should rank higher"
|
|
)
|
|
|
|
# Print results for manual verification (useful during development)
|
|
print("\nReranking results:")
|
|
for i, doc in enumerate(results, 1):
|
|
print(f"{i}. ID={doc.id}, Score={doc.score:.4f}")
|
|
if doc.fields:
|
|
content = doc.field("content")
|
|
if content:
|
|
print(f" Content: {content[:80]}...")
|
|
|
|
|
|
# ----------------------------
|
|
# DefaultLocalReRanker Test Case
|
|
# ----------------------------
|
|
class TestDefaultLocalReRanker:
|
|
"""Test cases for DefaultLocalReRanker."""
|
|
|
|
def test_init_without_query(self):
|
|
"""Test initialization fails without query."""
|
|
with pytest.raises(
|
|
ValueError, match="Query is required for DefaultLocalReRanker"
|
|
):
|
|
DefaultLocalReRanker(rerank_field="content")
|
|
|
|
def test_init_with_empty_query(self):
|
|
"""Test initialization fails with empty query."""
|
|
with pytest.raises(
|
|
ValueError, match="Query is required for DefaultLocalReRanker"
|
|
):
|
|
DefaultLocalReRanker(query="", rerank_field="content")
|
|
|
|
@patch("zvec.extension.sentence_transformer_rerank_function.require_module")
|
|
def test_init_success(self, mock_require_module):
|
|
"""Test successful initialization with mocked model."""
|
|
# Mock sentence_transformers module
|
|
mock_st = MagicMock()
|
|
mock_model = MagicMock()
|
|
mock_model.predict = MagicMock() # Cross-encoder has predict method
|
|
mock_model.device = "cpu"
|
|
mock_st.CrossEncoder.return_value = mock_model
|
|
mock_require_module.return_value = mock_st
|
|
|
|
reranker = DefaultLocalReRanker(
|
|
query="test query",
|
|
rerank_field="content",
|
|
model_name="cross-encoder/ms-marco-MiniLM-L6-v2",
|
|
)
|
|
|
|
assert reranker.query == "test query"
|
|
assert reranker.rerank_field == "content"
|
|
assert reranker.model_name == "cross-encoder/ms-marco-MiniLM-L6-v2"
|
|
assert reranker.model_source == "huggingface"
|
|
assert reranker.batch_size == 32
|
|
|
|
@pytest.mark.skipif(
|
|
not RUN_INTEGRATION_TESTS,
|
|
reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.",
|
|
)
|
|
@patch("zvec.extension.sentence_transformer_rerank_function.require_module")
|
|
def test_init_with_custom_params(self, mock_require_module):
|
|
"""Test initialization with custom parameters."""
|
|
mock_st = MagicMock()
|
|
mock_model = MagicMock()
|
|
mock_model.predict = MagicMock()
|
|
mock_model.device = "cuda"
|
|
mock_st.CrossEncoder.return_value = mock_model
|
|
mock_require_module.return_value = mock_st
|
|
|
|
reranker = DefaultLocalReRanker(
|
|
query="custom query",
|
|
rerank_field="title",
|
|
model_name="cross-encoder/ms-marco-MiniLM-L12-v2",
|
|
model_source="modelscope",
|
|
device="cuda",
|
|
batch_size=64,
|
|
)
|
|
|
|
assert reranker.query == "custom query"
|
|
assert reranker.rerank_field == "title"
|
|
assert reranker.model_name == "cross-encoder/ms-marco-MiniLM-L12-v2"
|
|
assert reranker.model_source == "modelscope"
|
|
assert reranker.batch_size == 64
|
|
|
|
@patch("zvec.extension.sentence_transformer_rerank_function.require_module")
|
|
def test_init_invalid_model(self, mock_require_module):
|
|
"""Test initialization fails with non-cross-encoder model."""
|
|
# Mock a model without predict method (not a cross-encoder)
|
|
mock_st = MagicMock()
|
|
mock_model = MagicMock(spec=[]) # No predict method
|
|
mock_st.CrossEncoder.return_value = mock_model
|
|
mock_require_module.return_value = mock_st
|
|
|
|
with pytest.raises(ValueError, match="does not appear to be a cross-encoder"):
|
|
DefaultLocalReRanker(query="test", rerank_field="content")
|
|
|
|
def test_query_property(self):
|
|
"""Test query property."""
|
|
mock_model = MagicMock()
|
|
mock_model.predict = MagicMock()
|
|
|
|
mock_st = MagicMock()
|
|
mock_st.CrossEncoder.return_value = mock_model
|
|
|
|
with patch(
|
|
"zvec.extension.sentence_transformer_rerank_function.require_module",
|
|
return_value=mock_st,
|
|
):
|
|
reranker = DefaultLocalReRanker(query="test query", rerank_field="content")
|
|
assert reranker.query == "test query"
|
|
|
|
def test_rerank_field_property(self):
|
|
"""Test rerank_field property."""
|
|
mock_model = MagicMock()
|
|
mock_model.predict = MagicMock()
|
|
|
|
mock_st = MagicMock()
|
|
mock_st.CrossEncoder.return_value = mock_model
|
|
|
|
with patch(
|
|
"zvec.extension.sentence_transformer_rerank_function.require_module",
|
|
return_value=mock_st,
|
|
):
|
|
reranker = DefaultLocalReRanker(query="test", rerank_field="title")
|
|
assert reranker.rerank_field == "title"
|
|
|
|
def test_batch_size_property(self):
|
|
"""Test batch_size property."""
|
|
mock_model = MagicMock()
|
|
mock_model.predict = MagicMock()
|
|
|
|
mock_st = MagicMock()
|
|
mock_st.CrossEncoder.return_value = mock_model
|
|
|
|
with patch(
|
|
"zvec.extension.sentence_transformer_rerank_function.require_module",
|
|
return_value=mock_st,
|
|
):
|
|
reranker = DefaultLocalReRanker(
|
|
query="test", rerank_field="content", batch_size=128
|
|
)
|
|
assert reranker.batch_size == 128
|
|
|
|
def test_rerank_empty_results(self):
|
|
"""Test rerank with empty query_results."""
|
|
mock_model = MagicMock()
|
|
mock_model.predict = MagicMock()
|
|
|
|
mock_st = MagicMock()
|
|
mock_st.CrossEncoder.return_value = mock_model
|
|
|
|
with patch(
|
|
"zvec.extension.sentence_transformer_rerank_function.require_module",
|
|
return_value=mock_st,
|
|
):
|
|
reranker = DefaultLocalReRanker(query="test", rerank_field="content")
|
|
results = reranker.rerank({})
|
|
assert results == []
|
|
|
|
def test_rerank_no_valid_documents(self):
|
|
"""Test rerank with documents missing rerank_field."""
|
|
mock_model = MagicMock()
|
|
mock_model.predict = MagicMock()
|
|
|
|
mock_st = MagicMock()
|
|
mock_st.CrossEncoder.return_value = mock_model
|
|
|
|
with patch(
|
|
"zvec.extension.sentence_transformer_rerank_function.require_module",
|
|
return_value=mock_st,
|
|
):
|
|
reranker = DefaultLocalReRanker(query="test", rerank_field="content")
|
|
|
|
# Document without the rerank_field
|
|
query_results = {"vector1": [Doc(id="1")]}
|
|
with pytest.raises(ValueError, match="No documents to rerank"):
|
|
reranker.rerank(query_results)
|
|
|
|
def test_rerank_skip_empty_content(self):
|
|
"""Test rerank skips documents with empty content."""
|
|
mock_model = MagicMock()
|
|
mock_model.predict = MagicMock()
|
|
|
|
mock_st = MagicMock()
|
|
mock_st.CrossEncoder.return_value = mock_model
|
|
|
|
with patch(
|
|
"zvec.extension.sentence_transformer_rerank_function.require_module",
|
|
return_value=mock_st,
|
|
):
|
|
reranker = DefaultLocalReRanker(query="test", rerank_field="content")
|
|
|
|
query_results = {
|
|
"vector1": [
|
|
Doc(id="1", fields={"content": ""}),
|
|
Doc(id="2", fields={"content": " "}),
|
|
]
|
|
}
|
|
with pytest.raises(ValueError, match="No documents to rerank"):
|
|
reranker.rerank(query_results)
|
|
|
|
def test_rerank_success(self):
|
|
"""Test successful rerank with mocked model."""
|
|
# Mock standard cross-encoder model
|
|
mock_model = MagicMock()
|
|
|
|
# Mock predict method to return scores
|
|
import numpy as np
|
|
|
|
mock_scores = np.array([0.95, 0.85, 0.75])
|
|
mock_model.predict.return_value = mock_scores
|
|
mock_model.device = "cpu"
|
|
|
|
# Mock sentence_transformers module
|
|
mock_st = MagicMock()
|
|
mock_st.CrossEncoder.return_value = mock_model
|
|
|
|
with patch(
|
|
"zvec.extension.sentence_transformer_rerank_function.require_module",
|
|
return_value=mock_st,
|
|
):
|
|
reranker = DefaultLocalReRanker(query="test query", rerank_field="content")
|
|
|
|
query_results = {
|
|
"vector1": [
|
|
Doc(id="1", score=0.8, fields={"content": "Document 1"}),
|
|
Doc(id="2", score=0.7, fields={"content": "Document 2"}),
|
|
Doc(id="3", score=0.6, fields={"content": "Document 3"}),
|
|
]
|
|
}
|
|
|
|
results = reranker.rerank(query_results, topn=3)
|
|
|
|
# Verify results
|
|
assert len(results) == 3
|
|
assert results[0].id == "1"
|
|
assert results[0].score == 0.95
|
|
assert results[1].id == "2"
|
|
assert results[1].score == 0.85
|
|
assert results[2].id == "3"
|
|
assert results[2].score == 0.75
|
|
|
|
# Verify model.predict was called correctly
|
|
assert mock_model.predict.called
|
|
call_args = mock_model.predict.call_args
|
|
pairs = call_args[0][0]
|
|
assert len(pairs) == 3
|
|
assert pairs[0] == ["test query", "Document 1"]
|
|
assert pairs[1] == ["test query", "Document 2"]
|
|
assert pairs[2] == ["test query", "Document 3"]
|
|
assert call_args[1]["batch_size"] == 32
|
|
assert call_args[1]["show_progress_bar"] is False
|
|
|
|
def test_rerank_with_topn_limit(self):
|
|
"""Test rerank respects topn limit."""
|
|
mock_model = MagicMock()
|
|
|
|
import numpy as np
|
|
|
|
mock_scores = np.array([0.9, 0.8, 0.7, 0.6, 0.5])
|
|
mock_model.predict.return_value = mock_scores
|
|
|
|
# Mock sentence_transformers module
|
|
mock_st = MagicMock()
|
|
mock_st.CrossEncoder.return_value = mock_model
|
|
|
|
with patch(
|
|
"zvec.extension.sentence_transformer_rerank_function.require_module",
|
|
return_value=mock_st,
|
|
):
|
|
reranker = DefaultLocalReRanker(query="test", rerank_field="content")
|
|
|
|
query_results = {
|
|
"vector1": [
|
|
Doc(id="1", fields={"content": "Doc 1"}),
|
|
Doc(id="2", fields={"content": "Doc 2"}),
|
|
Doc(id="3", fields={"content": "Doc 3"}),
|
|
Doc(id="4", fields={"content": "Doc 4"}),
|
|
Doc(id="5", fields={"content": "Doc 5"}),
|
|
]
|
|
}
|
|
|
|
results = reranker.rerank(query_results, topn=2)
|
|
|
|
# Should only return top 2
|
|
assert len(results) == 2
|
|
assert results[0].id == "1"
|
|
assert results[0].score == 0.9
|
|
assert results[1].id == "2"
|
|
assert results[1].score == 0.8
|
|
|
|
def test_rerank_deduplicate_documents(self):
|
|
"""Test rerank deduplicates documents across multiple vectors."""
|
|
mock_model = MagicMock()
|
|
|
|
import numpy as np
|
|
|
|
mock_scores = np.array([0.95, 0.85])
|
|
mock_model.predict.return_value = mock_scores
|
|
|
|
# Mock sentence_transformers module
|
|
mock_st = MagicMock()
|
|
mock_st.CrossEncoder.return_value = mock_model
|
|
|
|
with patch(
|
|
"zvec.extension.sentence_transformer_rerank_function.require_module",
|
|
return_value=mock_st,
|
|
):
|
|
reranker = DefaultLocalReRanker(query="test", rerank_field="content")
|
|
|
|
# Same document in multiple vector results
|
|
doc1 = Doc(id="1", fields={"content": "Document 1"})
|
|
doc2 = Doc(id="2", fields={"content": "Document 2"})
|
|
|
|
query_results = {
|
|
"vector1": [doc1, doc2],
|
|
"vector2": [doc1], # doc1 appears in both
|
|
}
|
|
|
|
results = reranker.rerank(query_results, topn=5)
|
|
|
|
# Should only process each document once
|
|
assert len(results) == 2
|
|
assert mock_model.predict.call_count == 1
|
|
|
|
call_args = mock_model.predict.call_args
|
|
pairs = call_args[0][0]
|
|
assert len(pairs) == 2 # Only 2 unique documents
|
|
|
|
def test_rerank_sorting(self):
|
|
"""Test rerank sorts documents by score in descending order."""
|
|
mock_model = MagicMock()
|
|
|
|
import numpy as np
|
|
|
|
# Return scores in non-sorted order
|
|
mock_scores = np.array([0.6, 0.9, 0.7])
|
|
mock_model.predict.return_value = mock_scores
|
|
|
|
# Mock sentence_transformers module
|
|
mock_st = MagicMock()
|
|
mock_st.CrossEncoder.return_value = mock_model
|
|
|
|
with patch(
|
|
"zvec.extension.sentence_transformer_rerank_function.require_module",
|
|
return_value=mock_st,
|
|
):
|
|
reranker = DefaultLocalReRanker(query="test", rerank_field="content")
|
|
|
|
query_results = {
|
|
"vector1": [
|
|
Doc(id="1", fields={"content": "Doc 1"}),
|
|
Doc(id="2", fields={"content": "Doc 2"}),
|
|
Doc(id="3", fields={"content": "Doc 3"}),
|
|
]
|
|
}
|
|
|
|
results = reranker.rerank(query_results, topn=3)
|
|
|
|
# Should be sorted by score (descending)
|
|
assert len(results) == 3
|
|
assert results[0].id == "2" # score 0.9
|
|
assert results[0].score == 0.9
|
|
assert results[1].id == "3" # score 0.7
|
|
assert results[1].score == 0.7
|
|
assert results[2].id == "1" # score 0.6
|
|
assert results[2].score == 0.6
|
|
|
|
def test_rerank_model_error(self):
|
|
"""Test rerank handles model prediction errors."""
|
|
mock_model = MagicMock()
|
|
|
|
# Mock predict to raise exception
|
|
mock_model.predict.side_effect = Exception("Model inference error")
|
|
|
|
# Mock sentence_transformers module
|
|
mock_st = MagicMock()
|
|
mock_st.CrossEncoder.return_value = mock_model
|
|
|
|
with patch(
|
|
"zvec.extension.sentence_transformer_rerank_function.require_module",
|
|
return_value=mock_st,
|
|
):
|
|
reranker = DefaultLocalReRanker(query="test", rerank_field="content")
|
|
|
|
query_results = {"vector1": [Doc(id="1", fields={"content": "Document 1"})]}
|
|
|
|
with pytest.raises(RuntimeError, match="Failed to compute rerank scores"):
|
|
reranker.rerank(query_results)
|
|
|
|
def test_rerank_with_custom_batch_size(self):
|
|
"""Test rerank uses custom batch_size."""
|
|
mock_model = MagicMock()
|
|
|
|
import numpy as np
|
|
|
|
mock_scores = np.array([0.9, 0.8])
|
|
mock_model.predict.return_value = mock_scores
|
|
|
|
# Mock sentence_transformers module
|
|
mock_st = MagicMock()
|
|
mock_st.CrossEncoder.return_value = mock_model
|
|
|
|
with patch(
|
|
"zvec.extension.sentence_transformer_rerank_function.require_module",
|
|
return_value=mock_st,
|
|
):
|
|
reranker = DefaultLocalReRanker(
|
|
query="test", rerank_field="content", batch_size=64
|
|
)
|
|
|
|
query_results = {
|
|
"vector1": [
|
|
Doc(id="1", fields={"content": "Doc 1"}),
|
|
Doc(id="2", fields={"content": "Doc 2"}),
|
|
]
|
|
}
|
|
|
|
reranker.rerank(query_results)
|
|
|
|
# Verify batch_size is passed to predict
|
|
call_args = mock_model.predict.call_args
|
|
assert call_args[1]["batch_size"] == 64
|
|
|
|
@pytest.mark.skipif(
|
|
not RUN_INTEGRATION_TESTS,
|
|
reason="Integration test skipped. Set ZVEC_RUN_INTEGRATION_TESTS=1 to run.",
|
|
)
|
|
def test_real_sentence_transformer_rerank(self):
|
|
"""Integration test with real SentenceTransformer cross-encoder model.
|
|
|
|
To run this test, set environment variable:
|
|
export ZVEC_RUN_INTEGRATION_TESTS=1
|
|
|
|
Note: This test requires sentence-transformers package and will
|
|
download the MS MARCO MiniLM model (~80MB) on first run.
|
|
"""
|
|
# Create reranker with real model (using default lightweight model)
|
|
reranker = DefaultLocalReRanker(
|
|
query="What is machine learning?",
|
|
rerank_field="content",
|
|
)
|
|
|
|
# Prepare test documents
|
|
query_results = {
|
|
"vector1": [
|
|
Doc(
|
|
id="1",
|
|
score=0.8,
|
|
fields={
|
|
"content": "Machine learning is a subset of artificial intelligence that focuses on building systems that can learn from data."
|
|
},
|
|
),
|
|
Doc(
|
|
id="2",
|
|
score=0.7,
|
|
fields={
|
|
"content": "The weather is nice today with clear skies and sunshine."
|
|
},
|
|
),
|
|
Doc(
|
|
id="3",
|
|
score=0.75,
|
|
fields={
|
|
"content": "Deep learning is a specialized branch of machine learning using neural networks with multiple layers."
|
|
},
|
|
),
|
|
],
|
|
"vector2": [
|
|
Doc(
|
|
id="4",
|
|
score=0.6,
|
|
fields={
|
|
"content": "Python is a popular programming language for data science and machine learning applications."
|
|
},
|
|
),
|
|
Doc(
|
|
id="5",
|
|
score=0.65,
|
|
fields={
|
|
"content": "A recipe for chocolate cake includes flour, sugar, eggs, and cocoa powder."
|
|
},
|
|
),
|
|
],
|
|
}
|
|
|
|
# Call real model
|
|
results = reranker.rerank(query_results, topn=3)
|
|
|
|
# Verify results
|
|
assert len(results) <= 3, "Should return at most topn documents"
|
|
assert len(results) > 0, "Should return at least one document"
|
|
|
|
# All results should have valid scores
|
|
for doc in results:
|
|
assert hasattr(doc, "score"), "Each document should have a score"
|
|
assert isinstance(doc.score, (int, float)), "Score should be numeric"
|
|
|
|
# Verify scores are in descending order
|
|
scores = [doc.score for doc in results]
|
|
assert scores == sorted(scores, reverse=True), (
|
|
"Results should be sorted by score in descending order"
|
|
)
|
|
|
|
# Verify relevant documents are ranked higher
|
|
# Documents 1, 3, and 4 are about machine learning, should rank higher
|
|
result_ids = [doc.id for doc in results]
|
|
|
|
# At least one of the ML-related documents should be in top results
|
|
ml_related_docs = {"1", "3", "4"}
|
|
assert any(doc_id in ml_related_docs for doc_id in result_ids[:2]), (
|
|
"ML-related documents should rank higher"
|
|
)
|
|
|
|
# Print results for manual verification (useful during development)
|
|
print("\nSentenceTransformer Reranking results:")
|
|
for i, doc in enumerate(results, 1):
|
|
print(f"{i}. ID={doc.id}, Score={doc.score:.4f}")
|
|
if doc.fields:
|
|
content = doc.field("content")
|
|
if content:
|
|
print(f" Content: {content[:80]}...")
|