from typing import List
import numpy as np
from qdrant_client.http import models
from qdrant_client.hybrid.fusion import reciprocal_rank_fusion, distribution_based_score_fusion
def test_reciprocal_rank_fusion() -> None:
responses = [
[
models.ScoredPoint(id="1", score=0.1, version=1),
models.ScoredPoint(id="2", score=0.2, version=1),
models.ScoredPoint(id="3", score=0.3, version=1),
],
[
models.ScoredPoint(id="5", score=12.0, version=1),
models.ScoredPoint(id="6", score=8.0, version=1),
models.ScoredPoint(id="7", score=5.0, version=1),
models.ScoredPoint(id="2", score=3.0, version=1),
],
]
fused = reciprocal_rank_fusion(responses)
assert fused[0].id == "2"
assert fused[1].id in ["1", "5"]
assert np.isclose(fused[1].score, 1 / 2)
assert fused[2].id in ["1", "5"]
assert np.isclose(fused[2].score, 1 / 2)
def test_distribution_based_score_fusion() -> None:
responses = [
[
models.ScoredPoint(id=1, version=0, score=85.0),
models.ScoredPoint(id=0, version=0, score=76.0),
models.ScoredPoint(id=5, version=0, score=68.0),
],
[
models.ScoredPoint(id=1, version=0, score=62.0),
models.ScoredPoint(id=0, version=0, score=61.0),
models.ScoredPoint(id=4, version=0, score=57.0),
models.ScoredPoint(id=3, version=0, score=51.0),
models.ScoredPoint(id=2, version=0, score=44.0),
],
]
fused = distribution_based_score_fusion(responses, limit=3)
assert fused[0].id == 1
assert fused[1].id == 0
assert fused[2].id == 4
def test_reciprocal_rank_fusion_empty_responses() -> None:
responses: List[List[models.ScoredPoint]] = [[]]
fused = reciprocal_rank_fusion(responses)
assert fused == []
responses = [
[
models.ScoredPoint(id="1", score=0.1, version=1),
models.ScoredPoint(id="2", score=0.2, version=1),
models.ScoredPoint(id="3", score=0.3, version=1),
],
[],
]
fused = reciprocal_rank_fusion(responses)
assert fused[0].id == "1"
assert np.isclose(fused[0].score, 1 / 2)
assert fused[1].id == "2"
assert np.isclose(fused[1].score, 1 / 3)
assert fused[2].id == "3"
assert np.isclose(fused[2].score, 1 / 4)
def test_distribution_based_score_fusion_empty_response() -> None:
responses: List[List[models.ScoredPoint]] = [[]]
fused = distribution_based_score_fusion(responses, limit=3)
assert fused == []
responses = [
[
models.ScoredPoint(id=1, version=0, score=85.0),
models.ScoredPoint(id=0, version=0, score=76.0),
models.ScoredPoint(id=5, version=0, score=68.0),
],
[],
]
fused = distribution_based_score_fusion(responses, limit=3)
assert fused[0].id == 1
assert fused[1].id == 0
assert fused[2].id == 5