hybrid search - Qdrant + BM25 + RRF, switch LLM to Groq

src/rag.py

@@ -1,13 +1,13 @@
 import os
-from google import genai
+from groq import Groq
 from dotenv import load_dotenv
 from src.embedder import embed_query
 from src.vector_store import query as vector_query
 
 load_dotenv()
-client = genai.Client(api_key=os.environ["GOOGLE_API_KEY"])
+client = Groq(api_key=os.environ["GROQ_API_KEY"])
 
-LLM_MODEL = "gemini-2.5-flash"
+LLM_MODEL = "llama-3.3-70b-versatile"
 
 
 def build_prompt(question: str, contexts: list[dict]) -> str:
@@ -30,13 +30,16 @@ TRẢ LỜI:"""
 def answer(question: str, top_k: int = 5) -> dict:
     """RAG pipeline: embed query -> retrieve -> generate."""
     query_vec = embed_query(question)
-    contexts = vector_query(query_vec, top_k=top_k)
+    contexts = vector_query(query_vec, query_text=question, top_k=top_k)
 
     prompt = build_prompt(question, contexts)
-    response = client.models.generate_content(model=LLM_MODEL, contents=prompt)
+    response = client.chat.completions.create(
+        model=LLM_MODEL,
+        messages=[{"role": "user", "content": prompt}],
+    )
 
     return {
-        "answer": response.text,
+        "answer": response.choices[0].message.content,
         "sources": [
             {"source": c["metadata"]["source"], "page": c["metadata"]["page"], "score": round(c["score"], 3)}
             for c in contexts

src/vector_store.py

@@ -1,57 +1,121 @@
 import time
-import chromadb
+import pickle
+from pathlib import Path
+from qdrant_client import QdrantClient
+from qdrant_client.models import Distance, VectorParams, PointStruct
+from rank_bm25 import BM25Okapi
 from tqdm import tqdm
 from src.embedder import embed_texts
 
-CHROMA_PATH = "chroma_db"
+QDRANT_PATH = "qdrant_storage"
 COLLECTION_NAME = "rag_docs"
+BM25_PATH = "bm25_index.pkl"
 BATCH_SIZE = 20
+VECTOR_DIM = 3072  # gemini-embedding-001
 
 
-def get_collection():
-    client = chromadb.PersistentClient(path=CHROMA_PATH)
-    return client.get_or_create_collection(
-        name=COLLECTION_NAME,
-        metadata={"hnsw:space": "cosine"},
-    )
+def get_client():
+    return QdrantClient(path=QDRANT_PATH)
+
+
+def get_or_create_collection(client):
+    existing = [c.name for c in client.get_collections().collections]
+    if COLLECTION_NAME not in existing:
+        client.create_collection(
+            collection_name=COLLECTION_NAME,
+            vectors_config=VectorParams(size=VECTOR_DIM, distance=Distance.COSINE),
+        )
+    return client.get_collection(COLLECTION_NAME)
 
 
 def add_chunks(chunks: list[dict]):
-    """Embed và lưu chunks vào Chroma."""
-    collection = get_collection()
+    """Embed và lưu chunks vào Qdrant + build BM25 index."""
+    client = get_client()
+    get_or_create_collection(client)
+
     texts = [c["text"] for c in chunks]
     metadatas = [c["metadata"] for c in chunks]
-    ids = [f"{m['source']}_p{m['page']}_c{m['chunk_index']}" for m in metadatas]
 
-    print(f"Embedding {len(chunks)} chunks (rate limit: ~80 req/min)...")
+    print(f"Embedding {len(chunks)} chunks...")
+    all_embeddings = []
     for i in tqdm(range(0, len(chunks), BATCH_SIZE)):
         batch_texts = texts[i:i + BATCH_SIZE]
-        batch_meta = metadatas[i:i + BATCH_SIZE]
-        batch_ids = ids[i:i + BATCH_SIZE]
         embeddings = embed_texts(batch_texts)
-        collection.add(
-            documents=batch_texts,
-            embeddings=embeddings,
-            metadatas=batch_meta,
-            ids=batch_ids,
-        )
-        # Free tier: 100 req/min, mỗi batch=BATCH_SIZE req → nghỉ để tránh vượt quota
+        all_embeddings.extend(embeddings)
         time.sleep(BATCH_SIZE * 1.5)
 
+    # Lưu vào Qdrant
+    points = [
+        PointStruct(id=i, vector=emb, payload={"text": texts[i], **metadatas[i]})
+        for i, emb in enumerate(all_embeddings)
+    ]
+    client.upsert(collection_name=COLLECTION_NAME, points=points)
+
+    # Build BM25 index
+    print("Building BM25 index...")
+    tokenized = [t.lower().split() for t in texts]
+    bm25 = BM25Okapi(tokenized)
+    with open(BM25_PATH, "wb") as f:
+        pickle.dump({"bm25": bm25, "texts": texts, "metadatas": metadatas}, f)
+
+    print(f"Đã lưu {len(chunks)} chunks vào Qdrant + BM25")
+
+
+def _load_bm25():
+    with open(BM25_PATH, "rb") as f:
+        return pickle.load(f)
+
+
+def _reciprocal_rank_fusion(rankings: list[list[int]], k: int = 60) -> list[tuple[int, float]]:
+    """Kết hợp nhiều bảng xếp hạng bằng RRF."""
+    scores = {}
+    for ranking in rankings:
+        for rank, doc_id in enumerate(ranking):
+            scores[doc_id] = scores.get(doc_id, 0) + 1 / (k + rank + 1)
+    return sorted(scores.items(), key=lambda x: x[1], reverse=True)
+
+
+def query(query_embedding: list[float], query_text: str, top_k: int = 5) -> list[dict]:
+    """Hybrid retrieval: vector + BM25 + RRF."""
+    client = get_client()
+    top_n = top_k * 4  # lấy nhiều hơn để RRF có đủ candidates
+
+    # Vector search
+    vector_results = client.query_points(
+        collection_name=COLLECTION_NAME,
+        query=query_embedding,
+        limit=top_n,
+    ).points
+    vector_ids = [r.id for r in vector_results]
+    vector_payload = {r.id: r.payload for r in vector_results}
+
+    # BM25 search
+    bm25_data = _load_bm25()
+    bm25 = bm25_data["bm25"]
+    texts = bm25_data["texts"]
+    metadatas = bm25_data["metadatas"]
+
+    tokenized_query = query_text.lower().split()
+    bm25_scores = bm25.get_scores(tokenized_query)
+    bm25_ids = sorted(range(len(bm25_scores)), key=lambda i: bm25_scores[i], reverse=True)[:top_n]
+
+    # RRF
+    fused = _reciprocal_rank_fusion([vector_ids, bm25_ids])[:top_k]
 
-def query(query_embedding: list[float], top_k: int = 5) -> list[dict]:
-    """Tìm top-k chunks gần nhất."""
-    collection = get_collection()
-    results = collection.query(
-        query_embeddings=[query_embedding],
-        n_results=top_k,
-        include=["documents", "metadatas", "distances"],
-    )
+    # Build kết quả
     hits = []
-    for doc, meta, dist in zip(
-        results["documents"][0],
-        results["metadatas"][0],
-        results["distances"][0],
-    ):
-        hits.append({"text": doc, "metadata": meta, "score": 1 - dist})
+    for doc_id, rrf_score in fused:
+        if doc_id in vector_payload:
+            payload = vector_payload[doc_id]
+            hits.append({
+                "text": payload["text"],
+                "metadata": {k: v for k, v in payload.items() if k != "text"},
+                "score": round(rrf_score, 4),
+            })
+        elif doc_id < len(texts):
+            hits.append({
+                "text": texts[doc_id],
+                "metadata": metadatas[doc_id],
+                "score": round(rrf_score, 4),
+            })
     return hits