Add mxbai dense + bm25 hybrid search to the backend

Dockerfile

@@ -38,7 +38,6 @@ COPY ./poetry.lock /app
 COPY --from=build-step /app/dist /app/static
 
 RUN poetry install --no-interaction --no-ansi --no-root --without dev
-RUN python -c 'from fastembed.embedding import DefaultEmbedding; DefaultEmbedding("sentence-transformers/all-MiniLM-L6-v2")'
 
 # Finally copy the application source code and install root
 COPY qdrant_demo /app/qdrant_demo

pyproject.toml

@@ -5,15 +5,15 @@ description = "Qdrant vector similarity engine demo"
 authors = ["Andrey Vasnetsov <andrey@vasnetsov.com>"]
 
 [tool.poetry.dependencies]
-python = "~3.11"
+python = ">=3.11,<3.15"
 fastapi = "^0.103.1"
 uvicorn = "^0.18.3"
 psutil = "^5.7.3"
 pandas = "^2.2.3"
 loguru = ">=0.7.2"
 requests = "^2.25.1"
 tqdm = "^4.66.1"
-qdrant-client = { extras = ["fastembed"], version = "1.14.2" }
+qdrant-client = "1.19.0"
 
 [tool.poetry.dev-dependencies]
 

qdrant_demo/config.py

@@ -8,7 +8,22 @@
 QDRANT_URL = os.environ.get("QDRANT_URL", "http://localhost:6333/")
 QDRANT_API_KEY = os.environ.get("QDRANT_API_KEY", "")
 
-COLLECTION_NAME = os.environ.get("COLLECTION_NAME", "text-demo")
-EMBEDDINGS_MODEL = os.environ.get("EMBEDDINGS_MODEL", "sentence-transformers/all-MiniLM-L6-v2")
+COLLECTION_NAME = os.environ.get("COLLECTION_NAME", "startups_hybrid")
+EMBEDDINGS_MODEL = os.environ.get("EMBEDDINGS_MODEL", "mixedbread-ai/mxbai-embed-large-v1")
+# Sparse keyword model. Qdrant/bm25 handles tokenization, stemming, and stopwords;
+# IDF is applied server-side via the collection's sparse modifier.
+SPARSE_EMBEDDINGS_MODEL = os.environ.get("SPARSE_EMBEDDINGS_MODEL", "Qdrant/bm25")
 
-TEXT_FIELD_NAME = "document"
+TEXT_FIELD_NAME = os.environ.get("TEXT_FIELD_NAME", "document")
+
+# Named vectors on the hybrid collection. Leave DENSE_VECTOR_NAME empty for a
+# collection with a single unnamed vector.
+DENSE_VECTOR_NAME = os.environ.get("DENSE_VECTOR_NAME", "dense")
+SPARSE_VECTOR_NAME = os.environ.get("SPARSE_VECTOR_NAME", "sparse")
+
+# Embed the query with Qdrant Cloud server-side inference. Defaults ON: the query
+# model is a 1024-d mxbai, too heavy to embed per-request on a small CPU box.
+# Parse leniently since some hosts keep the surrounding quotes on the value.
+CLOUD_INFERENCE = os.environ.get("CLOUD_INFERENCE", "1").strip().strip('"').strip("'").lower() in ("1", "true", "yes")
+RESULT_LIMIT = int(os.environ.get("RESULT_LIMIT", "20"))
+HYBRID_PREFETCH = int(os.environ.get("HYBRID_PREFETCH", "40"))

qdrant_demo/init_collection_startups.py

@@ -1,99 +1,95 @@
+"""Build the startups collection the search path expects: a named `dense` vector
+(mxbai) plus a `sparse` bm25 keyword vector with IDF, and a text index for keyword
+search. Payload fields are renamed once here to the schema the frontend reads
+(`document`, `logo_url`, `homepage_url`). Both vectors are embedded by Qdrant Cloud
+inference, so the query and document sides use the identical models by construction.
+Documents get no mxbai prefix (the query prefix is added at search time).
+Run:  python -m qdrant_demo.init_collection_startups
+"""
 import json
-import os.path
+import os
 from typing import Iterable
 
 from qdrant_client import QdrantClient, models
 from tqdm import tqdm
 
