Update rag_engine.py

rag_engine.py

@@ -1,33 +1,44 @@
 import os
 import glob
 from langchain_community.document_loaders import DirectoryLoader, TextLoader, PyPDFLoader, JSONLoader
-from langchain_community.vectorstores import Chroma
+from langchain_community.vectorstores import Qdrant
+from qdrant_client import QdrantClient
+from qdrant_client.http import models
 from langchain_huggingface import HuggingFaceEmbeddings
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_core.documents import Document
 
 class RAGEngine:
-    def __init__(self, knowledge_base_dir="./knowledge_base", persist_directory="./chroma_db"):
+    def __init__(self, knowledge_base_dir="./knowledge_base", persist_directory="./qdrant_db"):
         self.knowledge_base_dir = knowledge_base_dir
         self.persist_directory = persist_directory
+        self.collection_name = "phishing_knowledge"
         
         # Initialize Embeddings (using same model as before)
         self.embedding_fn = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
         
-        # Initialize Vector Store
-        self.vector_store = Chroma(
-            persist_directory=self.persist_directory,
-            embedding_function=self.embedding_fn,
-            collection_name="phishing_knowledge"
+        # Initialize Qdrant Client (Local mode)
+        self.client = QdrantClient(path=self.persist_directory)
+        
+        # Initialize Vector Store wrapper
+        self.vector_store = Qdrant(
+            client=self.client,
+            collection_name=self.collection_name,
+            embeddings=self.embedding_fn
         )
         
-        # Build index if empty or on init
-        if not self.vector_store.get()['ids']:
+        # Check if collection exists and has data
+        try:
+            count = self.client.count(collection_name=self.collection_name).count
+            if count == 0:
+                self._build_index()
+        except:
+            # Collection might not exist yet
             self._build_index()
 
     def _build_index(self):
         """Load documents and build index"""
-        print("🔄 Building Knowledge Base Index...")
+        print("🔄 Building Knowledge Base Index (Qdrant)...")
         
         documents = self._load_documents()
         if not documents:
@@ -43,9 +54,17 @@ class RAGEngine:
         chunks = text_splitter.split_documents(documents)
         
         if chunks:
-            # Add to vector store
-            self.vector_store.add_documents(chunks)
-            self.vector_store.persist()
+            # Re-create collection to ensure clean slate or add to it
+            # For simplicity in local build, we use Qdrant.from_documents which creates/replaces
+            self.vector_store = Qdrant.from_documents(
+                chunks,
+                self.embedding_fn,
+                path=self.persist_directory,
+                collection_name=self.collection_name,
+                force_recreate=True 
+            )
+            # Update the client reference after recreation
+            self.client = self.vector_store.client
             print(f"✅ Indexed {len(chunks)} chunks from {len(documents)} documents.")
         else:
             print("⚠️ No chunks created.")
@@ -95,14 +114,7 @@ class RAGEngine:
     def refresh_knowledge_base(self):
         """Force rebuild of the index"""
         print("♻️ Refreshing Knowledge Base...")
-        # Clear existing collection
-        self.vector_store.delete_collection()
-        self.vector_store = Chroma(
-            persist_directory=self.persist_directory,
-            embedding_function=self.embedding_fn,
-            collection_name="phishing_knowledge"
-        )
-        # Rebuild
+        # In Qdrant local, we can just rebuild with force_recreate=True which is handled in _build_index
         self._build_index()
         return "✅ Knowledge Base Refreshed!"