Upload 4 files

main/bm25_retriever.py

@@ -0,0 +1,180 @@
+from rank_bm25 import BM25Okapi
+from typing import List, Dict, Any, Tuple
+import pickle
+import os
+from utils.text_processor import VietnameseTextProcessor
+from config import Config
+from tqdm.auto import tqdm
+
+os.makedirs("index", exist_ok=True)
+
+
+class BM25Retriever:
+    """BM25 retriever for initial document retrieval"""
+
+    def __init__(self):
+        self.text_processor = VietnameseTextProcessor()
+        self.bm25 = None
+        self.documents = []
+        self.tokenized_corpus = []
+        self.index_file = "index/bm25_index.pkl"
+
+    def build_index(self, documents: List[Dict[str, Any]]):
+        """Build BM25 index from documents"""
+        print("Building BM25 index...")
+
+        self.documents = documents
+        self.tokenized_corpus = []
+
+        # Tokenize all documents
+        for doc in tqdm(documents):
+            content = doc.get("content", "")
+            title = doc.get("title", "")
+
+            # Combine title and content for better search
+            full_text = f"{title} {content}"
+
+            # Preprocess and tokenize
+            processed_text = self.text_processor.preprocess_for_search(full_text)
+            tokens = processed_text.split()
+
+            self.tokenized_corpus.append(tokens)
+
+        # Build BM25 index
+        self.bm25 = BM25Okapi(self.tokenized_corpus)
+
+        print(f"BM25 index built with {len(self.documents)} documents")
+
+    def save_index(self, filepath: str = None):
+        """Save BM25 index to file"""
+        if filepath is None:
+            filepath = self.index_file
+
+        try:
+            index_data = {
+                "bm25": self.bm25,
+                "documents": self.documents,
+                "tokenized_corpus": self.tokenized_corpus,
+            }
+
+            with open(filepath, "wb") as f:
+                pickle.dump(index_data, f)
+
+            print(f"BM25 index saved to {filepath}")
+        except Exception as e:
+            print(f"Error saving BM25 index: {e}")
+
+    def load_index(self, filepath: str = None):
+        """Load BM25 index from file"""
+        if filepath is None:
+            filepath = self.index_file
+
+        try:
+            if not os.path.exists(filepath):
+                print(f"Index file {filepath} not found")
+                return False
+
+            with open(filepath, "rb") as f:
+                index_data = pickle.load(f)
+
+            self.bm25 = index_data["bm25"]
+            self.documents = index_data["documents"]
+            self.tokenized_corpus = index_data["tokenized_corpus"]
+
+            print(f"BM25 index loaded from {filepath}")
+            return True
+        except UnicodeDecodeError as e:
+            print(f"Encoding error loading BM25 index: {e}")
+            print(f"Removing corrupted index file: {filepath}")
+            try:
+                os.remove(filepath)
+            except:
+                pass
+            return False
+        except (pickle.UnpicklingError, EOFError) as e:
+            print(f"Corrupted BM25 index file: {e}")
+            print(f"Removing corrupted index file: {filepath}")
+            try:
+                os.remove(filepath)
+            except:
+                pass
+            return False
+        except Exception as e:
+            print(f"Error loading BM25 index: {e}")
+            return False
+
+    def search(
+        self, query: str, top_k: int = None
+    ) -> List[Tuple[Dict[str, Any], float]]:
+        """Search documents using BM25"""
+        if top_k is None:
+            top_k = Config.BM25_TOP_K
+
+        if self.bm25 is None:
+            print("BM25 index not built. Please build index first.")
+            return []
+
+        # Preprocess query
+        processed_query = self.text_processor.preprocess_for_search(query)
+        query_tokens = processed_query.split()
+
+        if not query_tokens:
+            return []
+
+        # Get BM25 scores
+        scores = self.bm25.get_scores(query_tokens)
+
+        # Get top documents with scores
+        doc_score_pairs = [
+            (self.documents[i], scores[i]) for i in range(len(self.documents))
+        ]
+
+        # Sort by score (descending)
+        doc_score_pairs.sort(key=lambda x: x[1], reverse=True)
+
+        # Return top k
+        results = doc_score_pairs[:top_k]
+
+        print(f"BM25 search returned {len(results)} results for query: {query}")
+        return results
+
+    def get_relevant_documents(
+        self, query: str, top_k: int = None, min_score: float = 0.0
+    ) -> List[Dict[str, Any]]:
+        """Get relevant documents above minimum score threshold"""
+        results = self.search(query, top_k)
+
+        # Filter by minimum score
+        filtered_results = [
+            (doc, score) for doc, score in results if score >= min_score
+        ]
+
+        # Return only documents
+        return [doc for doc, score in filtered_results]
+
+    def search_with_keywords(
+        self, keywords: List[str], top_k: int = None
+    ) -> List[Tuple[Dict[str, Any], float]]:
+        """Search using multiple keywords"""
+        # Combine keywords into a single query
+        query = " ".join(keywords)
+        return self.search(query, top_k)
+
+    def get_index_stats(self) -> Dict[str, Any]:
+        """Get statistics about the BM25 index"""
+        if self.bm25 is None:
+            return {}
+
+        return {
+            "total_documents": len(self.documents),
+            "total_tokens": sum(len(tokens) for tokens in self.tokenized_corpus),
+            "average_document_length": sum(
+                len(tokens) for tokens in self.tokenized_corpus
+            )
+            / len(self.tokenized_corpus)
+            if self.tokenized_corpus
+            else 0,
+            "vocabulary_size": len(
+                set(token for tokens in self.tokenized_corpus for token in tokens)
+            ),
+        }

main/chatbot.py

