From 2f87a8be98db7ebae9e34469a25ac583ba0147b6 Mon Sep 17 00:00:00 2001 From: "manthapavankumar11@gmail.com" Date: Sat, 18 Apr 2026 20:24:05 +0530 Subject: [PATCH] fixed couple of issue for using the model while creating collection and sparse search --- README.md | 35 +++++++- src/qql/ast_nodes.py | 4 +- src/qql/cli.py | 3 + src/qql/executor.py | 46 ++++++++++- src/qql/parser.py | 31 ++++++- tests/test_executor.py | 180 +++++++++++++++++++++++++++++++++++++++++ tests/test_parser.py | 95 ++++++++++++++++++++++ 7 files changed, 387 insertions(+), 7 deletions(-) diff --git a/README.md b/README.md index 14f88c8..8c5837b 100644 --- a/README.md +++ b/README.md @@ -252,6 +252,7 @@ SEARCH SIMILAR TO '' LIMIT USING MODEL ' SIMILAR TO '' LIMIT [USING MODEL ''] WHERE SEARCH SIMILAR TO '' LIMIT USING HYBRID SEARCH SIMILAR TO '' LIMIT USING HYBRID [DENSE MODEL ''] [SPARSE MODEL ''] [WHERE ] +SEARCH SIMILAR TO '' LIMIT USING SPARSE [MODEL ''] SEARCH SIMILAR TO '' LIMIT EXACT SEARCH SIMILAR TO '' LIMIT [USING ...] [WHERE ] [RERANK] WITH { hnsw_ef: , exact: true|false, acorn: true|false } SEARCH SIMILAR TO '' LIMIT [USING ...] [WHERE ] RERANK [MODEL ''] @@ -284,6 +285,16 @@ Hybrid search with a WHERE filter: SEARCH articles SIMILAR TO 'transformers' LIMIT 10 USING HYBRID WHERE year >= 2020 ``` +Sparse-only search (queries only the `sparse` named vector — useful for pure keyword retrieval): +```sql +SEARCH medical_knowledge SIMILAR TO 'beta blocker contraindications' LIMIT 5 USING SPARSE +``` + +Sparse-only with a custom SPLADE model: +```sql +SEARCH medical_knowledge SIMILAR TO 'beta blocker contraindications' LIMIT 5 USING SPARSE MODEL 'prithivida/Splade_PP_en_v1' +``` + Exact search for recall debugging: ```sql SEARCH articles SIMILAR TO 'attention mechanism' LIMIT 10 EXACT @@ -498,7 +509,12 @@ Hybrid search combines **dense semantic vectors** and **sparse BM25 keyword vect A hybrid collection stores both a named dense vector (`"dense"`) and a named sparse vector (`"sparse"`): ```sql +-- Shorthand (backward compatible) CREATE COLLECTION articles HYBRID + +-- USING form — allows specifying a dense model +CREATE COLLECTION articles USING HYBRID +CREATE COLLECTION articles USING HYBRID DENSE MODEL 'BAAI/bge-base-en-v1.5' ``` This is equivalent to calling Qdrant with: @@ -695,21 +711,34 @@ Explicitly creates a new empty collection. Collections are also created automati ``` CREATE COLLECTION CREATE COLLECTION HYBRID +CREATE COLLECTION USING MODEL '' +CREATE COLLECTION USING HYBRID +CREATE COLLECTION USING HYBRID DENSE MODEL '' ``` **Examples:** -Dense-only collection (standard): +Dense-only collection (standard, uses default model dimensions): ```sql CREATE COLLECTION research_papers ``` -Hybrid collection (dense + sparse BM25): +Dense-only collection pinned to a specific model (768-dimensional): +```sql +CREATE COLLECTION research_papers USING MODEL 'BAAI/bge-base-en-v1.5' +``` + +Hybrid collection (dense + sparse BM25, default models): ```sql CREATE COLLECTION research_papers HYBRID ``` -The collection is created using the **default embedding model's dimensions** (384 for `all-MiniLM-L6-v2`) with **cosine distance**. +Hybrid collection with a custom dense model: +```sql +CREATE COLLECTION research_papers USING HYBRID DENSE MODEL 'BAAI/bge-base-en-v1.5' +``` + +When `USING MODEL` is omitted, the collection uses the **default embedding model's dimensions** (384 for `all-MiniLM-L6-v2`). Specify `USING MODEL` to pin the collection to a specific model's output size — this must match the model you use in INSERT and SEARCH. If the collection already exists, the command succeeds with a message and does nothing. diff --git a/src/qql/ast_nodes.py b/src/qql/ast_nodes.py index 392b4a0..49eadf1 100644 --- a/src/qql/ast_nodes.py +++ b/src/qql/ast_nodes.py @@ -133,6 +133,7 @@ class InsertStmt: class CreateCollectionStmt: collection: str hybrid: bool = False # if True, create with dense + sparse named vectors + model: str | None = None # dense model; None → use config default @dataclass(frozen=True) @@ -152,7 +153,8 @@ class SearchStmt: limit: int model: str | None # dense model; None → use config default hybrid: bool = False # if True, use prefetch+RRF hybrid search - sparse_model: str | None = None # sparse model for hybrid; None → SparseEmbedder.DEFAULT_MODEL + sparse_only: bool = False # if True, query only the sparse vector (no dense) + sparse_model: str | None = None # sparse model for hybrid/sparse-only; None → SparseEmbedder.DEFAULT_MODEL query_filter: FilterExpr | None = None # optional WHERE clause; default keeps existing tests valid rerank: bool = False # if True, apply cross-encoder reranking post-Qdrant rerank_model: str | None = None # cross-encoder model; None → CrossEncoderEmbedder.DEFAULT_MODEL diff --git a/src/qql/cli.py b/src/qql/cli.py index 32594df..bf1b3fe 100644 --- a/src/qql/cli.py +++ b/src/qql/cli.py @@ -30,6 +30,8 @@ [yellow]CREATE COLLECTION[/yellow] [[yellow]HYBRID[/yellow]] Create a new collection. Add HYBRID for dense+sparse BM25 vectors. + Optional: [yellow]USING MODEL[/yellow] '' + Optional: [yellow]USING HYBRID[/yellow] [DENSE MODEL ''] [yellow]DROP COLLECTION[/yellow] Delete a collection and all its points. @@ -41,6 +43,7 @@ Semantic search by vector similarity. Optional: [yellow]USING MODEL[/yellow] '' Optional: [yellow]USING HYBRID[/yellow] [DENSE MODEL ''] [SPARSE MODEL ''] + Optional: [yellow]USING SPARSE[/yellow] [MODEL ''] sparse-vector-only search Optional: [yellow]WHERE[/yellow] (e.g. WHERE year > 2020 AND status = 'ok') Optional: [yellow]RERANK[/yellow] [MODEL ''] rerank results with a cross-encoder Optional: [yellow]EXACT[/yellow] bypass HNSW and perform exact search diff --git a/src/qql/executor.py b/src/qql/executor.py index 68ac535..0144420 100644 --- a/src/qql/executor.py +++ b/src/qql/executor.py @@ -177,9 +177,11 @@ def _execute_create(self, node: CreateCollectionStmt) -> ExecutionResult: message=f"Collection '{node.collection}' already exists", ) + dense_model_name = node.model or self._config.default_model + # ── Hybrid collection: named dense + sparse vectors ──────────────── if node.hybrid: - embedder = Embedder(self._config.default_model) + embedder = Embedder(dense_model_name) dims = embedder.dimensions self._client.create_collection( collection_name=node.collection, @@ -199,7 +201,7 @@ def _execute_create(self, node: CreateCollectionStmt) -> ExecutionResult: ) # ── Standard dense-only collection ───────────────────────────────── - embedder = Embedder(self._config.default_model) + embedder = Embedder(dense_model_name) dims = embedder.dimensions self._client.create_collection( collection_name=node.collection, @@ -302,6 +304,46 @@ def _execute_search(self, node: SearchStmt) -> ExecutionResult: data=results, ) + # ── Sparse-only SEARCH: query the "sparse" named vector directly ───── + if node.sparse_only: + sparse_model_name = node.sparse_model or SparseEmbedder.DEFAULT_MODEL + sparse_embedder = SparseEmbedder(sparse_model_name) + sparse_obj = sparse_embedder.query_embed(node.query_text) + sparse_vector = SparseVector( + indices=sparse_obj["indices"], + values=sparse_obj["values"], + ) + + try: + response = self._client.query_points( + collection_name=node.collection, + query=sparse_vector, + using="sparse", + limit=fetch_limit, + query_filter=qdrant_filter, + ) + except UnexpectedResponse as e: + raise QQLRuntimeError(f"Qdrant error during SEARCH: {e}") from e + + results = [ + {"id": str(h.id), "score": round(h.score, 4), "payload": h.payload} + for h in response.points + ] + + if node.rerank: + results = self._apply_reranking(node.query_text, results, node.limit, node.rerank_model) + return ExecutionResult( + success=True, + message=f"Found {len(results)} result(s) (sparse, reranked)", + data=results, + ) + + return ExecutionResult( + success=True, + message=f"Found {len(results)} result(s) (sparse)", + data=results, + ) + # ── Standard dense-only SEARCH ───────────────────────────────────── model_name = node.model or self._config.default_model embedder = Embedder(model_name) diff --git a/src/qql/parser.py b/src/qql/parser.py index 9d81fd1..4e5b340 100644 --- a/src/qql/parser.py +++ b/src/qql/parser.py @@ -107,10 +107,31 @@ def _parse_create(self) -> CreateCollectionStmt: self._expect(TokenKind.COLLECTION) collection = self._parse_identifier() hybrid: bool = False + model: str | None = None + if self._peek().kind == TokenKind.HYBRID: + # Bare HYBRID shorthand — backward compat self._advance() hybrid = True - return CreateCollectionStmt(collection=collection, hybrid=hybrid) + elif self._peek().kind == TokenKind.USING: + self._advance() # consume USING + if self._peek().kind == TokenKind.HYBRID: + self._advance() # consume HYBRID + hybrid = True + # Optional DENSE MODEL sub-clause + if self._peek().kind == TokenKind.DENSE: + self._advance() # consume DENSE + self._expect(TokenKind.MODEL) + model = self._expect(TokenKind.STRING).value + else: + self._expect(TokenKind.MODEL) + model = self._expect(TokenKind.STRING).value + + return CreateCollectionStmt( + collection=collection, + hybrid=hybrid, + model=model, + ) def _parse_drop(self) -> DropCollectionStmt: self._expect(TokenKind.DROP) @@ -139,6 +160,7 @@ def _parse_search(self) -> SearchStmt: model: str | None = None hybrid: bool = False + sparse_only: bool = False sparse_model: str | None = None if self._peek().kind == TokenKind.USING: self._advance() # consume USING @@ -154,6 +176,12 @@ def _parse_search(self) -> SearchStmt: model = m else: sparse_model = m + elif self._peek().kind == TokenKind.SPARSE: + self._advance() # consume SPARSE + sparse_only = True + if self._peek().kind == TokenKind.MODEL: + self._advance() # consume MODEL + sparse_model = self._expect(TokenKind.STRING).value else: self._expect(TokenKind.MODEL) model = self._expect(TokenKind.STRING).value @@ -196,6 +224,7 @@ def _parse_search(self) -> SearchStmt: limit=limit, model=model, hybrid=hybrid, + sparse_only=sparse_only, sparse_model=sparse_model, query_filter=query_filter, rerank=rerank, diff --git a/tests/test_executor.py b/tests/test_executor.py index 8320977..04fc0a7 100644 --- a/tests/test_executor.py +++ b/tests/test_executor.py @@ -107,6 +107,70 @@ def test_create_existing_collection_is_noop(self, executor, mock_client): assert "already exists" in result.message +class TestCreateWithModel: + def test_create_with_model_passes_model_to_embedder(self, mock_client, cfg, mocker): + mock_emb = mocker.MagicMock() + mock_emb.dimensions = 768 + embedder_cls = mocker.patch("qql.executor.Embedder", return_value=mock_emb) + executor = Executor(mock_client, cfg) + node = CreateCollectionStmt(collection="col", model="BAAI/bge-base-en-v1.5") + executor.execute(node) + embedder_cls.assert_called_once_with("BAAI/bge-base-en-v1.5") + + def test_create_without_model_uses_default_model(self, mock_client, cfg, mocker): + mock_emb = mocker.MagicMock() + mock_emb.dimensions = 384 + embedder_cls = mocker.patch("qql.executor.Embedder", return_value=mock_emb) + executor = Executor(mock_client, cfg) + node = CreateCollectionStmt(collection="col") + executor.execute(node) + embedder_cls.assert_called_once_with(cfg.default_model) + + def test_create_hybrid_with_model_uses_named_vectors(self, mock_client, cfg, mocker): + mock_emb = mocker.MagicMock() + mock_emb.dimensions = 768 + embedder_cls = mocker.patch("qql.executor.Embedder", return_value=mock_emb) + executor = Executor(mock_client, cfg) + node = CreateCollectionStmt(collection="col", hybrid=True, model="BAAI/bge-base-en-v1.5") + executor.execute(node) + embedder_cls.assert_called_once_with("BAAI/bge-base-en-v1.5") + kw = mock_client.create_collection.call_args.kwargs + assert isinstance(kw["vectors_config"], dict) + assert "dense" in kw["vectors_config"] + assert "sparse_vectors_config" in kw + + def test_create_hybrid_without_model_uses_default(self, mock_client, cfg, mocker): + mock_emb = mocker.MagicMock() + mock_emb.dimensions = 384 + embedder_cls = mocker.patch("qql.executor.Embedder", return_value=mock_emb) + executor = Executor(mock_client, cfg) + node = CreateCollectionStmt(collection="col", hybrid=True) + executor.execute(node) + embedder_cls.assert_called_once_with(cfg.default_model) + kw = mock_client.create_collection.call_args.kwargs + assert isinstance(kw["vectors_config"], dict) + + def test_create_dense_with_model_uses_scalar_vectors(self, mock_client, cfg, mocker): + from qdrant_client.models import VectorParams + mock_emb = mocker.MagicMock() + mock_emb.dimensions = 768 + mocker.patch("qql.executor.Embedder", return_value=mock_emb) + executor = Executor(mock_client, cfg) + node = CreateCollectionStmt(collection="col", model="BAAI/bge-base-en-v1.5") + executor.execute(node) + kw = mock_client.create_collection.call_args.kwargs + assert isinstance(kw["vectors_config"], VectorParams) + assert "sparse_vectors_config" not in kw + + def test_create_existing_noop_with_model(self, executor, mock_client): + mock_client.collection_exists.return_value = True + node = CreateCollectionStmt(collection="col", model="some/model") + result = executor.execute(node) + mock_client.create_collection.assert_not_called() + assert result.success is True + assert "already exists" in result.message + + class TestDrop: def test_drop_existing_collection(self, executor, mock_client): mock_client.collection_exists.return_value = True @@ -938,3 +1002,119 @@ def test_rerank_hybrid_search_message( result = executor.execute(node) assert "hybrid" in result.message assert "reranked" in result.message + + +class TestSparseOnlySearch: + @pytest.fixture + def mock_sparse(self, mocker): + mock = mocker.MagicMock() + mock.query_embed.return_value = FAKE_SPARSE + mocker.patch("qql.executor.SparseEmbedder", return_value=mock) + return mock + + def test_sparse_only_calls_query_embed( + self, executor, mock_client, mock_sparse, mocker + ): + mock_client.collection_exists.return_value = True + mock_resp = mocker.MagicMock() + mock_resp.points = [] + mock_client.query_points.return_value = mock_resp + + node = SearchStmt( + collection="col", query_text="q", limit=5, model=None, + sparse_only=True, + ) + executor.execute(node) + mock_sparse.query_embed.assert_called_once_with("q") + + def test_sparse_only_queries_sparse_vector_name( + self, executor, mock_client, mock_sparse, mocker + ): + mock_client.collection_exists.return_value = True + mock_resp = mocker.MagicMock() + mock_resp.points = [] + mock_client.query_points.return_value = mock_resp + + node = SearchStmt( + collection="col", query_text="q", limit=5, model=None, + sparse_only=True, + ) + executor.execute(node) + kw = mock_client.query_points.call_args.kwargs + assert kw["using"] == "sparse" + + def test_sparse_only_message_contains_sparse( + self, executor, mock_client, mock_sparse, mocker + ): + mock_client.collection_exists.return_value = True + mock_resp = mocker.MagicMock() + mock_resp.points = [] + mock_client.query_points.return_value = mock_resp + + node = SearchStmt( + collection="col", query_text="q", limit=5, model=None, + sparse_only=True, + ) + result = executor.execute(node) + assert "sparse" in result.message + + def test_sparse_only_uses_custom_model( + self, executor, mock_client, mocker + ): + mock_client.collection_exists.return_value = True + mock_resp = mocker.MagicMock() + mock_resp.points = [] + mock_client.query_points.return_value = mock_resp + + mock_sparse = mocker.MagicMock() + mock_sparse.query_embed.return_value = FAKE_SPARSE + sparse_cls = mocker.patch("qql.executor.SparseEmbedder", return_value=mock_sparse) + + node = SearchStmt( + collection="col", query_text="q", limit=5, model=None, + sparse_only=True, sparse_model="prithivida/Splade_PP_en_v1", + ) + executor.execute(node) + sparse_cls.assert_called_once_with("prithivida/Splade_PP_en_v1") + + def test_sparse_only_uses_default_model_when_none( + self, executor, mock_client, mocker + ): + mock_client.collection_exists.return_value = True + mock_resp = mocker.MagicMock() + mock_resp.points = [] + mock_client.query_points.return_value = mock_resp + + mock_sparse = mocker.MagicMock() + mock_sparse.query_embed.return_value = FAKE_SPARSE + sparse_cls = mocker.patch("qql.executor.SparseEmbedder", return_value=mock_sparse) + # Make DEFAULT_MODEL on the mock class resolve to the real value so the + # executor's `node.sparse_model or SparseEmbedder.DEFAULT_MODEL` uses it. + sparse_cls.DEFAULT_MODEL = "Qdrant/bm25" + + node = SearchStmt( + collection="col", query_text="q", limit=5, model=None, + sparse_only=True, + ) + executor.execute(node) + sparse_cls.assert_called_once_with("Qdrant/bm25") + + def test_sparse_only_with_rerank_message( + self, executor, mock_client, mock_sparse, mocker + ): + mock_client.collection_exists.return_value = True + mock_resp = mocker.MagicMock() + mock_resp.points = [] + mock_client.query_points.return_value = mock_resp + + mock_ce = mocker.MagicMock() + mock_ce.rerank.return_value = [] + mocker.patch("qql.executor.CrossEncoderEmbedder", return_value=mock_ce) + + node = SearchStmt( + collection="col", query_text="q", limit=5, model=None, + sparse_only=True, rerank=True, + ) + result = executor.execute(node) + assert "sparse" in result.message + assert "reranked" in result.message diff --git a/tests/test_parser.py b/tests/test_parser.py index 6fa1d8d..002be79 100644 --- a/tests/test_parser.py +++ b/tests/test_parser.py @@ -347,6 +347,48 @@ def test_create_hybrid_case_insensitive(self): assert node.hybrid is True +class TestCreateUsing: + def test_create_using_model(self): + node = parse("CREATE COLLECTION articles USING MODEL 'BAAI/bge-base-en-v1.5'") + assert isinstance(node, CreateCollectionStmt) + assert node.hybrid is False + assert node.model == "BAAI/bge-base-en-v1.5" + + def test_create_using_hybrid(self): + node = parse("CREATE COLLECTION articles USING HYBRID") + assert isinstance(node, CreateCollectionStmt) + assert node.hybrid is True + assert node.model is None + + def test_create_using_hybrid_dense_model(self): + node = parse("CREATE COLLECTION articles USING HYBRID DENSE MODEL 'BAAI/bge-base-en-v1.5'") + assert node.hybrid is True + assert node.model == "BAAI/bge-base-en-v1.5" + + def test_create_bare_hybrid_backward_compat(self): + node = parse("CREATE COLLECTION articles HYBRID") + assert node.hybrid is True + assert node.model is None + + def test_create_plain_backward_compat(self): + node = parse("CREATE COLLECTION articles") + assert node.hybrid is False + assert node.model is None + + def test_create_using_model_sets_collection_name(self): + node = parse("CREATE COLLECTION my_col USING MODEL 'some/model'") + assert isinstance(node, CreateCollectionStmt) + assert node.collection == "my_col" + + def test_create_using_hybrid_case_insensitive(self): + node = parse("create collection articles using hybrid") + assert node.hybrid is True + + def test_create_using_model_case_insensitive(self): + node = parse("create collection articles using model 'some/model'") + assert node.model == "some/model" + + class TestHybridInsert: def test_insert_using_hybrid_sets_flag(self): node = parse("INSERT INTO COLLECTION col VALUES {'text': 'hi'} USING HYBRID") @@ -616,3 +658,56 @@ def test_with_unknown_keyword_raises(self): def test_with_trailing_comma(self): node = parse("SEARCH col SIMILAR TO 'q' LIMIT 5 WITH { hnsw_ef: 256, }") assert node.with_clause.hnsw_ef == 256 + + +class TestSparseOnlySearch: + def test_using_sparse_sets_flag(self): + node = parse("SEARCH col SIMILAR TO 'q' LIMIT 5 USING SPARSE") + assert node.sparse_only is True + assert node.hybrid is False + assert node.sparse_model is None + + def test_using_sparse_with_model(self): + node = parse( + "SEARCH col SIMILAR TO 'q' LIMIT 5 USING SPARSE MODEL 'prithivida/Splade_PP_en_v1'" + ) + assert node.sparse_only is True + assert node.sparse_model == "prithivida/Splade_PP_en_v1" + + def test_using_sparse_default_flags(self): + """All other fields remain at their defaults when USING SPARSE is used.""" + node = parse("SEARCH col SIMILAR TO 'q' LIMIT 5 USING SPARSE") + assert node.hybrid is False + assert node.model is None + assert node.rerank is False + assert node.query_filter is None + + def test_using_sparse_with_where(self): + node = parse( + "SEARCH col SIMILAR TO 'q' LIMIT 5 USING SPARSE WHERE year > 2020" + ) + assert node.sparse_only is True + assert node.query_filter is not None + + def test_using_sparse_with_rerank(self): + node = parse("SEARCH col SIMILAR TO 'q' LIMIT 5 USING SPARSE RERANK") + assert node.sparse_only is True + assert node.rerank is True + + def test_using_sparse_with_model_and_rerank(self): + node = parse( + "SEARCH col SIMILAR TO 'q' LIMIT 5 " + "USING SPARSE MODEL 'prithivida/Splade_PP_en_v1' RERANK" + ) + assert node.sparse_only is True + assert node.sparse_model == "prithivida/Splade_PP_en_v1" + assert node.rerank is True + + def test_sparse_only_false_by_default(self): + node = parse("SEARCH col SIMILAR TO 'q' LIMIT 5") + assert node.sparse_only is False + + def test_sparse_only_false_for_hybrid(self): + node = parse("SEARCH col SIMILAR TO 'q' LIMIT 5 USING HYBRID") + assert node.sparse_only is False + assert node.hybrid is True