wehub-resource-sync
392 lines
14 KiB
Python
392 lines
14 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.
|
|
"""Tests for FTS + vector hybrid retrieval via multi-query with reranker."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import pytest
|
|
import zvec
|
|
from zvec import (
|
|
Collection,
|
|
CollectionOption,
|
|
DataType,
|
|
Doc,
|
|
FieldSchema,
|
|
FtsIndexParam,
|
|
HnswIndexParam,
|
|
VectorSchema,
|
|
)
|
|
from zvec.extension.multi_vector_reranker import RrfReRanker, WeightedReRanker
|
|
from zvec.model.param.query import Fts, Query
|
|
|
|
|
|
DIM = 16
|
|
|
|
|
|
# ==================== Fixtures ====================
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
def hybrid_collection(tmp_path_factory) -> Collection:
|
|
"""Collection with one vector field + one FTS field."""
|
|
temp_dir = tmp_path_factory.mktemp("zvec_hybrid")
|
|
collection_path = temp_dir / "hybrid_collection"
|
|
|
|
schema = zvec.CollectionSchema(
|
|
name="hybrid_test",
|
|
fields=[
|
|
FieldSchema("title", DataType.STRING, nullable=False),
|
|
FieldSchema(
|
|
"content",
|
|
DataType.STRING,
|
|
nullable=False,
|
|
index_param=FtsIndexParam(
|
|
tokenizer_name="standard",
|
|
filters=["lowercase"],
|
|
),
|
|
),
|
|
],
|
|
vectors=[
|
|
VectorSchema(
|
|
"embedding",
|
|
DataType.VECTOR_FP32,
|
|
dimension=DIM,
|
|
index_param=HnswIndexParam(),
|
|
),
|
|
],
|
|
)
|
|
|
|
coll = zvec.create_and_open(
|
|
path=str(collection_path),
|
|
schema=schema,
|
|
option=CollectionOption(read_only=False, enable_mmap=True),
|
|
)
|
|
assert coll is not None
|
|
|
|
try:
|
|
yield coll
|
|
finally:
|
|
try:
|
|
coll.destroy()
|
|
except Exception as e:
|
|
print(f"Warning: failed to destroy collection: {e}")
|
|
|
|
|
|
def _make_docs() -> list[Doc]:
|
|
"""Corpus with both text content and vectors.
|
|
|
|
Docs 0-2: AI/ML topic, vectors clustered in one region.
|
|
Docs 3-4: retrieval topic, vectors clustered in another region.
|
|
Doc 5: unrelated topic.
|
|
"""
|
|
# AI cluster vectors
|
|
ai_vec = [1.0] * 8 + [0.0] * 8
|
|
# Retrieval cluster vectors
|
|
ret_vec = [0.0] * 8 + [1.0] * 8
|
|
# Unrelated vector
|
|
other_vec = [0.5] * 16
|
|
|
|
return [
|
|
Doc(
|
|
id="pk_0",
|
|
fields={
|
|
"title": "ML Intro",
|
|
"content": "machine learning is a branch of artificial intelligence",
|
|
},
|
|
vectors={"embedding": ai_vec},
|
|
),
|
|
Doc(
|
|
id="pk_1",
|
|
fields={
|
|
"title": "Deep Learning",
|
|
"content": "deep learning uses neural networks for pattern recognition",
|
|
},
|
|
vectors={"embedding": [0.9] * 8 + [0.1] * 8},
|
|
),
|
|
Doc(
|
|
id="pk_2",
|
|
fields={
|
|
"title": "NLP",
|
|
"content": "natural language processing handles text with artificial intelligence",
|
|
},
|
|
vectors={"embedding": [0.8] * 8 + [0.2] * 8},
|
|
),
|
|
Doc(
|
|
id="pk_3",
|
|
fields={
|
|
"title": "Search Engine",
|
|
"content": "search engine uses inverted index for text retrieval",
|
|
},
|
|
vectors={"embedding": ret_vec},
|
|
),
|
|
Doc(
|
|
id="pk_4",
|
|
fields={
|
|
"title": "Vector DB",
|
|
"content": "vector database enables similarity retrieval and search",
|
|
},
|
|
vectors={"embedding": [0.1] * 8 + [0.9] * 8},
|
|
),
|
|
Doc(
|
|
id="pk_5",
|
|
fields={
|
|
"title": "Cooking",
|
|
"content": "baking bread requires flour water yeast and salt",
|
|
},
|
|
vectors={"embedding": other_vec},
|
|
),
|
|
]
|
|
|
|
|
|
@pytest.fixture(scope="function")
|
|
def hybrid_collection_with_docs(hybrid_collection: Collection) -> Collection:
|
|
"""Hybrid collection pre-populated with test documents."""
|
|
results = hybrid_collection.insert(_make_docs())
|
|
assert all(r.ok() for r in results)
|
|
return hybrid_collection
|
|
|
|
|
|
# ==================== Tests ====================
|
|
|
|
|
|
class TestFtsVectorHybridQuery:
|
|
"""Test FTS + vector hybrid retrieval using multi-query with RRF reranker."""
|
|
|
|
def test_hybrid_fts_and_vector_basic(self, hybrid_collection_with_docs: Collection):
|
|
"""FTS + vector multi-query with RRF reranker returns results."""
|
|
reranker = RrfReRanker(rank_constant=60)
|
|
result = hybrid_collection_with_docs.query(
|
|
queries=[
|
|
Query(field_name="content", fts=Fts(match_string="retrieval")),
|
|
Query(field_name="embedding", vector=[0.0] * 8 + [1.0] * 8),
|
|
],
|
|
topk=5,
|
|
reranker=reranker,
|
|
)
|
|
assert len(result) > 0
|
|
assert len(result) <= 5
|
|
# Results should have scores
|
|
for doc in result:
|
|
assert doc.score > 0
|
|
|
|
def test_hybrid_fts_and_vector_ranking(
|
|
self, hybrid_collection_with_docs: Collection
|
|
):
|
|
"""Docs relevant in both FTS and vector should rank higher."""
|
|
reranker = RrfReRanker(rank_constant=60)
|
|
# FTS: "retrieval search" matches pk_3, pk_4
|
|
# Vector: ret_vec cluster matches pk_3, pk_4
|
|
# Both signals agree: pk_3 and pk_4 should rank top
|
|
result = hybrid_collection_with_docs.query(
|
|
queries=[
|
|
Query(field_name="content", fts=Fts(match_string="retrieval search")),
|
|
Query(field_name="embedding", vector=[0.0] * 8 + [1.0] * 8),
|
|
],
|
|
topk=5,
|
|
reranker=reranker,
|
|
)
|
|
top_ids = {doc.id for doc in result[:3]}
|
|
assert "pk_3" in top_ids or "pk_4" in top_ids
|
|
|
|
def test_hybrid_scores_descending(self, hybrid_collection_with_docs: Collection):
|
|
"""Hybrid query results must be sorted by score descending."""
|
|
reranker = RrfReRanker(rank_constant=60)
|
|
result = hybrid_collection_with_docs.query(
|
|
queries=[
|
|
Query(field_name="content", fts=Fts(match_string="intelligence")),
|
|
Query(field_name="embedding", vector=[1.0] * 8 + [0.0] * 8),
|
|
],
|
|
topk=6,
|
|
reranker=reranker,
|
|
)
|
|
assert len(result) >= 2
|
|
scores = [doc.score for doc in result]
|
|
assert scores == sorted(scores, reverse=True)
|
|
|
|
def test_hybrid_with_filter(self, hybrid_collection_with_docs: Collection):
|
|
"""Hybrid query respects SQL filter."""
|
|
reranker = RrfReRanker(rank_constant=60)
|
|
result = hybrid_collection_with_docs.query(
|
|
queries=[
|
|
Query(field_name="content", fts=Fts(match_string="learning")),
|
|
Query(field_name="embedding", vector=[1.0] * 8 + [0.0] * 8),
|
|
],
|
|
topk=10,
|
|
reranker=reranker,
|
|
filter="title like '%Learning%'",
|
|
)
|
|
for doc in result:
|
|
assert "Learning" in doc.fields["title"]
|
|
|
|
def test_hybrid_fts_no_match_still_returns_vector_results(
|
|
self, hybrid_collection_with_docs: Collection
|
|
):
|
|
"""When FTS matches nothing, vector results still appear."""
|
|
reranker = RrfReRanker(rank_constant=60)
|
|
result = hybrid_collection_with_docs.query(
|
|
queries=[
|
|
Query(
|
|
field_name="content",
|
|
fts=Fts(match_string="nonexistent_term_xyz"),
|
|
),
|
|
Query(field_name="embedding", vector=[1.0] * 8 + [0.0] * 8),
|
|
],
|
|
topk=5,
|
|
reranker=reranker,
|
|
)
|
|
# Vector query alone should still produce results
|
|
assert len(result) > 0
|
|
|
|
def test_hybrid_query_string_syntax(self, hybrid_collection_with_docs: Collection):
|
|
"""Hybrid query works with FTS query_string (advanced syntax)."""
|
|
reranker = RrfReRanker(rank_constant=60)
|
|
result = hybrid_collection_with_docs.query(
|
|
queries=[
|
|
Query(
|
|
field_name="content",
|
|
fts=Fts(query_string="artificial AND intelligence"),
|
|
),
|
|
Query(field_name="embedding", vector=[1.0] * 8 + [0.0] * 8),
|
|
],
|
|
topk=5,
|
|
reranker=reranker,
|
|
)
|
|
assert len(result) > 0
|
|
# pk_0 and pk_2 contain "artificial intelligence"
|
|
hit_ids = {doc.id for doc in result}
|
|
assert "pk_0" in hit_ids or "pk_2" in hit_ids
|
|
|
|
|
|
class TestFtsVectorHybridValidation:
|
|
"""Test validation rules for FTS + vector hybrid queries."""
|
|
|
|
def test_hybrid_requires_reranker(self, hybrid_collection_with_docs: Collection):
|
|
"""Multi-query with FTS + vector without reranker should raise."""
|
|
with pytest.raises(ValueError, match="[Rr]eranker"):
|
|
hybrid_collection_with_docs.query(
|
|
queries=[
|
|
Query(field_name="content", fts=Fts(match_string="learning")),
|
|
Query(field_name="embedding", vector=[1.0] * DIM),
|
|
],
|
|
topk=5,
|
|
)
|
|
|
|
def test_duplicate_field_name_allowed(
|
|
self, hybrid_collection_with_docs: Collection
|
|
):
|
|
"""Multi-query with duplicate field names is allowed and returns results."""
|
|
reranker = RrfReRanker(rank_constant=60)
|
|
result = hybrid_collection_with_docs.query(
|
|
queries=[
|
|
Query(field_name="content", fts=Fts(match_string="learning")),
|
|
Query(field_name="content", fts=Fts(match_string="intelligence")),
|
|
],
|
|
topk=5,
|
|
reranker=reranker,
|
|
)
|
|
assert len(result) > 0
|
|
assert len(result) <= 5
|
|
|
|
def test_multiple_vectors_allowed(self, hybrid_collection_with_docs: Collection):
|
|
"""Two vector queries on the same field are allowed with a reranker."""
|
|
reranker = RrfReRanker(rank_constant=60)
|
|
result = hybrid_collection_with_docs.query(
|
|
queries=[
|
|
Query(field_name="embedding", vector=[1.0] * DIM),
|
|
Query(field_name="embedding", vector=[0.5] * DIM),
|
|
],
|
|
topk=5,
|
|
reranker=reranker,
|
|
)
|
|
assert len(result) > 0
|
|
assert len(result) <= 5
|
|
|
|
|
|
class TestFtsVectorHybridWeightedReranker:
|
|
"""Test FTS + vector hybrid retrieval using WeightedReranker."""
|
|
|
|
def test_weighted_reranker_fts_and_vector(
|
|
self, hybrid_collection_with_docs: Collection
|
|
):
|
|
"""WeightedReranker correctly normalizes FTS scores alongside vector scores."""
|
|
weights = [0.5, 0.5]
|
|
reranker = WeightedReRanker(weights=weights)
|
|
result = hybrid_collection_with_docs.query(
|
|
queries=[
|
|
Query(field_name="content", fts=Fts(match_string="retrieval search")),
|
|
Query(field_name="embedding", vector=[0.0] * 8 + [1.0] * 8),
|
|
],
|
|
topk=5,
|
|
reranker=reranker,
|
|
)
|
|
assert len(result) > 0
|
|
assert len(result) <= 5
|
|
for doc in result:
|
|
assert doc.score > 0
|
|
|
|
def test_weighted_reranker_scores_descending(
|
|
self, hybrid_collection_with_docs: Collection
|
|
):
|
|
"""WeightedReranker hybrid results are sorted by score descending."""
|
|
weights = [0.4, 0.6]
|
|
reranker = WeightedReRanker(weights=weights)
|
|
result = hybrid_collection_with_docs.query(
|
|
queries=[
|
|
Query(field_name="content", fts=Fts(match_string="intelligence")),
|
|
Query(field_name="embedding", vector=[1.0] * 8 + [0.0] * 8),
|
|
],
|
|
topk=6,
|
|
reranker=reranker,
|
|
)
|
|
assert len(result) >= 2
|
|
scores = [doc.score for doc in result]
|
|
assert scores == sorted(scores, reverse=True)
|
|
|
|
def test_weighted_reranker_fts_weight_influence(
|
|
self, hybrid_collection_with_docs: Collection
|
|
):
|
|
"""Higher FTS weight should boost FTS-relevant docs in ranking."""
|
|
# High FTS weight: FTS signal dominates
|
|
weights_fts_heavy = [0.9, 0.1]
|
|
reranker_fts = WeightedReRanker(weights=weights_fts_heavy)
|
|
result_fts = hybrid_collection_with_docs.query(
|
|
queries=[
|
|
Query(field_name="content", fts=Fts(match_string="retrieval")),
|
|
Query(field_name="embedding", vector=[1.0] * 8 + [0.0] * 8),
|
|
],
|
|
topk=5,
|
|
reranker=reranker_fts,
|
|
)
|
|
|
|
# High vector weight: vector signal dominates
|
|
weights_vec_heavy = [0.1, 0.9]
|
|
reranker_vec = WeightedReRanker(weights=weights_vec_heavy)
|
|
result_vec = hybrid_collection_with_docs.query(
|
|
queries=[
|
|
Query(field_name="content", fts=Fts(match_string="retrieval")),
|
|
Query(field_name="embedding", vector=[1.0] * 8 + [0.0] * 8),
|
|
],
|
|
topk=5,
|
|
reranker=reranker_vec,
|
|
)
|
|
|
|
# Both should return results
|
|
assert len(result_fts) > 0
|
|
assert len(result_vec) > 0
|
|
# With FTS-heavy weight, FTS-relevant docs (pk_3, pk_4) should rank higher
|
|
fts_top = [doc.id for doc in result_fts[:2]]
|
|
vec_top = [doc.id for doc in result_vec[:2]]
|
|
# The rankings should differ due to weight difference
|
|
assert fts_top != vec_top or len(result_fts) == len(result_vec) == 1
|