Implemented Hybrid Vectors and Reranking

Author: pavanjava

README.md

@@ -1,6 +1,6 @@
 [project]
 name = "qql-cli"
-version = "0.2.0"
+version = "1.0.0"
 description = "A SQL-like query language CLI wrapper for Qdrant vector database"
 readme = "README.md"
 license = { file = "LICENSE" }
@@ -32,6 +32,7 @@ dependencies = [
     "qdrant-client[fastembed]>=1.13.0",
     "click>=8.1.0",
     "rich>=13.0.0",
+    "prompt_toolkit>=3.0.0",
 ]
 
 [project.urls]

pyproject.toml

@@ -116,12 +116,15 @@ class NotExpr:
 class InsertStmt:
     collection: str
     values: dict[str, Any]  # must contain "text" key
-    model: str | None       # None → use default
+    model: str | None       # dense model; None → use config default
+    hybrid: bool = False            # if True, also embed + store sparse BM25 vector
+    sparse_model: str | None = None # sparse model; None → SparseEmbedder.DEFAULT_MODEL
 
 
 @dataclass(frozen=True)
 class CreateCollectionStmt:
     collection: str
+    hybrid: bool = False    # if True, create with dense + sparse named vectors
 
 
 @dataclass(frozen=True)
@@ -139,7 +142,9 @@ class SearchStmt:
     collection: str
     query_text: str
     limit: int
-    model: str | None
+    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
     query_filter: FilterExpr | None = None  # optional WHERE clause; default keeps existing tests valid
 
 

src/qql/ast_nodes.py

@@ -3,8 +3,10 @@
 import sys
 
 import click
+from prompt_toolkit import PromptSession
+from prompt_toolkit.formatted_text import HTML
+from prompt_toolkit.history import InMemoryHistory
 from rich.console import Console
-from rich.prompt import Prompt
 from rich.table import Table
 
 from .config import delete_config, load_config, save_config, QQLConfig
@@ -24,9 +26,10 @@
   [yellow]INSERT INTO COLLECTION[/yellow] <name> [yellow]VALUES[/yellow] {[yellow]'text'[/yellow]: '...', ...}
       Insert a point. 'text' is required and auto-vectorized.
       Optional: [yellow]USING MODEL[/yellow] '<model>'
+      Optional: [yellow]USING HYBRID[/yellow] [DENSE MODEL '<model>'] [SPARSE MODEL '<model>']
 
-  [yellow]CREATE COLLECTION[/yellow] <name>
-      Create a new collection (uses default model dimensions).
+  [yellow]CREATE COLLECTION[/yellow] <name> [[yellow]HYBRID[/yellow]]
+      Create a new collection. Add HYBRID for dense+sparse BM25 vectors.
 
   [yellow]DROP COLLECTION[/yellow] <name>
       Delete a collection and all its points.
@@ -37,10 +40,19 @@
   [yellow]SEARCH[/yellow] <name> [yellow]SIMILAR TO[/yellow] '<text>' [yellow]LIMIT[/yellow] <n>
       Semantic search by vector similarity.
       Optional: [yellow]USING MODEL[/yellow] '<model>'
+      Optional: [yellow]USING HYBRID[/yellow] [DENSE MODEL '<model>'] [SPARSE MODEL '<model>']
+      Optional: [yellow]WHERE[/yellow] <filter>   (e.g. WHERE year > 2020 AND status = 'ok')
 
   [yellow]DELETE FROM[/yellow] <name> [yellow]WHERE id =[/yellow] '<id>'
       Delete a point by its ID.
 