@@ -0,0 +1,652 @@
+from typing import List, Dict, Any, Optional
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain.prompts import PromptTemplate
+from langchain.schema import HumanMessage, SystemMessage
+
+from main.vector_store import QdrantVectorStore
+from main.bm25_retriever import BM25Retriever
+from main.reranker import DocumentReranker
+from utils.text_processor import VietnameseTextProcessor
+from utils.google_search import GoogleSearchTool
+from utils.question_refiner import VietnameseLegalQuestionRefiner
+from config import Config
+
+class VietnameseLegalRAG:
+    """Vietnamese Legal RAG System"""
+    
+    def __init__(self):
+        self.vector_store = None
+        self.bm25_retriever = None
+        self.reranker = None
+        self.llm = None
+        self.text_processor = VietnameseTextProcessor()
+        self.google_search = GoogleSearchTool()
+        self.question_refiner = VietnameseLegalQuestionRefiner()
+        
+        self._initialize_components()
+    
+    def _initialize_components(self):
+        """Initialize RAG components"""
+        try:
+            # Initialize LLM
+            if Config.GOOGLE_API_KEY:
+                self.llm = ChatGoogleGenerativeAI(
+                    model=Config.MODEL_GEN,
+                    google_api_key=Config.GOOGLE_API_KEY,
+                    temperature=0.1
+                )
+                print("Google Gemini LLM initialized")
+            else:
+                print("Warning: Google API key not found")
+            
+            # Initialize vector store
+            self.vector_store = QdrantVectorStore()
+            
+            # Initialize BM25 retriever
+            self.bm25_retriever = BM25Retriever()
+            
+            # Initialize reranker if enabled
+            if Config.ENABLE_RERANKING:
+                self.reranker = DocumentReranker(model_name=Config.RERANKER_MODEL)
+            else:
+                print("Reranking disabled in configuration")
+            
+        except Exception as e:
+            print(f"Error initializing RAG components: {e}")
+    
+    def setup_indices(self, documents: List[Dict[str, Any]], force_rebuild: bool = False):
+        """Setup both vector and BM25 indices"""
+        print("Setting up RAG indices...")
+        
+        try:
+            # Setup vector store
+            if self.vector_store:
+                # Check if we need to create collection
+                try:
+                    # First, do a simple existence check
+                    collections = self.vector_store.client.get_collections().collections
+                    collection_exists = any(col.name == self.vector_store.collection_name for col in collections)
+                    print(f"Collection existence check: {collection_exists}")
+                    
+                    if collection_exists:
+                        # Collection exists, try to get detailed info
+                        try:
+                            collection_info = self.vector_store.get_collection_info()
+                            has_documents = collection_info.get('points_count', 0) > 0
+                            
+                            if force_rebuild:
+                                print("Force rebuild requested - recreating vector store...")
+                                self.vector_store.create_collection(force_recreate=True)
+                                self.vector_store.add_documents(documents)
+                            elif not has_documents:
+                                print("Collection exists but is empty - adding documents...")
+                                self.vector_store.add_documents(documents)
+                            else:
+                                print(f"Vector store collection already exists with {collection_info.get('points_count', 0)} documents")
+                        except Exception as info_e:
+                            print(f"Could not get collection info: {info_e}")
+                            if force_rebuild:
+                                print("Force rebuild requested - recreating vector store...")
+                                self.vector_store.create_collection(force_recreate=True)
+                                self.vector_store.add_documents(documents)
+                            else:
+                                print("Assuming collection has documents - skipping setup")
+                    else:
+                        # Collection doesn't exist, create it
+                        print("Collection does not exist - creating new collection...")
+                        self.vector_store.create_collection(force_recreate=False)
+                        self.vector_store.add_documents(documents)
+                        
+                except Exception as e:
+                    print(f"Error during vector store setup: {e}")
+                    print("Attempting to create collection...")
+                    self.vector_store.create_collection()
+                    self.vector_store.add_documents(documents)
+            
+            # Setup BM25 index
+            if self.bm25_retriever:
+                # Try to load existing index
+                if not self.bm25_retriever.load_index() or force_rebuild:
+                    self.bm25_retriever.build_index(documents)
+                    self.bm25_retriever.save_index()
+                else:
+                    print("BM25 index loaded from file")
+            
+            print("RAG indices setup completed")
+            
+        except Exception as e:
+            print(f"Error setting up indices: {e}")
+            raise
+    
+    def retrieve_documents(self, query: str, use_hybrid: bool = True, use_reranking: bool = None) -> List[Dict[str, Any]]:
+        """Retrieve relevant documents using hybrid approach with optional reranking"""
+        retrieved_docs = []
+        
+        # Use config default if not specified
+        if use_reranking is None:
+            use_reranking = Config.ENABLE_RERANKING
+        
+        # Adjust retrieval counts if reranking is enabled
+        bm25_top_k = Config.RERANK_BEFORE_RETRIEVAL_TOP_K if use_reranking else Config.BM25_TOP_K
+        vector_top_k = Config.RERANK_BEFORE_RETRIEVAL_TOP_K if use_reranking else Config.TOP_K_RETRIEVAL
+        
+        try:
+            if use_hybrid and self.bm25_retriever and self.vector_store:
+                # Hybrid retrieval: BM25 + Vector Search
+                
+                # BM25 retrieval
+                bm25_results = self.bm25_retriever.get_relevant_documents(
+                    query, top_k=bm25_top_k
+                )
+                
+                # Vector search
+                vector_results = self.vector_store.search_similar_documents(
+                    query, top_k=vector_top_k
+                )
+                
+                # Combine and deduplicate results
+                all_docs = {}
+                
+                # Add BM25 results
+                for doc in bm25_results:
+                    doc_id = doc.get('id', '')
+                    if doc_id:
+                        all_docs[doc_id] = {**doc, 'retrieval_method': 'bm25'}
+                
+                # Add vector results
+                for doc in vector_results:
+                    doc_id = doc.get('id', '')
+                    if doc_id:
+                        if doc_id in all_docs:
+                            # Combine scores if document found by both methods
+                            all_docs[doc_id]['retrieval_method'] = 'hybrid'
+                            all_docs[doc_id]['vector_score'] = doc.get('score', 0)
+                        else:
+                            all_docs[doc_id] = {**doc, 'retrieval_method': 'vector'}
+                
+                retrieved_docs = list(all_docs.values())
+                
+            elif self.vector_store:
+                # Vector search only
+                retrieved_docs = self.vector_store.search_similar_documents(query, top_k=vector_top_k)
+                
+            elif self.bm25_retriever:
+                # BM25 only
+                retrieved_docs = self.bm25_retriever.get_relevant_documents(query, top_k=bm25_top_k)
+            
+            # Improved similarity filtering logic
+            if retrieved_docs:
+                # Check for high-quality documents first
+                high_quality_docs = []
+                moderate_quality_docs = []
+                
+                for doc in retrieved_docs:
+                    score = doc.get('score', 0)
+                    if score >= Config.SIMILARITY_THRESHOLD:
+                        high_quality_docs.append(doc)
+                    elif score >= Config.MIN_SIMILARITY_FOR_LEGAL_DOCS:
+                        moderate_quality_docs.append(doc)
+                
+                # Return high quality docs if available
+                if high_quality_docs:
+                    print(f"Retrieved {len(high_quality_docs)} high-quality documents")
+                    return high_quality_docs[:Config.TOP_K_RETRIEVAL]
+                
+                # Return moderate quality docs if no high quality ones
+                elif moderate_quality_docs:
+                    print(f"Retrieved {len(moderate_quality_docs)} moderate-quality documents")
+                    return moderate_quality_docs[:Config.TOP_K_RETRIEVAL]
+                
+                else:
+                    print("No documents found with sufficient similarity scores")
+                    return []
+            
+            # Apply reranking if enabled
+            if use_reranking and self.reranker and retrieved_docs:
+                print(f"Applying reranking to {len(retrieved_docs)} documents...")
+                
+                if Config.USE_SCORE_FUSION:
+                    # Use score fusion for better results
+                    retrieved_docs = self.reranker.rerank_with_fusion(
+                        query, 
+                        retrieved_docs, 
+                        alpha=Config.RERANKER_FUSION_ALPHA,
+                        top_k=Config.RERANKER_TOP_K
+                    )
+                else:
+                    # Use pure reranker scores
+                    retrieved_docs = self.reranker.rerank_documents(
+                        query, 
+                        retrieved_docs, 
+                        top_k=Config.RERANKER_TOP_K
+                    )
+                print(f"Reranking completed, returning {len(retrieved_docs)} documents")
+            else:
+                # No reranking, limit results to original config
+                retrieved_docs = retrieved_docs[:Config.TOP_K_RETRIEVAL]
+            
+            print(f"Retrieved {len(retrieved_docs)} documents for query")
+            return retrieved_docs
+            
+        except Exception as e:
+            print(f"Error retrieving documents: {e}")
+            return []
+    
+    def format_context(self, documents: List[Dict[str, Any]]) -> str:
+        """Format retrieved documents as context for LLM"""
+        if not documents:
+            return "Không có tài liệu pháp luật liên quan được tìm thấy."
+        
+        context_parts = []
+        
+        for i, doc in enumerate(documents, 1):
+            title = doc.get('title', 'Không có tiêu đề')
+            content = doc.get('content', '')
+            doc_id = doc.get('id', '')
+            metadata = doc.get('metadata', {})
+            law_id = metadata.get('law_id', '')
+            article_id = metadata.get('article_id', '')
+            
+            # Limit content length
+            if len(content) > 500:
+                content = content[:500] + "..."
+            
+            # Format law and article information
+            law_info = f"Luật: {law_id}" if law_id else ""
+            article_info = f"Điều {article_id}" if article_id else f"ID: {doc_id}"
+            
+            context_part = f"""
+Tài liệu {i}:
+{law_info}
+{article_info}: {title}
+Nội dung: {content}
+"""
+            context_parts.append(context_part.strip())
+        
+        return "\n\n".join(context_parts)
+    
+    def generate_answer(self, query: str, context: str, is_fallback: bool = False) -> str:
+        """Generate answer using LLM with context"""
+        if not self.llm:
+            return "Lỗi: Không thể kết nối với mô hình ngôn ngữ."
+        
+        try:
+            # Create prompt based on scenario
+            if is_fallback:
+                prompt = Config.FALLBACK_SYSTEM_PROMPT.format(
+                    context=context,
+                    question=query
+                )
+            else:
+                prompt = Config.SYSTEM_PROMPT.format(
+                    context=context,
+                    question=query
+                )
+            
+            # Generate response
+            messages = [HumanMessage(content=prompt)]
+            response = self.llm(messages)
+            
+            answer = response.content
+            
+            # Add legal disclaimer if enabled and it's a fallback response
+            if Config.SHOW_LEGAL_DISCLAIMER and is_fallback:
+                disclaimer = "\n\nLưu ý quan trọng: Để đảm bảo quyền lợi của mình, người lao động nên tìm đến các chuyên gia pháp lý hoặc cơ quan chức năng có thẩm quyền để được tư vấn cụ thể và chính xác nhất dựa trên tình hình thực tế của mình."
+                answer += disclaimer
+            
+            return answer
+            
+        except Exception as e:
+            print(f"Error generating answer: {e}")
+            return f"Lỗi khi tạo câu trả lời: {str(e)}"
+    
+    def _is_negative_response(self, response: str) -> bool:
+        """Check if the response is a negative/unable to answer response"""
+        negative_indicators = [
+            "không thể trả lời",
+            "không tìm thấy",
+            "không có thông tin",
+            "xin lỗi",
+            "không thể tìm thấy",
+            "không có dữ liệu",
+            "không rõ",
+            "không biết",
+            "không đủ thông tin",
+            "thiếu thông tin",
+            "không có trong",
+            "ngoài phạm vi",
+            # Add the specific pattern mentioned by user
+            "không có đủ thông tin trong tài liệu tham khảo được cung cấp để trả lời trực tiếp câu hỏi này",
+            "cần tham khảo thêm các văn bản pháp luật khác",
+            "tìm kiếm thông tin chuyên sâu hơn về",
+            "tài liệu tham khảo không chứa thông tin đầy đủ"
+        ]
+        
+        response_lower = response.lower()
+        return any(indicator in response_lower for indicator in negative_indicators)
+    
+    def _enhanced_search_fallback(self, query: str) -> List[Dict[str, Any]]:
+        """Enhanced search fallback with multiple strategies"""
+        print("Attempting enhanced search fallback...")
+        
+        all_docs = []
+        
+        # Strategy 1: Lower similarity threshold
+        try:
+            if self.bm25_retriever and self.vector_store:
+                # BM25 with more results
+                bm25_results = self.bm25_retriever.get_relevant_documents(
+                    query, top_k=15  # Increased from default
+                )
+                
+                # Vector search with lower threshold
+                vector_results = self.vector_store.search_similar_documents(
+                    query, top_k=10  # Increased from default
+                )
+                
+                # Combine results with very low threshold
+                combined_docs = {}
+                
+                for doc in bm25_results:
+                    doc_id = doc.get('id', '')
+                    if doc_id and doc.get('score', 0) >= 0.1:  # Very low threshold
+                        combined_docs[doc_id] = {**doc, 'retrieval_method': 'bm25_enhanced'}
+                
+                for doc in vector_results:
+                    doc_id = doc.get('id', '')
+                    if doc_id and doc.get('score', 0) >= 0.15:  # Very low threshold
+                        if doc_id in combined_docs:
+                            combined_docs[doc_id]['retrieval_method'] = 'hybrid_enhanced'
+                        else:
+                            combined_docs[doc_id] = {**doc, 'retrieval_method': 'vector_enhanced'}
+                
+                all_docs = list(combined_docs.values())
+                
+        except Exception as e:
+            print(f"Error in enhanced search: {e}")
+        
+        # Strategy 2: Keyword-based search if still no results
+        if not all_docs:
+            try:
+                # Extract keywords from query
+                keywords = self.text_processor.extract_keywords(query)
+                if keywords:
+                    keyword_query = " ".join(keywords[:5])  # Use top 5 keywords
+                    
+                    if self.bm25_retriever:
+                        keyword_results = self.bm25_retriever.get_relevant_documents(
+                            keyword_query, top_k=10
+                        )
+                        all_docs.extend([{**doc, 'retrieval_method': 'keyword_fallback'} 
+                                       for doc in keyword_results if doc.get('score', 0) >= 0.05])
+                        
+            except Exception as e:
+                print(f"Error in keyword fallback: {e}")
+        
+        print(f"Enhanced search found {len(all_docs)} documents")
+        return all_docs[:Config.TOP_K_RETRIEVAL]
+    
+    def answer_question(self, query: str, use_fallback: bool = True, refine_question: bool = True) -> Dict[str, Any]:
+        """Answer a legal question using RAG with enhanced negative response handling and question refinement"""
+        print(f"Processing question: {query}")
+        
+        # Step 1: Refine the question if enabled
+        original_query = query
+        refinement_result = None
+        
+        if refine_question and Config.ENABLE_QUESTION_REFINEMENT and self.question_refiner:
+            print("🔧 Refining question for better search accuracy...")
+            refinement_result = self.question_refiner.refine_question(query, use_llm=Config.USE_LLM_FOR_REFINEMENT)
+            
+            if refinement_result["refined_question"] != query:
+                refined_query = refinement_result["refined_question"]
+                print(f"📝 Original: {query}")
+                print(f"✨ Refined: {refined_query}")
+                query = refined_query
+        
+        # Step 2: Retrieve relevant documents using refined query
+        retrieved_docs = self.retrieve_documents(query)
+        
+        # Check if we have relevant documents
+        if not retrieved_docs and Config.ENABLE_GOOGLE_SEARCH and use_fallback:
+            print("No relevant legal documents found, using Google search fallback")
+            
+            # Use Google search as fallback
+            search_results = self.google_search.search_legal_info(query)
+            
+            if search_results:
+                fallback_context = self.google_search.format_search_results(search_results)
+                
+                # Generate answer with fallback context
+                fallback_answer = self.generate_answer(query, fallback_context, True)
+                
+                return {
+                    'answer': fallback_answer,
+                    'retrieved_documents': [],
+                    'fallback_used': True,
+                    'search_results': search_results,
+                    'context': fallback_context,
+                    'search_results_html': self.google_search.format_search_results_for_display(search_results),
+                    'original_question': original_query,
+                    'refined_question': query,
+                    'question_refinement': refinement_result
+                }
+            else:
+                # Try enhanced search before giving up
+                enhanced_docs = self._enhanced_search_fallback(query)
+                if enhanced_docs:
+                    enhanced_context = self.format_context(enhanced_docs)
+                    enhanced_answer = self.generate_answer(query, enhanced_context, False)
+                    
+                    return {
+                        'answer': enhanced_answer,
+                        'retrieved_documents': enhanced_docs,
+                        'fallback_used': True,
+                        'context': enhanced_context,
+                        'search_results': [],
+                        'search_results_html': "",
+                        'enhanced_search_used': True,
+                        'original_question': original_query,
+                        'refined_question': query,
+                        'question_refinement': refinement_result
+                    }
+                else:
+                    return {
+                        'answer': "Xin lỗi, tôi không tìm thấy thông tin pháp luật liên quan đến câu hỏi của bạn trong cơ sở dữ liệu nội bộ và cũng không thể tìm kiếm thông tin trên web. Vui lòng thử lại với câu hỏi khác hoặc liên hệ với chuyên gia pháp lý.",
+                        'retrieved_documents': [],
+                        'fallback_used': True,
+                        'search_results': [],
+                        'context': "",
+                        'search_results_html': "",
+                        'enhanced_search_used': False,
+                        'original_question': original_query,
+                        'refined_question': query,
+                        'question_refinement': refinement_result
+                    }
+        elif not retrieved_docs:
+            # Try enhanced search before giving negative response
+            enhanced_docs = self._enhanced_search_fallback(query)
+            if enhanced_docs:
+                enhanced_context = self.format_context(enhanced_docs)
+                enhanced_answer = self.generate_answer(query, enhanced_context, False)
+                
+                return {
+                    'answer': enhanced_answer,
+                    'retrieved_documents': enhanced_docs,
+                    'fallback_used': False,
+                    'context': enhanced_context,
+                    'search_results': [],
+                    'search_results_html': "",
+                    'enhanced_search_used': True,
+                    'original_question': original_query,
+                    'refined_question': query,
+                    'question_refinement': refinement_result
+                }
+            else:
+                return {
+                    'answer': "Xin lỗi, tôi không tìm thấy thông tin pháp luật liên quan đến câu hỏi của bạn trong cơ sở dữ liệu.",
+                    'retrieved_documents': [],
+                    'fallback_used': False,
+                    'context': "",
+                    'search_results': [],
+                    'search_results_html': "",
+                    'enhanced_search_used': False,
+                    'original_question': original_query,
+                    'refined_question': query,
+                    'question_refinement': refinement_result
+                }
+        
+        # Format context
+        context = self.format_context(retrieved_docs)
+        
+        # Generate answer
+        answer = self.generate_answer(query, context, False)
+        
+        # Check if the generated answer is negative and retry with enhanced search
+        if self._is_negative_response(answer) and use_fallback:
+            print("🔍 Detected insufficient information response, activating search tools...")
+            
+            # Inform user that search is being performed
+            search_notification = f"\n\n*🔍 Đang tìm kiếm thông tin bổ sung để trả lời câu hỏi của bạn...*"
+            
+            # Try enhanced search first
+            enhanced_docs = self._enhanced_search_fallback(query)
+            
+            if enhanced_docs:
+                # Combine original and enhanced docs
+                all_docs = retrieved_docs + [doc for doc in enhanced_docs 
+                                           if doc.get('id') not in [d.get('id') for d in retrieved_docs]]
+                enhanced_context = self.format_context(all_docs[:Config.TOP_K_RETRIEVAL])
+                enhanced_answer = self.generate_answer(query, enhanced_context, False)
+                
+                # If still negative, try Google search
+                if self._is_negative_response(enhanced_answer) and Config.ENABLE_GOOGLE_SEARCH:
+                    print("📡 Enhanced search still insufficient, trying web search...")
+                    search_results = self.google_search.search_legal_info(query)
+                    
+                    if search_results:
+                        fallback_context = self.google_search.format_search_results(search_results)
+                        final_answer = self.generate_answer(query, fallback_context, True)
+                        
+                        # Add notification about web search usage (if enabled)
+                        if Config.SHOW_SOURCE_INFO:
+                            final_answer = f"{final_answer}\n\n*🌐 Thông tin này được tìm kiếm từ web do không tìm thấy đủ thông tin trong cơ sở dữ liệu pháp luật nội bộ.*"
+                        
+                        return {
+                            'answer': final_answer,
+                            'retrieved_documents': all_docs,
+                            'fallback_used': True,
+                            'context': enhanced_context,
+                            'search_results': search_results,
+                            'search_results_html': self.google_search.format_search_results_for_display(search_results),
+                            'enhanced_search_used': True,
+                            'search_triggered': True,
+                            'original_question': original_query,
+                            'refined_question': query,
+                            'question_refinement': refinement_result
+                        }
+                    else:
+                        # Web search failed, return enhanced answer with notification (if enabled)
+                        if Config.SHOW_SOURCE_INFO:
+                            enhanced_answer = f"{enhanced_answer}\n\n*🔍 Đã sử dụng tìm kiếm nâng cao trong cơ sở dữ liệu.*"
+                        
+                        return {
+                            'answer': enhanced_answer,
+                            'retrieved_documents': all_docs,
+                            'fallback_used': False,
+                            'context': enhanced_context,
+                            'search_results': [],
+                            'search_results_html': "",
+                            'enhanced_search_used': True,
+                            'search_triggered': True,
+                            'original_question': original_query,
+                            'refined_question': query,
+                            'question_refinement': refinement_result
+                        }
+                else:
+                    # Enhanced search was sufficient (if enabled)
+                    if Config.SHOW_SOURCE_INFO:
+                        enhanced_answer = f"{enhanced_answer}\n\n*🔍 Đã sử dụng tìm kiếm nâng cao để tìm thông tin này.*"
+                    
+                    return {
+                        'answer': enhanced_answer,
+                        'retrieved_documents': all_docs,
+                        'fallback_used': False,
+                        'context': enhanced_context,
+                        'search_results': [],
+                        'search_results_html': "",
+                        'enhanced_search_used': True,
+                        'search_triggered': True,
+                        'original_question': original_query,
+                        'refined_question': query,
+                        'question_refinement': refinement_result
+                    }
+            elif Config.ENABLE_GOOGLE_SEARCH:
+                # Try Google search as last resort
+                print("📡 Database search failed, trying web search as last resort...")
+                search_results = self.google_search.search_legal_info(query)
+                
+                if search_results:
+                    fallback_context = self.google_search.format_search_results(search_results)
+                    final_answer = self.generate_answer(query, fallback_context, True)
+                    
+                    # Add notification about web search usage (if enabled)
+                    if Config.SHOW_SOURCE_INFO:
+                        final_answer = f"{final_answer}\n\n*🌐 Thông tin này được tìm kiếm từ web do không tìm thấy trong cơ sở dữ liệu pháp luật nội bộ.*"
+                    
+                    return {
+                        'answer': final_answer,
+                        'retrieved_documents': retrieved_docs,
+                        'fallback_used': True,
+                        'context': context,
+                        'search_results': search_results,
+                        'search_results_html': self.google_search.format_search_results_for_display(search_results),
+                        'enhanced_search_used': False,
+                        'search_triggered': True,
+                        'original_question': original_query,
+                        'refined_question': query,
+                        'question_refinement': refinement_result
+                    }
+        
+        return {
+            'answer': answer,
+            'retrieved_documents': retrieved_docs,
+            'fallback_used': False,
+            'context': context,
+            'search_results': [],
+            'search_results_html': "",
+            'enhanced_search_used': False,
+            'search_triggered': False,
+            'original_question': original_query,
+            'refined_question': query,
+            'question_refinement': refinement_result
+        }
+    
+    def get_system_status(self) -> Dict[str, Any]:
+        """Get status of RAG system components"""
+        status = {
+            'llm_available': self.llm is not None,
+            'vector_store_available': self.vector_store is not None,
+            'bm25_available': self.bm25_retriever is not None,
+            'reranker_available': self.reranker is not None and self.reranker.model is not None,
+            'reranking_enabled': Config.ENABLE_RERANKING,
+            'google_api_configured': bool(Config.GOOGLE_API_KEY),
+            'qdrant_configured': bool(Config.QDRANT_URL and Config.QDRANT_API_KEY)
+        }
+        
+        # Get collection info if available
+        if self.vector_store:
+            try:
+                status['vector_store_info'] = self.vector_store.get_collection_info()
+            except:
+                status['vector_store_info'] = {}
+        
+        # Get BM25 stats if available
+        if self.bm25_retriever:
+            status['bm25_stats'] = self.bm25_retriever.get_index_stats()
+        
+        # Get reranker info if available
+        if self.reranker:
+            status['reranker_info'] = self.reranker.get_model_info()
+        
+        return status 

