Implement hybrid search with semantic + BM25 weighted fusion

Author: ToastCoder

core/rag_engine.py

@@ -1,40 +1,93 @@
 # HRBuddy
 # core/rag_engine.py
 
-# Required Libraries
 from langchain_community.vectorstores import Chroma
+from langchain_community.retrievers import BM25Retriever
 from langchain_ollama import OllamaEmbeddings
+from langchain_core.documents import Document
 import ollama
 from core.logger import log
 from core.config_loader import cfg
 
-# Main Class    
+
+def _prepare_bm25_chunks(chunks):
+    bm25_chunks = []
+    for doc in chunks:
+        heading = doc.metadata.get("heading", "")
+        text = f"{heading}\n{doc.page_content}" if heading else doc.page_content
+        bm25_chunks.append(Document(page_content=text, metadata=doc.metadata))
+    return bm25_chunks
+
+
+class _HybridRetriever:
+    def __init__(self, retrievers, weights=None, top_k=6):
+        self.retrievers = retrievers
+        self.weights = weights if weights else [1.0 / len(retrievers)] * len(retrievers)
+        self.top_k = top_k
+
+    def invoke(self, query):
+        all_docs = []
+        for retriever, weight in zip(self.retrievers, self.weights):
+            docs = retriever.invoke(query)
+            for rank, doc in enumerate(docs):
+                doc.metadata["_score"] = doc.metadata.get("_score", 0) + weight / (rank + 1)
+                all_docs.append(doc)
+
+        seen = set()
+        unique = []
+        for doc in all_docs:
+            key = doc.page_content[:200]
+            if key not in seen:
+                seen.add(key)
+                unique.append(doc)
+
+        unique.sort(key=lambda d: d.metadata.get("_score", 0), reverse=True)
+        return unique[:self.top_k]
+
+
 class HRBuddyEngine:
     def __init__(self, chunks):
-        """
-        Initialization parameters are pulled entirely from cfg (JSON).
-        """
+        hs = cfg.get("hybrid_search", {})
+        top_k = hs.get("top_k", 6)
+        fetch_k = hs.get("fetch_k", top_k * 3)
+
         log.info(f"[RAG] Initializing with model: {cfg['vector_store']['embedding_model']}")
-        
-        # Embeddings
+
         self.embeddings = OllamaEmbeddings(
             model=cfg["vector_store"]["embedding_model"],
             base_url=cfg.get("ollama_base_url", "http://localhost:11434")
         )
-        
-        # Vector DB
+
         self.vectorstore = Chroma.from_documents(
-            documents=chunks, 
+            documents=chunks,
             embedding=self.embeddings
         )
-        
-        # k (number of documents)
-        search_k = cfg["vector_store"].get("search_top_k", 2)
-        self.retriever = self.vectorstore.as_retriever(search_kwargs={"k": search_k})
-        
-        log.info(f"[RAG] Search depth set to k={search_k}")
-
-    # Generate Response
+        semantic_retriever = self.vectorstore.as_retriever(
+            search_type="mmr",
+            search_kwargs={"k": top_k, "fetch_k": fetch_k}
+        )
+
+        bm25_chunks = _prepare_bm25_chunks(chunks)
+        bm25_retriever = BM25Retriever.from_documents(
+            documents=bm25_chunks, k=top_k, preprocess_func=str.lower
+        )
+
+        if hs.get("enabled", True):
+            semantic_w = hs.get("semantic_weight", 0.7)
+            bm25_w = hs.get("bm25_weight", 0.3)
+            log.info(
+                f"[RAG] Hybrid search enabled "
+                f"(semantic={semantic_w}, bm25={bm25_w}, top_k={top_k})"
+            )
+            self.retriever = _HybridRetriever(
+                retrievers=[semantic_retriever, bm25_retriever],
+                weights=[semantic_w, bm25_w],
+                top_k=top_k
+            )
+        else:
+            log.info("[RAG] Hybrid search disabled, falling back to semantic only")
+            self.retriever = semantic_retriever
+
     def generate_response(self, user_input, history, session_id):
         docs = self.retriever.invoke(user_input)
         context_block = "\n\n".join([doc.page_content for doc in docs])