Skip to content

Commit a5f9543

Browse files
authored
feat: enable MMR diversity for hybrid search (#39)
* feat: enable MMR diversity for hybrid search MMR (Maximal Marginal Relevance) was artificially blocked for hybrid search in QQL, even though the Qdrant SDK supports it: Prefetch.query accepts NearestQuery, which carries an MMR field. - Remove the hybrid guard from _validate_search_mmr_usage - Wire _build_dense_query() into the dense prefetch of both flat and GROUP BY hybrid paths, so MMR params produce a NearestQuery(mmr=...) instead of a raw vector - Keep the sparse-only and recommend guards (MMR is a dense-space concept and RecommendInput has no MMR field in the Qdrant API) - Replace test_hybrid_search_with_mmr_raises with a new test that asserts NearestQuery + MMR is passed in the dense prefetch * feat: add hybrid search test with MMR grouped prefetch functionality --------- Co-authored-by: srimon12 <srimon12mckv.gmail.com>
1 parent ed49eb1 commit a5f9543

2 files changed

Lines changed: 60 additions & 8 deletions

File tree

src/qql/executor.py

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -872,7 +872,7 @@ def _execute_search(self, node: SearchStmt) -> ExecutionResult:
872872
collection_name=node.collection,
873873
prefetch=[
874874
Prefetch(
875-
query=dense_vector,
875+
query=self._build_dense_query(dense_vector, node.with_clause),
876876
using=topology.dense_using(node.dense_vector),
877877
limit=node.limit * _HYBRID_PREFETCH_MULTIPLIER,
878878
params=search_params,
@@ -1460,8 +1460,6 @@ def _has_mmr(self, with_clause: SearchWith | None) -> bool:
14601460
def _validate_search_mmr_usage(self, node: SearchStmt) -> None:
14611461
if not self._has_mmr(node.with_clause):
14621462
return
1463-
if node.hybrid:
1464-
raise QQLRuntimeError("MMR is not supported with USING HYBRID yet")
14651463
if node.sparse_only:
14661464
raise QQLRuntimeError("MMR is not supported with USING SPARSE yet")
14671465

@@ -1635,7 +1633,7 @@ def _execute_search_groups(
16351633
group_by=node.group_by,
16361634
prefetch=[
16371635
Prefetch(
1638-
query=dense_vector,
1636+
query=self._build_dense_query(dense_vector, node.with_clause),
16391637
using=topology.dense_using(node.dense_vector),
16401638
limit=node.limit * _HYBRID_PREFETCH_MULTIPLIER,
16411639
params=search_params,

tests/test_executor.py

Lines changed: 58 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1242,18 +1242,72 @@ def test_dense_search_with_mmr_uses_nearest_query(self, executor, mock_client, m
12421242
assert query.mmr.diversity == pytest.approx(0.4)
12431243
assert query.mmr.candidates_limit == 25
12441244

1245-
def test_hybrid_search_with_mmr_raises(self, executor, mock_client):
1245+
def test_hybrid_search_with_mmr_uses_nearest_query_in_prefetch(self, executor, mock_client, mocker):
1246+
from qdrant_client.models import NearestQuery
1247+
1248+
mocker.patch("qql.executor.Embedder", return_value=mocker.MagicMock())
1249+
mocker.patch("qql.executor.SparseEmbedder", return_value=mocker.MagicMock())
12461250
mock_client.collection_exists.return_value = True
1251+
1252+
collection_info = mocker.MagicMock()
1253+
collection_info.config.params.vectors = {"dense": {}}
1254+
collection_info.config.params.sparse_vectors = {"sparse": {}}
1255+
mock_client.get_collection.return_value = collection_info
1256+
1257+
mock_response = mocker.MagicMock()
1258+
mock_response.points = []
1259+
mock_client.query_points.return_value = mock_response
1260+
12471261
node = SearchStmt(
12481262
collection="notes",
12491263
query_text="hello",
12501264
limit=5,
12511265
model=None,
12521266
hybrid=True,
1253-
with_clause=SearchWith(mmr_diversity=0.5),
1267+
with_clause=SearchWith(mmr_diversity=0.5, mmr_candidates=30),
12541268
)
1255-
with pytest.raises(QQLRuntimeError, match="MMR is not supported with USING HYBRID yet"):
1256-
executor.execute(node)
1269+
executor.execute(node)
1270+
1271+
prefetch = mock_client.query_points.call_args.kwargs["prefetch"]
1272+
assert prefetch is not None
1273+
dense_prefetch = next(p for p in prefetch if p.using == "dense")
1274+
dense_query = dense_prefetch.query
1275+
assert isinstance(dense_query, NearestQuery)
1276+
assert dense_query.mmr is not None
1277+
assert dense_query.mmr.diversity == pytest.approx(0.5)
1278+
assert dense_query.mmr.candidates_limit == 30
1279+
1280+
def test_hybrid_search_with_mmr_grouped_uses_nearest_query_in_prefetch(self, executor, mock_client, mocker):
1281+
from qdrant_client.models import NearestQuery
1282+
1283+
_mock_hybrid_collection(mock_client)
1284+
mock_response = mocker.MagicMock()
1285+
mock_response.groups = []
1286+
mock_client.query_points_groups.return_value = mock_response
1287+
1288+
mock_sparse_embedder = mocker.MagicMock()
1289+
mock_sparse_embedder.query_embed.return_value = {"indices": [0, 1], "values": [0.5, 0.5]}
1290+
mocker.patch("qql.executor.SparseEmbedder", return_value=mock_sparse_embedder)
1291+
1292+
node = SearchStmt(
1293+
collection="articles",
1294+
query_text="hello",
1295+
limit=5,
1296+
model=None,
1297+
hybrid=True,
1298+
group_by="category",
1299+
with_clause=SearchWith(mmr_diversity=0.5, mmr_candidates=30),
1300+
)
1301+
executor.execute(node)
1302+
1303+
prefetch = mock_client.query_points_groups.call_args.kwargs["prefetch"]
1304+
assert prefetch is not None
1305+
dense_prefetch = next(p for p in prefetch if p.using == "dense")
1306+
dense_query = dense_prefetch.query
1307+
assert isinstance(dense_query, NearestQuery)
1308+
assert dense_query.mmr is not None
1309+
assert dense_query.mmr.diversity == pytest.approx(0.5)
1310+
assert dense_query.mmr.candidates_limit == 30
12571311

12581312
def test_sparse_search_with_mmr_raises(self, executor, mock_client):
12591313
mock_client.collection_exists.return_value = True

0 commit comments

Comments
 (0)