main/reranker.py

@@ -0,0 +1,199 @@
+from typing import List, Dict, Any, Tuple
+import numpy as np
+from sentence_transformers import CrossEncoder
+import logging
+
+class DocumentReranker:
+    """Document reranker using cross-encoder models for improved relevance scoring"""
+    
+    def __init__(self, model_name: str = "cross-encoder/ms-marco-MiniLM-L-6-v2"):
+        """
+        Initialize the reranker with a cross-encoder model
+        
+        Args:
+            model_name: Name of the cross-encoder model to use
+                       Default is a multilingual model that works well with Vietnamese
+        """
+        self.model_name = model_name
+        self.model = None
+        self._initialize_model()
+    
+    def _initialize_model(self):
+        """Initialize the cross-encoder model"""
+        try:
+            print(f"Loading reranker model: {self.model_name}")
+            self.model = CrossEncoder(self.model_name)
+            print("Reranker model loaded successfully")
+        except Exception as e:
+            print(f"Error loading reranker model: {e}")
+            # Fallback to a lighter model
+            try:
+                fallback_model = "cross-encoder/ms-marco-TinyBERT-L-2-v2"
+                print(f"Trying fallback model: {fallback_model}")
+                self.model = CrossEncoder(fallback_model)
+                self.model_name = fallback_model
+                print("Fallback reranker model loaded successfully")
+            except Exception as e2:
+                print(f"Error loading fallback model: {e2}")
+                self.model = None
+    
+    def rerank_documents(self, 
+                        query: str, 
+                        documents: List[Dict[str, Any]], 
+                        top_k: int = None) -> List[Dict[str, Any]]:
+        """
+        Rerank documents based on relevance to the query
+        
+        Args:
+            query: The search query
+            documents: List of documents to rerank
+            top_k: Number of top documents to return (None for all)
+            
+        Returns:
+            List of reranked documents with updated scores
+        """
+        if not self.model or not documents:
+            return documents
+        
+        try:
+            # Prepare query-document pairs for the cross-encoder
+            pairs = []
+            valid_docs = []
+            
+            for doc in documents:
+                content = doc.get('content', '')
+                title = doc.get('title', '')
+                
+                # Combine title and content for better matching
+                doc_text = f"{title}. {content}" if title else content
+                
+                # Truncate very long documents to avoid model limits
+                max_length = 512
+                if len(doc_text) > max_length:
+                    doc_text = doc_text[:max_length] + "..."
+                
+                pairs.append([query, doc_text])
+                valid_docs.append(doc)
+            
+            if not pairs:
+                return documents
+            
+            # Get relevance scores from cross-encoder
+            scores = self.model.predict(pairs)
+            
+            # Update documents with reranker scores
+            reranked_docs = []
+            for doc, score in zip(valid_docs, scores):
+                # Create a copy to avoid modifying original
+                reranked_doc = doc.copy()
+                
+                # Store both original and reranker scores
+                reranked_doc['reranker_score'] = float(score)
+                reranked_doc['original_score'] = doc.get('score', 0.0)
+                
+                # Update the main score with reranker score
+                reranked_doc['score'] = float(score)
+                
+                # Add reranking method info
+                if 'retrieval_method' in reranked_doc:
+                    reranked_doc['retrieval_method'] += '_reranked'
+                else:
+                    reranked_doc['retrieval_method'] = 'reranked'
+                
+                reranked_docs.append(reranked_doc)
+            
+            # Sort by reranker score (descending)
+            reranked_docs.sort(key=lambda x: x['reranker_score'], reverse=True)
+            
+            # Return top_k if specified
+            if top_k:
+                reranked_docs = reranked_docs[:top_k]
+            
+            print(f"Reranked {len(reranked_docs)} documents")
+            
+            # Log top scores for debugging
+            if reranked_docs:
+                top_scores = [doc['reranker_score'] for doc in reranked_docs[:3]]
+                print(f"Top reranker scores: {top_scores}")
+            
+            return reranked_docs
+            
+        except Exception as e:
+            print(f"Error during reranking: {e}")
+            # Return original documents if reranking fails
+            return documents
+    
+    def rerank_with_fusion(self, 
+                          query: str, 
+                          documents: List[Dict[str, Any]], 
+                          alpha: float = 0.7,
+                          top_k: int = None) -> List[Dict[str, Any]]:
+        """
+        Rerank documents using score fusion between original and reranker scores
+        
+        Args:
+            query: The search query
+            documents: List of documents to rerank
+            alpha: Weight for reranker score (0-1), higher means more weight on reranker
+            top_k: Number of top documents to return
+            
+        Returns:
+            List of reranked documents with fused scores
+        """
+        if not self.model or not documents:
+            return documents
+        
+        try:
+            # First get reranker scores
+            reranked_docs = self.rerank_documents(query, documents, top_k=None)
+            
+            if not reranked_docs:
+                return documents
+            
+            # Normalize original scores to 0-1 range
+            original_scores = [doc.get('original_score', 0.0) for doc in reranked_docs]
+            if max(original_scores) > 0:
+                original_scores_norm = [s / max(original_scores) for s in original_scores]
+            else:
+                original_scores_norm = [0.0] * len(original_scores)
+            
+            # Normalize reranker scores to 0-1 range
+            reranker_scores = [doc.get('reranker_score', 0.0) for doc in reranked_docs]
+            min_reranker = min(reranker_scores)
+            max_reranker = max(reranker_scores)
+            
+            if max_reranker > min_reranker:
+                reranker_scores_norm = [(s - min_reranker) / (max_reranker - min_reranker) 
+                                       for s in reranker_scores]
+            else:
+                reranker_scores_norm = [0.5] * len(reranker_scores)
+            
+            # Compute fused scores
+            for i, doc in enumerate(reranked_docs):
+                fused_score = (alpha * reranker_scores_norm[i] + 
+                              (1 - alpha) * original_scores_norm[i])
+                doc['fused_score'] = fused_score
+                doc['score'] = fused_score  # Update main score
+            
+            # Sort by fused score
+            reranked_docs.sort(key=lambda x: x['fused_score'], reverse=True)
+            
+            # Return top_k if specified
+            if top_k:
+                reranked_docs = reranked_docs[:top_k]
+            
+            print(f"Score fusion reranking completed for {len(reranked_docs)} documents")
+            
+            return reranked_docs
+            
+        except Exception as e:
+            print(f"Error during fusion reranking: {e}")
+            return documents
+    
+    def get_model_info(self) -> Dict[str, Any]:
+        """Get information about the loaded model"""
+        return {
+            'model_name': self.model_name,
+            'model_loaded': self.model is not None,
+            'model_type': 'cross-encoder'
+        } 