-from qdrant_demo.config import DATA_DIR, QDRANT_URL, QDRANT_API_KEY, COLLECTION_NAME, TEXT_FIELD_NAME, EMBEDDINGS_MODEL
+from qdrant_demo.config import (
+    DATA_DIR, QDRANT_URL, QDRANT_API_KEY, COLLECTION_NAME, TEXT_FIELD_NAME,
+    EMBEDDINGS_MODEL, SPARSE_EMBEDDINGS_MODEL, DENSE_VECTOR_NAME, SPARSE_VECTOR_NAME,
+)
 
+DENSE_DIM = 1024  # mxbai-embed-large-v1
 
-def read_points() -> Iterable[models.PointStruct]:
-    payload_path = os.path.join(DATA_DIR, 'startups_demo.json')
-    with open(payload_path) as fd:
-        for idx, line in enumerate(fd):
-            obj = json.loads(line)
 
-            # Rename fields to unified schema
-            obj["logo_url"] = obj.pop("images")
-            obj["homepage_url"] = obj.pop("link")
-            obj["document"] = obj.pop("description")
-            yield models.PointStruct(
-                id=idx,
-                vector=models.Document(
-                    text=obj["document"],
-                    model=EMBEDDINGS_MODEL,
-                ),
-                payload=obj,
-            )
+def _prepare(obj: dict) -> dict:
+    # Rename to the unified schema the frontend and search path read.
+    obj["logo_url"] = obj.pop("images", None)
+    obj["homepage_url"] = obj.pop("link", None)
+    obj[TEXT_FIELD_NAME] = obj.pop("description", "")
+    return obj
 
 
-def upload_embeddings():
-    client = QdrantClient(
-        url=QDRANT_URL,
-        api_key=QDRANT_API_KEY,
-        prefer_grpc=True,
-    )
+def _records() -> Iterable[dict]:
+    path = os.path.join(DATA_DIR, "startups_demo.json")
+    with open(path, encoding="utf-8") as fd:
+        for line in fd:
+            line = line.strip()
+            if line:
+                yield _prepare(json.loads(line))
 
-    client.set_model(EMBEDDINGS_MODEL)
 
-    payload_path = os.path.join(DATA_DIR, 'startups_demo.json')
-    payload = []
-    documents = []
+def _doc_text(obj: dict) -> str:
+    # Same text the searcher matches against: name + the document body.
+    return f"{obj.get('name', '')}. {obj.get(TEXT_FIELD_NAME, '')}".strip()
 
-    with open(payload_path) as fd:
-        for line in fd:
-            obj = json.loads(line)
-            # Rename fields to unified schema
-            documents.append(obj.pop('description'))
-            obj["logo_url"] = obj.pop("images")
-            obj["homepage_url"] = obj.pop("link")
-            payload.append(obj)
+
+def build():
+    client = QdrantClient(url=QDRANT_URL, api_key=QDRANT_API_KEY, cloud_inference=True)
 
     if client.collection_exists(COLLECTION_NAME):
-        print(f"Collection {COLLECTION_NAME} already exists. Remove it first.")
+        print(f"{COLLECTION_NAME} exists, recreating.")
         client.delete_collection(COLLECTION_NAME)
 
     client.create_collection(
-        collection_name=COLLECTION_NAME,
-        vectors_config=models.VectorParams(
-            size=client.get_embedding_size(EMBEDDINGS_MODEL),
-            distance=models.Distance.COSINE,
-            on_disk=True,
-        ),
-        # Quantization is optional, but it can significantly reduce the memory usage
+        COLLECTION_NAME,
+        vectors_config={
+            DENSE_VECTOR_NAME: models.VectorParams(
+                size=DENSE_DIM, distance=models.Distance.COSINE, on_disk=True,
+            )
+        },
+        sparse_vectors_config={
+            # bm25 values carry term frequency; IDF is applied at query time.
+            SPARSE_VECTOR_NAME: models.SparseVectorParams(modifier=models.Modifier.IDF)
+        },
         quantization_config=models.ScalarQuantization(
             scalar=models.ScalarQuantizationConfig(
-                type=models.ScalarType.INT8,
-                quantile=0.99,
-                always_ram=True
+                type=models.ScalarType.INT8, quantile=0.99, always_ram=True,
             )
-        )
+        ),
     )
