#!/usr/bin/env python3 """ Vietnamese Legal Chatbot - Setup Script This script initializes the RAG system and processes the legal documents. """ import os import sys from typing import Dict, Any from utils.data_loader import LegalDataLoader from main.chatbot import VietnameseLegalRAG from config import Config def check_environment(): """Check if all required environment variables are set""" print("Checking environment configuration...") missing_vars = [] if not Config.GOOGLE_API_KEY: missing_vars.append("GOOGLE_API_KEY") if not Config.QDRANT_URL: missing_vars.append("QDRANT_URL") if not Config.QDRANT_API_KEY: missing_vars.append("QDRANT_API_KEY") if missing_vars: print(f"❌ Missing environment variables: {', '.join(missing_vars)}") print("\nPlease set the following environment variables:") for var in missing_vars: print(f" export {var}=your_value_here") print("\nOr create a .env file with these variables.") return False print("✅ Environment configuration OK") return True def check_data_files(): """Check if required data files exist""" print("Checking data files...") required_files = [ Config.CORPUS_PATH, Config.STOPWORDS_PATH, ] missing_files = [] for file_path in required_files: if not os.path.exists(file_path): missing_files.append(file_path) if missing_files: print(f"❌ Missing data files: {', '.join(missing_files)}") return False print("✅ Data files OK") return True def setup_rag_system(force_rebuild: bool = False): """Setup the RAG system with indices""" print("Setting up RAG system...") try: # Initialize data loader print("Initializing data loader...") data_loader = LegalDataLoader() # Load legal documents print("Loading legal corpus...") legal_docs = data_loader.load_legal_corpus() if not legal_docs: print("❌ No legal documents loaded") return None # Prepare documents for indexing print("Preparing documents for indexing...") documents = data_loader.prepare_documents_for_indexing() if not documents: print("❌ No documents prepared for indexing") return None print(f"📚 Prepared {len(documents)} documents for indexing") # Initialize RAG system with component-by-component error handling print("Initializing RAG system components...") try: print(" - Initializing text processor...") from utils.text_processor import VietnameseTextProcessor text_processor = VietnameseTextProcessor() print(" ✅ Text processor initialized") except Exception as e: print(f" ❌ Text processor error: {e}") raise try: print(" - Initializing vector store...") from main.vector_store import QdrantVectorStore vector_store = QdrantVectorStore() print(" ✅ Vector store initialized") except Exception as e: print(f" ❌ Vector store error: {e}") raise try: print(" - Initializing BM25 retriever...") from main.bm25_retriever import BM25Retriever bm25_retriever = BM25Retriever() print(" ✅ BM25 retriever initialized") except Exception as e: print(f" ❌ BM25 retriever error: {e}") raise try: print(" - Initializing complete RAG system...") rag_system = VietnameseLegalRAG() print(" ✅ RAG system initialized") except Exception as e: print(f" ❌ RAG system initialization error: {e}") raise # Setup indices print("Building indices (this may take a while)...") rag_system.setup_indices(documents, force_rebuild=force_rebuild) print("✅ RAG system setup completed") return rag_system except UnicodeDecodeError as e: print(f"❌ Encoding error setting up RAG system: {e}") print("💡 Try running: python cleanup.py") print("💡 Then run setup again: python setup_system.py") return None except Exception as e: print(f"❌ Error setting up RAG system: {e}") print("💡 For encoding issues, try: python cleanup.py") import traceback print("Full error traceback:") traceback.print_exc() return None def test_system(rag_system): """Test the RAG system with sample questions""" print("\nTesting RAG system...") test_questions = [ "Quyền và nghĩa vụ của người lao động là gì?", "Thời gian làm việc theo quy định của pháp luật?", "Điều kiện kết hôn theo luật hôn nhân và gia đình?" ] for i, question in enumerate(test_questions, 1): print(f"\n--- Test {i}: {question} ---") try: result = rag_system.answer_question(question, use_fallback=False) print(f"Answer: {result['answer'][:200]}...") print(f"Retrieved docs: {len(result['retrieved_documents'])}") print(f"Fallback used: {result['fallback_used']}") except Exception as e: print(f"Error answering question: {e}") def display_system_status(rag_system): """Display system status and statistics""" print("\n" + "="*50) print("SYSTEM STATUS") print("="*50) status = rag_system.get_system_status() print(f"🤖 LLM Available: {'✅' if status['llm_available'] else '❌'}") print(f"🔍 Vector Store: {'✅' if status['vector_store_available'] else '❌'}") print(f"📊 BM25 Retriever: {'✅' if status['bm25_available'] else '❌'}") print(f"🔑 Google API: {'✅' if status['google_api_configured'] else '❌'}") print(f"☁️ QDrant Cloud: {'✅' if status['qdrant_configured'] else '❌'}") # Vector store info if 'vector_store_info' in status and status['vector_store_info']: info = status['vector_store_info'] print(f"\n📚 Vector Store Info:") print(f" - Collection: {info.get('name', 'N/A')}") print(f" - Documents: {info.get('points_count', 0):,}") vectors_count = info.get('vectors_count') print(f" - Vectors: {vectors_count if vectors_count is not None else 0:,}") # BM25 stats if 'bm25_stats' in status and status['bm25_stats']: stats = status['bm25_stats'] print(f"\n📊 BM25 Index Stats:") print(f" - Documents: {stats.get('total_documents', 0):,}") print(f" - Vocabulary: {stats.get('vocabulary_size', 0):,}") print(f" - Avg Doc Length: {stats.get('average_document_length', 0):.1f}") def main(): """Main setup function""" print("🏛️ Vietnamese Legal Chatbot - Setup") print("="*50) # Check prerequisites if not check_environment(): print("\n❌ Environment check failed. Please configure your environment variables.") sys.exit(1) if not check_data_files(): print("\n❌ Data file check failed. Please ensure all data files are present.") sys.exit(1) # Parse command line arguments force_rebuild = "--rebuild" in sys.argv or "-r" in sys.argv run_tests = "--test" in sys.argv or "-t" in sys.argv if force_rebuild: print("\n🔄 Force rebuild mode enabled") # Setup RAG system rag_system = setup_rag_system(force_rebuild=force_rebuild) if not rag_system: print("\n❌ RAG system setup failed") sys.exit(1) # Display system status display_system_status(rag_system) # Run tests if requested if run_tests: test_system(rag_system) print("\n✅ Setup completed successfully!") print("\nYou can now run the Streamlit app:") print(" streamlit run app.py") print("\nUsage:") print(" python setup_system.py # Normal setup") print(" python setup_system.py --rebuild # Force rebuild indices") print(" python setup_system.py --test # Run with tests") if __name__ == "__main__": main()