+Keyboard shortcuts:
+  ← → arrows   move cursor within the current line
+  ↑ ↓ arrows   scroll through command history
+  Ctrl-A / Ctrl-E   jump to beginning / end of line
+  Ctrl-C   cancel current input
+  Ctrl-D   exit shell
+
 Type [bold]exit[/bold] or [bold]quit[/bold] to leave the shell.
 """
 
@@ -115,10 +127,16 @@ def _launch_repl(cfg: QQLConfig) -> None:
     console.print(f"[bold cyan]QQL Interactive Shell[/bold cyan]  •  {cfg.url}")
     console.print("Type [bold]help[/bold] for available commands or [bold]exit[/bold] to quit.\n")
 
+    session: PromptSession[str] = PromptSession(history=InMemoryHistory())
+
     while True:
         try:
-            query = Prompt.ask("[bold green]qql>[/bold green]").strip()
-        except (EOFError, KeyboardInterrupt):
+            query = session.prompt(HTML("<ansigreen><b>qql&gt;</b></ansigreen> ")).strip()
+        except KeyboardInterrupt:
+            # Ctrl-C clears the current line; continue the loop
+            continue
+        except EOFError:
+            # Ctrl-D exits
             console.print("\nBye.")
             break
 

src/qql/cli.py

@@ -31,3 +31,37 @@ def embed_batch(self, texts: list[str]) -> list[list[float]]:
     def dimensions(self) -> int:
         """Return the vector dimensionality by embedding a dummy string."""
         return len(self.embed("probe"))
+
+
+class SparseEmbedder:
+    """Sparse BM25-style embedder using fastembed.SparseTextEmbedding.
+
+    Returns dicts with "indices" and "values" lists (not numpy arrays),
+    ready for direct construction of qdrant_client SparseVector objects.
+
+    Uses asymmetric embedding: embed() for document indexing, query_embed()
+    for query-time encoding (BM25 IDF weighting differs at query vs. index time).
+    """
+
+    DEFAULT_MODEL = "Qdrant/bm25"
+
+    # Class-level cache mirrors Embedder's pattern
+    _cache: dict[str, object] = {}
+
+    def __init__(self, model_name: str = DEFAULT_MODEL) -> None:
+        self._model_name = model_name
+        if model_name not in SparseEmbedder._cache:
+            from fastembed import SparseTextEmbedding
+
+            SparseEmbedder._cache[model_name] = SparseTextEmbedding(model_name)
+        self._model = SparseEmbedder._cache[model_name]
+
+    def embed(self, text: str) -> dict[str, list]:
+        """Embed a document string. Returns {"indices": [...], "values": [...]}."""
+        result = next(iter(self._model.embed([text])))  # type: ignore[attr-defined]
+        return {"indices": result.indices.tolist(), "values": result.values.tolist()}
+
+    def query_embed(self, text: str) -> dict[str, list]:
+        """Embed a query string (BM25 applies different IDF weighting at query time)."""
+        result = next(iter(self._model.query_embed(text)))  # type: ignore[attr-defined]
+        return {"indices": result.indices.tolist(), "values": result.values.tolist()}

src/qql/embedder.py

@@ -10,6 +10,8 @@
     Distance,
     FieldCondition,
     Filter,
+    Fusion,
+    FusionQuery,
     IsEmptyCondition,
     IsNullCondition,
     MatchAny,
@@ -18,9 +20,13 @@
     MatchText,
     MatchTextAny,
     MatchValue,
+    Modifier,
     PayloadField,
     PointStruct,
+    Prefetch,
     Range,
+    SparseVector,
+    SparseVectorParams,
     VectorParams,
 )
 
@@ -49,7 +55,7 @@
     ShowCollectionsStmt,
 )
 from .config import QQLConfig
-from .embedder import Embedder
+from .embedder import Embedder, SparseEmbedder
 from .exceptions import QQLRuntimeError
 
 
@@ -86,6 +92,56 @@ def _execute_insert(self, node: InsertStmt) -> ExecutionResult:
         if "text" not in node.values:
             raise QQLRuntimeError("INSERT requires a 'text' field in VALUES")
 
+        # ── Hybrid INSERT: dense + sparse vectors ──────────────────────────
+        if node.hybrid:
+            dense_model = node.model or self._config.default_model
+            sparse_model_name = node.sparse_model or SparseEmbedder.DEFAULT_MODEL
+            dense_embedder = Embedder(dense_model)
+            sparse_embedder = SparseEmbedder(sparse_model_name)
+
+            dense_vector = dense_embedder.embed(node.values["text"])
+            sparse_obj = sparse_embedder.embed(node.values["text"])
+            sparse_vector = SparseVector(
+                indices=sparse_obj["indices"],
+                values=sparse_obj["values"],
+            )
+
+            # Auto-create hybrid collection if it doesn't exist yet
+            if not self._client.collection_exists(node.collection):
+                self._client.create_collection(
+                    collection_name=node.collection,
+                    vectors_config={
+                        "dense": VectorParams(
+                            size=len(dense_vector), distance=Distance.COSINE
+                        )
+                    },
+                    sparse_vectors_config={
+                        "sparse": SparseVectorParams(modifier=Modifier.IDF)
+                    },
+                )
+
+            point_id = str(uuid.uuid4())
+            try:
+                self._client.upsert(
+                    collection_name=node.collection,
+                    points=[
+                        PointStruct(
+                            id=point_id,
+                            vector={"dense": dense_vector, "sparse": sparse_vector},
+                            payload=dict(node.values),
+                        )
+                    ],
+                )
+            except UnexpectedResponse as e:
+                raise QQLRuntimeError(f"Qdrant error during INSERT: {e}") from e
+
+            return ExecutionResult(
+                success=True,
+                message=f"Inserted 1 point [{point_id}] (hybrid)",
+                data={"id": point_id, "collection": node.collection},
+            )
+
+        # ── Standard dense-only INSERT ─────────────────────────────────────
         model_name = node.model or self._config.default_model
         embedder = Embedder(model_name)
         vector = embedder.embed(node.values["text"])
@@ -115,6 +171,29 @@ def _execute_create(self, node: CreateCollectionStmt) -> ExecutionResult:
                 success=True,
                 message=f"Collection '{node.collection}' already exists",
             )
+
+        # ── Hybrid collection: named dense + sparse vectors ────────────────
+        if node.hybrid:
+            embedder = Embedder(self._config.default_model)
+            dims = embedder.dimensions
+            self._client.create_collection(
+                collection_name=node.collection,
+                vectors_config={
+                    "dense": VectorParams(size=dims, distance=Distance.COSINE)
+                },
+                sparse_vectors_config={
+                    "sparse": SparseVectorParams(modifier=Modifier.IDF)
+                },
+            )
+            return ExecutionResult(
+                success=True,
+                message=(
+                    f"Collection '{node.collection}' created "
+                    f"(hybrid: {dims}-dim dense + BM25 sparse, cosine distance)"
+                ),
+            )
+
+        # ── Standard dense-only collection ─────────────────────────────────
         embedder = Embedder(self._config.default_model)
         dims = embedder.dimensions
         self._client.create_collection(
@@ -148,16 +227,64 @@ def _execute_search(self, node: SearchStmt) -> ExecutionResult:
         if not self._client.collection_exists(node.collection):
             raise QQLRuntimeError(f"Collection '{node.collection}' does not exist")
 
-        model_name = node.model or self._config.default_model
-        embedder = Embedder(model_name)
-        vector = embedder.embed(node.query_text)
-
+        # Build WHERE filter (shared by both hybrid and dense-only paths)
         qdrant_filter: Filter | None = None
         if node.query_filter is not None:
             qdrant_filter = self._wrap_as_filter(
                 self._build_qdrant_filter(node.query_filter)
             )
 
+        # ── Hybrid SEARCH: prefetch dense+sparse, fuse with RRF ───────────
+        if node.hybrid:
+            dense_model = node.model or self._config.default_model
+            sparse_model_name = node.sparse_model or SparseEmbedder.DEFAULT_MODEL
+            dense_embedder = Embedder(dense_model)
+            sparse_embedder = SparseEmbedder(sparse_model_name)
+
+            dense_vector = dense_embedder.embed(node.query_text)
+            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,
+                    prefetch=[
+                        Prefetch(
+                            query=dense_vector,
+                            using="dense",
+                            limit=node.limit * 4,
+                        ),
+                        Prefetch(
+                            query=sparse_vector,
+                            using="sparse",
+                            limit=node.limit * 4,
+                        ),
+                    ],
+                    query=FusionQuery(fusion=Fusion.RRF),
+                    limit=node.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
+            ]
+            return ExecutionResult(
+                success=True,
+                message=f"Found {len(results)} result(s) (hybrid)",
+                data=results,
+            )
+
+        # ── Standard dense-only SEARCH ─────────────────────────────────────
+        model_name = node.model or self._config.default_model
+        embedder = Embedder(model_name)
+        vector = embedder.embed(node.query_text)
+
         try:
             response = self._client.query_points(
                 collection_name=node.collection,
@@ -293,16 +420,26 @@ def _wrap_as_filter(self, qdrant_expr: Any) -> Filter:
     # ── Collection helpers ────────────────────────────────────────────────
 
     def _ensure_collection(self, name: str, vector_size: int) -> None:
-        """Create the collection if it doesn't exist. Raises on dimension mismatch."""
+        """Create the collection if it doesn't exist. Raises on dimension mismatch.
+
+        For named-vector (hybrid) collections the validation is skipped — those
+        collections are managed directly by the hybrid insert/create paths.
+        """
         if self._client.collection_exists(name):
             info = self._client.get_collection(name)