-
-    # Create a payload index for text field.
-    # This index enables text search by the TEXT_FIELD_NAME field.
     client.create_payload_index(
-        collection_name=COLLECTION_NAME,
-        field_name=TEXT_FIELD_NAME,
+        COLLECTION_NAME, field_name=TEXT_FIELD_NAME,
         field_schema=models.TextIndexParams(
-            type=models.TextIndexType.TEXT,
-            tokenizer=models.TokenizerType.WORD,
-            min_token_len=2,
-            max_token_len=20,
-            lowercase=True,
-        )
+            type=models.TextIndexType.TEXT, tokenizer=models.TokenizerType.WORD,
+            min_token_len=2, max_token_len=20, lowercase=True,
+        ),
     )
 
-    # Upload points to the collection
-    # Embeddings will be automatically generated from the Document model
-    client.upload_points(
-        collection_name=COLLECTION_NAME,
-        points=tqdm(read_points()),
-        parallel=4,
-        batch_size=16,
-    )
+    def points() -> Iterable[models.PointStruct]:
+        for idx, obj in enumerate(_records()):
+            text = _doc_text(obj)
+            yield models.PointStruct(
+                id=idx,
+                vector={
+                    DENSE_VECTOR_NAME: models.Document(text=text, model=EMBEDDINGS_MODEL),
+                    SPARSE_VECTOR_NAME: models.Document(text=text, model=SPARSE_EMBEDDINGS_MODEL),
+                },
+                payload=obj,
+            )
+
+    client.upload_points(COLLECTION_NAME, points=tqdm(points()), batch_size=64)
+    print(f"built {COLLECTION_NAME}: {client.count(COLLECTION_NAME).count} points")
 
 
-if __name__ == '__main__':
-    upload_embeddings()
+if __name__ == "__main__":
+    build()

qdrant_demo/neural_searcher.py

@@ -1,27 +1,76 @@
+import os
 import time
-from typing import List
 
 from qdrant_client import QdrantClient, models
 
-from qdrant_demo.config import QDRANT_URL, QDRANT_API_KEY, EMBEDDINGS_MODEL
+from qdrant_demo.config import (
+    QDRANT_URL, QDRANT_API_KEY, EMBEDDINGS_MODEL, SPARSE_EMBEDDINGS_MODEL,
+    DENSE_VECTOR_NAME, SPARSE_VECTOR_NAME, RESULT_LIMIT, HYBRID_PREFETCH,
+    CLOUD_INFERENCE,
+)
 
 
 class NeuralSearcher:
+    """Dense (semantic) and hybrid (dense + bm25 keyword, fused with RRF) search
+    over a collection with a named dense vector and a bm25 sparse vector."""
 
     def __init__(self, collection_name: str):
         self.collection_name = collection_name
-        self.qdrant_client = QdrantClient(url=QDRANT_URL, api_key=QDRANT_API_KEY, prefer_grpc=True)
+        # Fail loudly instead of hanging: measured novel-query latency is well
+        # under a second, so 15s is a generous ceiling that still surfaces a real
+        # problem quickly rather than masking it for a minute.
+        timeout = int(os.environ.get("QDRANT_TIMEOUT", "15"))
+        self.qdrant_client = QdrantClient(
+            url=QDRANT_URL, api_key=QDRANT_API_KEY,
+            cloud_inference=CLOUD_INFERENCE, timeout=timeout,
+        )
+
+    # mxbai is an asymmetric retrieval model: the query gets a prompt prefix, the
+    # stored documents do not, so applying it only here needs no re-indexing.
+    # Qdrant Cloud inference already applies this prompt server-side, so on the
+    # Cloud path the prefix is a verified no-op (identical scores). It matters for
+    # the self-hosted / local-embedding path, where nothing else adds it.
+    QUERY_PREFIX = "Represent this sentence for searching relevant passages: "
+
+    def _dense(self, text: str):
+        return models.Document(text=self.QUERY_PREFIX + text, model=EMBEDDINGS_MODEL)
+
+    def _sparse(self, text: str):
+        return models.Document(text=text, model=SPARSE_EMBEDDINGS_MODEL)
 