main/vector_store.py

@@ -0,0 +1,244 @@
+import os
+import uuid
+from typing import List, Dict, Any, Optional
+from qdrant_client import QdrantClient
+from qdrant_client.models import (
+    Distance,
+    VectorParams,
+    PointStruct,
+    Filter,
+    FieldCondition,
+    MatchValue,
+)
+from sentence_transformers import SentenceTransformer
+import numpy as np
+from config import Config
+from tqdm import tqdm
+
+
+class QdrantVectorStore:
+    """QDrant vector store for legal document embeddings"""
+
+    def __init__(self):
+        self.client = None
+        self.embedding_model = None
+        self.collection_name = Config.COLLECTION_NAME
+        self._initialize_client()
+        self._initialize_embedding_model()
+
+    def _initialize_client(self):
+        """Initialize QDrant client"""
+        try:
+            if Config.QDRANT_URL and Config.QDRANT_API_KEY:
+                # Cloud QDrant
+                self.client = QdrantClient(
+                    url=Config.QDRANT_URL, api_key=Config.QDRANT_API_KEY
+                )
+            else:
+                # Local QDrant (fallback)
+                self.client = QdrantClient(host="localhost", port=6333)
+
+            print("QDrant client initialized successfully")
+        except Exception as e:
+            print(f"Error initializing QDrant client: {e}")
+            raise
+
+    def _initialize_embedding_model(self):
+        """Initialize embedding model"""
+        try:
+            # Clear any potentially corrupted cache
+            import tempfile
+            import shutil
+
+            cache_dir = os.path.join(tempfile.gettempdir(), "sentence_transformers")
+
+            self.embedding_model = SentenceTransformer(Config.EMBEDDING_MODEL)
+            print(f"Embedding model {Config.EMBEDDING_MODEL} loaded successfully")
+        except UnicodeDecodeError as e:
+            print(f"Encoding error loading embedding model: {e}")
+            print("Trying to clear sentence-transformers cache...")
+            try:
+                import tempfile
+                import shutil
+
+                cache_dir = os.path.join(tempfile.gettempdir(), "sentence_transformers")
+                if os.path.exists(cache_dir):
+                    shutil.rmtree(cache_dir)
+                    print("Cache cleared, retrying...")
+
+                self.embedding_model = SentenceTransformer(Config.EMBEDDING_MODEL)
+                print(
+                    f"Embedding model {Config.EMBEDDING_MODEL} loaded successfully after cache clear"
+                )
+            except Exception as retry_e:
+                print(
+                    f"Failed to load embedding model even after cache clear: {retry_e}"
+                )
+                raise
+        except Exception as e:
+            print(f"Error loading embedding model: {e}")
+            raise
+
+    def create_collection(self, vector_size: int = 384, force_recreate: bool = False):
+        """Create collection in QDrant"""
+        try:
+            # Check if collection exists
+            collections = self.client.get_collections().collections
+            collection_exists = any(
+                col.name == self.collection_name for col in collections
+            )
+            print(f"Collection exists: {collection_exists}")
+
+            if collection_exists:
+                if force_recreate:
+                    print(f"Force recreating collection: {self.collection_name}")
+                    self.client.delete_collection(self.collection_name)
+                    print(f"Deleted existing collection: {self.collection_name}")
+                    # Create collection
+                    self.client.create_collection(
+                        collection_name=self.collection_name,
+                        vectors_config=VectorParams(
+                            size=vector_size, distance=Distance.COSINE
+                        ),
+                    )
+                    print(f"Successfully created collection: {self.collection_name}")
+                else:
+                    print(
+                        f"Collection {self.collection_name} already exists - skipping creation"
+                    )
+                    return
+            else:
+                print(
+                    f"Collection {self.collection_name} does not exist - creating new collection"
+                )
+                # Create collection
+                self.client.create_collection(
+                    collection_name=self.collection_name,
+                    vectors_config=VectorParams(
+                        size=vector_size, distance=Distance.COSINE
+                    ),
+                )
+                print(f"Successfully created collection: {self.collection_name}")
+
+        except Exception as e:
+            print(f"Error creating collection: {e}")
+            raise
+
+    def embed_text(self, text: str) -> List[float]:
+        """Generate embedding for text"""
+        try:
+            embedding = self.embedding_model.encode(text, convert_to_tensor=False)
+            return embedding.tolist()
+        except Exception as e:
+            print(f"Error generating embedding: {e}")
+            return []
+
+    def add_documents(self, documents: List[Dict[str, Any]]):
+        """Add documents to vector store"""
+        try:
+            points = []
+
+            for doc in tqdm(documents):
+                # Generate embedding
+                content = doc.get("content", "")
+                if not content:
+                    continue
+
+                embedding = self.embed_text(content)
+                if not embedding:
+                    continue
+
+                # Create point
+                point = PointStruct(
+                    id=str(uuid.uuid4()),
+                    vector=embedding,
+                    payload={
+                        "article_id": doc.get("id", ""),
+                        "title": doc.get("title", ""),
+                        "content": content,
+                        "metadata": doc.get("metadata", {}),
+                    },
+                )
+                points.append(point)
+
+            # Batch upload
+            batch_size = 100
+            for i in range(0, len(points), batch_size):
+                batch = points[i : i + batch_size]
+                self.client.upsert(collection_name=self.collection_name, points=batch)
+                print(
+                    f"Uploaded batch {i//batch_size + 1}/{(len(points) + batch_size - 1)//batch_size}"
+                )
+
+            print(f"Successfully added {len(points)} documents to vector store")
+
+        except Exception as e:
+            print(f"Error adding documents: {e}")
+            raise
+
+    def search_similar_documents(
+        self, query: str, top_k: int = None, score_threshold: float = None
+    ) -> List[Dict[str, Any]]:
+        """Search for similar documents"""
+        if top_k is None:
+            top_k = Config.TOP_K_RETRIEVAL
+        if score_threshold is None:
+            score_threshold = Config.SIMILARITY_THRESHOLD
+
+        try:
+            # Generate query embedding
+            query_embedding = self.embed_text(query)
+            if not query_embedding:
+                return []
+
+            # Search
+            search_results = self.client.search(
+                collection_name=self.collection_name,
+                query_vector=query_embedding,
+                limit=top_k,
+                score_threshold=score_threshold,
+            )
+
+            # Format results
+            results = []
+            for result in search_results:
+                results.append(
+                    {
+                        "id": result.payload.get("article_id", ""),
+                        "title": result.payload.get("title", ""),
+                        "content": result.payload.get("content", ""),
+                        "score": result.score,
+                        "metadata": result.payload.get("metadata", {}),
+                    }
+                )
+
+            print(f"Found {len(results)} similar documents")
+            return results
+
+        except Exception as e:
+            print(f"Error searching documents: {e}")
+            return []
+
+    def get_collection_info(self) -> Dict[str, Any]:
+        """Get collection information"""
+        try:
+            info = self.client.get_collection(self.collection_name)
+            result = {
+                "name": self.collection_name,  # Use the collection name we know
+                "vectors_count": info.vectors_count,
+                "indexed_vectors_count": info.indexed_vectors_count,
+                "points_count": info.points_count,
+            }
+            print(f"Collection info: {result}")
+            return result
+        except Exception as e:
+            print(f"Collection '{self.collection_name}' does not exist: {e}")
+            return {}
+
+    def delete_collection(self):
+        """Delete collection"""
+        try:
+            self.client.delete_collection(self.collection_name)
+            print(f"Deleted collection: {self.collection_name}")
+        except Exception as e:
+            print(f"Error deleting collection: {e}")