Update rag_engine.py

rag_engine.py

@@ -1,44 +1,33 @@
 import os
 import glob
 from langchain_community.document_loaders import DirectoryLoader, TextLoader, PyPDFLoader, JSONLoader
-from langchain_community.vectorstores import Qdrant
-from qdrant_client import QdrantClient
-from qdrant_client.http import models
+from langchain_community.vectorstores import Chroma
 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="./qdrant_db"):
+    def __init__(self, knowledge_base_dir="./knowledge_base", persist_directory="./chroma_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 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
+        # Initialize Vector Store
+        self.vector_store = Chroma(
+            persist_directory=self.persist_directory,
+            embedding_function=self.embedding_fn,
+            collection_name="phishing_knowledge"
         )
         
-        # 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
+        # Build index if empty or on init
+        if not self.vector_store.get()['ids']:
             self._build_index()
 
     def _build_index(self):
         """Load documents and build index"""
-        print("🔄 Building Knowledge Base Index (Qdrant)...")
+        print("🔄 Building Knowledge Base Index...")
         
         documents = self._load_documents()
         if not documents:
@@ -54,17 +43,9 @@ class RAGEngine:
         chunks = text_splitter.split_documents(documents)
         
         if chunks:
-            # 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
+            # Add to vector store
+            self.vector_store.add_documents(chunks)
+            self.vector_store.persist()
             print(f"✅ Indexed {len(chunks)} chunks from {len(documents)} documents.")
         else:
             print("⚠️ No chunks created.")
@@ -114,30 +95,23 @@ class RAGEngine:
     def refresh_knowledge_base(self):
         """Force rebuild of the index"""
         print("♻️ Refreshing Knowledge Base...")
-        # In Qdrant local, we can just rebuild with force_recreate=True which is handled in _build_index
+        # 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
         self._build_index()
         return "✅ Knowledge Base Refreshed!"
 
-    def retrieve(self, query, n_results=3, use_mmr=True):
-        """
-        Retrieve relevant context
-        Args:
-            query: Câu truy vấn
-            n_results: Số lượng kết quả trả về
-            use_mmr: Sử dụng MMR (True) hay Similarity Search thường (False)
-        """
-        if use_mmr:
-            results = self.vector_store.max_marginal_relevance_search(
-                query, 
-                k=n_results, 
-                fetch_k=n_results*3, 
-                lambda_mult=0.6
-            )
-        else:
-            # Standard Similarity Search
-            results = self.vector_store.similarity_search(query, k=n_results)
+    def retrieve(self, query, n_results=3):
+        """Retrieve relevant context"""
+        # Search
+        results = self.vector_store.similarity_search(query, k=n_results)
         
         # Format results
         if results:
             return [doc.page_content for doc in results]
-        return []
+        return []