Semantic Search for Enterprise Knowledge Bases: Engineering Beyond Full-Text

import os
from dataclasses import dataclass
from typing import Optional

from rank_bm25 import BM25Okapi
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams, PointStruct
from openai import OpenAI
from sentence_transformers import CrossEncoder
import numpy as np


@dataclass
class RetrievedChunk:
    doc_id: str
    content: str
    bm25_rank: Optional[int]
    vector_rank: Optional[int]
    rrf_score: float
    rerank_score: Optional[float]


class HybridRetrievalPipeline:
    """BM25 + Qdrant vector search fused via RRF, re-ranked by cross-encoder."""

    RRF_K = 60
    COLLECTION_NAME = "enterprise_kb"
    EMBEDDING_MODEL = "text-embedding-3-large"
    EMBEDDING_DIMS = 3072

    def __init__(self, qdrant_url: str = "http://localhost:6333"):
        self.openai = OpenAI(api_key=os.environ["OPENAI_API_KEY"])
        self.qdrant = QdrantClient(url=qdrant_url)
        self.reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")
        self._corpus: list[dict] = []  # [{doc_id, content, tokens}]
        self._bm25: Optional[BM25Okapi] = None

    def build_index(self, documents: list[dict]) -> None:
        """Index documents into both BM25 and Qdrant. documents: [{id, content}]"""
        self._corpus = [
            {"doc_id": d["id"], "content": d["content"], "tokens": d["content"].lower().split()}
            for d in documents
        ]
        self._bm25 = BM25Okapi([c["tokens"] for c in self._corpus])

        # Ensure Qdrant collection exists
        existing = [c.name for c in self.qdrant.get_collections().collections]
        if self.COLLECTION_NAME not in existing:
            self.qdrant.create_collection(
                self.COLLECTION_NAME,
                vectors_config=VectorParams(size=self.EMBEDDING_DIMS, distance=Distance.COSINE),
            )

        # Batch-embed and upsert
        contents = [d["content"] for d in documents]
        response = self.openai.embeddings.create(model=self.EMBEDDING_MODEL, input=contents)
        points = [
            PointStruct(id=i, vector=r.embedding, payload={"doc_id": documents[i]["id"], "content": documents[i]["content"]})
            for i, r in enumerate(response.data)
        ]
        self.qdrant.upsert(collection_name=self.COLLECTION_NAME, points=points)
        print(f"Indexed {len(documents)} documents.")

    def _embed_query(self, query: str) -> list[float]:
        response = self.openai.embeddings.create(model=self.EMBEDDING_MODEL, input=[query])
        return response.data[0].embedding

    def _rrf_fuse(
        self,
        bm25_hits: list[dict],   # [{doc_id, content, rank}]
        vector_hits: list[dict], # [{doc_id, content, rank}]
    ) -> list[RetrievedChunk]:
        scores: dict[str, dict] = {}
        for hit in bm25_hits:
            doc_id = hit["doc_id"]
            scores.setdefault(doc_id, {"content": hit["content"], "bm25_rank": None, "vector_rank": None})
            scores[doc_id]["bm25_rank"] = hit["rank"]
        for hit in vector_hits:
            doc_id = hit["doc_id"]
            scores.setdefault(doc_id, {"content": hit["content"], "bm25_rank": None, "vector_rank": None})
            scores[doc_id]["vector_rank"] = hit["rank"]

        results = []
        for doc_id, meta in scores.items():
            rrf = 0.0
            if meta["bm25_rank"] is not None:
                rrf += 1.0 / (self.RRF_K + meta["bm25_rank"])
            if meta["vector_rank"] is not None:
                rrf += 1.0 / (self.RRF_K + meta["vector_rank"])
            results.append(RetrievedChunk(
                doc_id=doc_id,
                content=meta["content"],
                bm25_rank=meta["bm25_rank"],
                vector_rank=meta["vector_rank"],
                rrf_score=rrf,
            ))
        return sorted(results, key=lambda x: x.rrf_score, reverse=True)

    def retrieve(
        self,
        query: str,
        top_k_retrieval: int = 50,
        top_k_rerank: int = 10,
        use_reranker: bool = True,
    ) -> list[RetrievedChunk]:
        if self._bm25 is None:
            raise RuntimeError("Index not built. Call build_index() first.")

        # BM25 retrieval
        bm25_scores = self._bm25.get_scores(query.lower().split())
        top_bm25_indices = np.argsort(bm25_scores)[::-1][:top_k_retrieval]
        bm25_hits = [
            {"doc_id": self._corpus[i]["doc_id"], "content": self._corpus[i]["content"], "rank": rank + 1}
            for rank, i in enumerate(top_bm25_indices)
        ]

        # Vector retrieval
        query_vector = self._embed_query(query)
        qdrant_results = self.qdrant.search(
            collection_name=self.COLLECTION_NAME,
            query_vector=query_vector,
            limit=top_k_retrieval,
        )
        vector_hits = [
            {"doc_id": r.payload["doc_id"], "content": r.payload["content"], "rank": rank + 1}
            for rank, r in enumerate(qdrant_results)
        ]

        # RRF fusion
        fused = self._rrf_fuse(bm25_hits, vector_hits)
        candidates = fused[:top_k_rerank * 2]  # Re-rank over 2x the desired output

        if not use_reranker:
            return candidates[:top_k_rerank]

        # Cross-encoder re-ranking
        pairs = [[query, c.content] for c in candidates]
        rerank_scores = self.reranker.predict(pairs)
        for chunk, score in zip(candidates, rerank_scores):
            chunk.rerank_score = float(score)

        return sorted(candidates, key=lambda x: x.rerank_score or 0.0, reverse=True)[:top_k_rerank]