-            existing_size = info.config.params.vectors.size  # type: ignore[union-attr]
-            if existing_size != vector_size:
-                raise QQLRuntimeError(
-                    f"Vector dimension mismatch: collection '{name}' expects "
-                    f"{existing_size} dims, but model produces {vector_size} dims. "
-                    f"Specify a compatible model with USING MODEL '<model>'."
-                )
+            vectors = info.config.params.vectors  # type: ignore[union-attr]
+            if isinstance(vectors, dict):
+                # Named-vector (hybrid) collection — skip validation here;
+                # the hybrid insert path manages its own collection creation.
+                pass
+            else:
+                # Unnamed single-vector collection: validate dimensions
+                if vectors.size != vector_size:
+                    raise QQLRuntimeError(
+                        f"Vector dimension mismatch: collection '{name}' expects "
+                        f"{vectors.size} dims, but model produces {vector_size} dims. "
+                        f"Specify a compatible model with USING MODEL '<model>'."
+                    )
         else:
             self._client.create_collection(
                 collection_name=name,

src/qql/executor.py

@@ -12,6 +12,9 @@ class TokenKind(Enum):
     VALUES = auto()
     USING = auto()
     MODEL = auto()
+    HYBRID = auto()
+    DENSE = auto()
+    SPARSE = auto()
     CREATE = auto()
     DROP = auto()
     SHOW = auto()
@@ -69,6 +72,9 @@ class TokenKind(Enum):
     "VALUES": TokenKind.VALUES,
     "USING": TokenKind.USING,
     "MODEL": TokenKind.MODEL,
+    "HYBRID": TokenKind.HYBRID,
+    "DENSE": TokenKind.DENSE,
+    "SPARSE": TokenKind.SPARSE,
     "CREATE": TokenKind.CREATE,
     "DROP": TokenKind.DROP,
     "SHOW": TokenKind.SHOW,