-    def search(self, text: str, filter_: dict = None) -> List[dict]:
-        start_time = time.time()
-        hits = self.qdrant_client.query_points(
+    def _dense_only(self, text: str):
+        return self.qdrant_client.query_points(
             collection_name=self.collection_name,
-            query=models.Document(
-                text=text,
-                model=EMBEDDINGS_MODEL,
-            ),
-            query_filter=models.Filter(**filter_) if filter_ else None,
-            limit=5
-        )
-        print(f"Search took {time.time() - start_time} seconds")
-        return [hit.payload for hit in hits.points]
+            query=self._dense(text),
+            using=DENSE_VECTOR_NAME or None,
+            limit=RESULT_LIMIT,
+        ).points
+
+    def search(self, text: str, hybrid: bool = False) -> dict:
+        t0 = time.perf_counter()
+        mode = "hybrid" if hybrid else "semantic"
+
+        if hybrid:
+            hits = self.qdrant_client.query_points(
+                collection_name=self.collection_name,
+                prefetch=[
+                    models.Prefetch(query=self._dense(text), using=DENSE_VECTOR_NAME or None, limit=HYBRID_PREFETCH),
+                    models.Prefetch(query=self._sparse(text), using=SPARSE_VECTOR_NAME, limit=HYBRID_PREFETCH),
+                ],
+                query=models.FusionQuery(fusion=models.Fusion.RRF),
+                limit=RESULT_LIMIT,
+            ).points
+        else:
+            hits = self._dense_only(text)
+
+        latency_ms = round((time.perf_counter() - t0) * 1000)
+        results = [{**hit.payload, "score": hit.score} for hit in hits]
+        stats = {
+            "mode": mode,
+            "embedding_model": EMBEDDINGS_MODEL,
+            # RRF fusion scores (~1/60) are not on the same scale as cosine (~0..1).
+            "score_type": "rrf" if mode == "hybrid" else "cosine",
+            "latency_ms": latency_ms,
+            "results": len(results),
+        }
+        return {"results": results, "stats": stats}

qdrant_demo/service.py

@@ -1,14 +1,20 @@
 import os
+import logging
+from typing import Optional
 
-from fastapi import FastAPI
+from fastapi import FastAPI, HTTPException
 from fastapi.staticfiles import StaticFiles
 
-from qdrant_demo.config import COLLECTION_NAME, STATIC_DIR
+from qdrant_demo.config import (
+    COLLECTION_NAME, STATIC_DIR, RESULT_LIMIT, CLOUD_INFERENCE, EMBEDDINGS_MODEL,
+)
 from qdrant_demo.neural_searcher import NeuralSearcher
 from qdrant_demo.text_searcher import TextSearcher
 
 from fastapi.middleware.cors import CORSMiddleware
 
+logger = logging.getLogger(__name__)
+
 app = FastAPI()
 
 app.add_middleware(
@@ -24,11 +30,39 @@
 
 
 @app.get("/api/search")
-async def read_item(q: str, neural: bool = True):
-    return {
-        "result": neural_searcher.search(text=q)
-        if neural else text_searcher.search(query=q)
-    }
+async def read_item(q: str, mode: Optional[str] = None, neural: Optional[bool] = None):
+    """mode = semantic (dense) | keyword (full-text) | hybrid (dense + keyword).
+
+    Back-compat with the older frontend, which passes `neural` (bool):
+    neural=true -> semantic (its original meaning), neural=false -> keyword.
+    When neither is given, default to hybrid. Explicit `mode` always wins."""
+    if mode is None:
+        mode = "semantic" if neural is True else "keyword" if neural is False else "hybrid"
+    if not q.strip():
+        return {"result": [], "stats": {"mode": mode}}
+    try:
+        if mode == "keyword":
+            return {"result": text_searcher.search(query=q, top=RESULT_LIMIT),
+                    "stats": {"mode": "keyword"}}
+        out = neural_searcher.search(text=q, hybrid=(mode == "hybrid"))
+        return {"result": out["results"], "stats": out["stats"]}
+    except Exception as e:
+        logger.exception("search failed for q=%r mode=%s", q, mode)
+        raise HTTPException(status_code=502, detail="Search is temporarily unavailable.")
+
+
+@app.get("/api/stats")
+async def stats():
+    """Live collection size, so the frontend can show off the scale."""
+    try:
+        count = neural_searcher.qdrant_client.count(COLLECTION_NAME).count
+        return {
+            "count": count, "collection": COLLECTION_NAME,
+            "cloud_inference": CLOUD_INFERENCE, "model": EMBEDDINGS_MODEL,
+        }
+    except Exception:
+        logger.exception("stats failed")
+        raise HTTPException(status_code=502, detail="Stats are temporarily unavailable.")
 
 
 # Mount the static files directory once the search endpoint is defined