Sharada25/dhammaai
1
1"""lightrag_wrapper.py2 3Provides a resilient RAG interface for the Vipassana agent.4Now includes:5- Cross-Encoder Reranking for improved precision and accuracy6- BM25 keyword-based search for hybrid retrieval (BM25 + FAISS)7"""8import os9import json10import pickle11import faiss12import numpy as np13import openai14import re15import unicodedata16from typing import List, Dict, Optional, Tuple, Any17 18# Conditional imports for RAG components19try:20 from sentence_transformers import SentenceTransformer, CrossEncoder21 from sklearn.preprocessing import normalize22except ImportError:23 SentenceTransformer = None24 CrossEncoder = None25 normalize = None26 27try:28 from openai import OpenAI29except ImportError:30 OpenAI = None31 32try:33 from huggingface_hub import InferenceClient34except ImportError:35 InferenceClient = None36 37try:38 from rank_bm25 import BM25Okapi39except ImportError:40 BM25Okapi = None41 42try:43 from multilingual_prompts import (44 detect_language, get_system_prompt, get_user_prompt,45 get_english_query_hints, add_translation_context,46 VIPASSANA_TERMS_MULTILINGUAL47 )48except ImportError:49 print("Warning: multilingual_prompts module not found. Fallback to basic prompts.")50 detect_language = None51 get_system_prompt = None52 get_user_prompt = None53 get_english_query_hints = None54 add_translation_context = None55 VIPASSANA_TERMS_MULTILINGUAL = None56 57# ============================================================================58# PLUG-AND-PLAY LLM CONFIGURATION59# ============================================================================60# Switch between OpenAI and HuggingFace easily:61# - Set LLM_PROVIDER to "openai" to use OpenAI GPT models (requires OPENAI_API_KEY)62# * Cost: ~$0.002-0.01 per chat (very cheap)63# * Quality: Excellent (GPT-3.5-turbo)64# * Speed: Fast65# - Set LLM_PROVIDER to "huggingface" to use HF Inference API (requires HF_API_TOKEN)66# * Cost: FREE67# * Quality: Good (Qwen2.5-7B-Instruct)68# * Speed: Moderate69#70# For HuggingFace Spaces with OpenAI: Set LLM_PROVIDER=openai + OPENAI_API_KEY71# For HuggingFace Spaces without OpenAI: Set LLM_PROVIDER=huggingface + HF_API_TOKEN72# For Local Development: Set to "openai" if you have OpenAI API key73# ============================================================================74 75LLM_PROVIDER = os.getenv("LLM_PROVIDER", "openai") # "openai" or "huggingface" - default to openai for better quality76 77# OpenAI Configuration78OPENAI_MODEL = os.getenv("OPENAI_MODEL", "gpt-3.5-turbo")79 80# HuggingFace Configuration81# Recommended models for best multilingual support (Hindi, Marathi, English):82# - "Qwen/Qwen2.5-7B-Instruct" (BEST multilingual quality - 18T tokens trained)83# - "mistralai/Mistral-7B-Instruct-v0.2" (good English, weak Hindi/Marathi)84# - "meta-llama/Meta-Llama-3.1-8B-Instruct" (decent multilingual)85# - "google/gemma-2-9b-it" (good multilingual, larger model)86HF_MODEL = os.getenv("HF_MODEL", "Qwen/Qwen2.5-7B-Instruct")87 88# --- Other Configuration ---89EMBED_MODEL = os.getenv("EMBED_MODEL", "sentence-transformers/all-MiniLM-L6-v2")90# Using a powerful, free, public Cross-Encoder for better relevance scoring91RERANK_MODEL = os.getenv("RERANK_MODEL", "cross-encoder/ms-marco-MiniLM-L-6-v2")92INDEX_PATH = os.getenv("VIPASSANA_INDEX_PATH", "data/vector_store/index.faiss")93META_PATH = os.getenv("VIPASSANA_META_PATH", "data/vector_store/meta.json")94BM25_PATH = os.getenv("VIPASSANA_BM25_PATH", "data/vector_store/bm25.pkl")95 96 97class VipassanaRAGAgent:98 """Unified RAG wrapper with Multi-Query Search, BM25, and Cross-Encoder Reranking."""99 100 def __init__(self, openai_api_key: Optional[str] = None, hf_api_token: Optional[str] = None):101 # --- CRITICAL FIX: Initialize status attributes ---102 self.is_ready = False103 self.index_loaded = False104 # --------------------------------------------------105 106 self.model = None # Sentence Transformer for embedding107 self.reranker = None # Cross Encoder for re-scoring108 self.bm25 = None # BM25 for keyword search109 self.index = None110 self.metadatas = []111 112 # LLM clients (only one will be used based on LLM_PROVIDER)113 self.openai_client = None114 self.hf_client = None115 self.llm_provider = LLM_PROVIDER116 117 # 1. Check for required packages118 if SentenceTransformer is None or CrossEncoder is None or faiss is None or np is None:119 print("ERROR: Missing required packages (sentence-transformers, faiss, numpy). Cannot initialize RAG components.")120 return121 122 # 2. Setup LLM Client based on LLM_PROVIDER123 print(f"[LLM] Initializing LLM provider: {self.llm_provider}")124 125 if self.llm_provider == "openai":126 # ========== OPENAI SETUP ==========127 if OpenAI is None:128 print("ERROR: openai package not installed. Run: pip install openai")129 return130 131 try:132 # Use provided key or environment key133 api_key = openai_api_key or os.getenv("OPENAI_API_KEY")134 if not api_key:135 print("ERROR: OPENAI_API_KEY not provided. Cannot initialize OpenAI client.")136 print("HINT: Set OPENAI_API_KEY in your .env file or HF Spaces secrets")137 return138 139 # Initialize with longer timeout for HF Spaces network140 self.openai_client = OpenAI(141 api_key=api_key,142 timeout=60.0, # 60 second timeout143 max_retries=3 # Retry up to 3 times on connection errors144 )145 print(f"[LLM] OpenAI client initialized successfully (model: {OPENAI_MODEL})")146 147 except Exception as e:148 print(f"ERROR: Failed to initialize OpenAI client: {e}")149 return150 151 elif self.llm_provider == "huggingface":152 # ========== HUGGINGFACE SETUP ==========153 if InferenceClient is None:154 print("ERROR: huggingface_hub package not installed. Run: pip install huggingface_hub")155 return156 157 try:158 # Use provided token or environment token (optional for public models)159 token = hf_api_token or os.getenv("HF_API_TOKEN") or os.getenv("HF_TOKEN")160 161 # Initialize HF Inference Client162 self.hf_client = InferenceClient(163 model=HF_MODEL,164 token=token,165 timeout=60.0166 )167 print(f"[LLM] HuggingFace Inference client initialized successfully")168 print(f"[LLM] Using model: {HF_MODEL}")169 if not token:170 print("[LLM] No HF_API_TOKEN provided - using public inference (may have rate limits)")171 172 except Exception as e:173 print(f"ERROR: Failed to initialize HuggingFace client: {e}")174 return175 else:176 print(f"ERROR: Invalid LLM_PROVIDER: {self.llm_provider}. Must be 'openai' or 'huggingface'")177 return 178 179 # 3. Load Models and Index180 try:181 print(f"Initializing embedding model: {EMBED_MODEL}")182 self.model = SentenceTransformer(EMBED_MODEL)183 184 print(f"Initializing reranking model: {RERANK_MODEL}")185 # Reranker model is a CrossEncoder, which expects (query, document) pairs186 self.reranker = CrossEncoder(RERANK_MODEL)187 188 self._load_index()189 190 except Exception as e:191 print(f"Error loading RAG components (models or index): {e}")192 return193 194 # 4. Set final readiness status195 if self.index_loaded:196 self.is_ready = True197 print("VipassanaRAGAgent initialized successfully with Reranking enabled.")198 199 200 def _load_index(self):201 """Loads the FAISS index, BM25 index, and metadata files."""202 print(f"Attempting to load index from {INDEX_PATH}...")203 try:204 self.index = faiss.read_index(INDEX_PATH)205 with open(META_PATH, "r", encoding="utf-8") as f:206 self.metadatas = json.load(f)207 208 # Load BM25 index if available209 if os.path.exists(BM25_PATH) and BM25Okapi is not None:210 with open(BM25_PATH, "rb") as f:211 self.bm25 = pickle.load(f)212 print(f"BM25 index loaded successfully from {BM25_PATH}")213 else:214 print(f"Warning: BM25 index not found at {BM25_PATH}. Falling back to FAISS only.")215 self.bm25 = None216 217 self.index_loaded = True218 print(f"Index loaded successfully with {len(self.metadatas)} chunks.")219 except Exception as e:220 print(f"Warning: Could not load index from {INDEX_PATH}. Error: {e}")221 self.index_loaded = False222 self.index = None223 self.metadatas = []224 self.bm25 = None225 226 def _expand_query(self, query: str, detected_language: str = 'english') -> List[str]:227 """228 Expand query into a list of query variations for better initial recall.229 ENHANCED: Now supports multilingual query expansion for Hindi/Marathi.230 """231 query_lower = query.lower()232 query_list = [query_lower]233 234 # Get English query hints for non-English queries (for better retrieval)235 if detected_language != 'english' and get_english_query_hints:236 english_hints = get_english_query_hints(query, detected_language)237 query_list.extend(english_hints)238 239 # Use multilingual terms map if available240 if VIPASSANA_TERMS_MULTILINGUAL:241 vipassana_terms_map = {k: v.get('en', []) for k, v in VIPASSANA_TERMS_MULTILINGUAL.items()}242 else:243 vipassana_terms_map = {244 'meditation': ['practice', 'technique', 'sadhana'],245 'vipassana': ['insight meditation', 'mindfulness', 'awareness'],246 'anapana': ['breathing', 'breath', 'respiration'],247 'goenka': ['s.n. goenka', 'teacher', 'acharya'],248 'dhamma': ['dharma', 'teaching', 'truth'],249 'suffering': ['dukkha', 'pain', 'misery', 'stress'],250 'anicca': ['impermanence', 'change'],251 }252 253 # 1. Keyword Substitution Queries254 for key, terms in vipassana_terms_map.items():255 if key in query_lower:256 for term in terms[:2]: # Limit to 2 synonyms per term257 query_list.append(query_lower.replace(key, term))258 259 # 2. Simple Rephrase Heuristic for English260 if "?" in query or query_lower.startswith(("what is", "how to", "why is")):261 simple_rephrase = re.sub(r"what is|how to|why is", "", query_lower, 1).strip().replace("?", "")262 if simple_rephrase:263 query_list.append(simple_rephrase)264 265 # 3. Hindi/Marathi question word removal for better matching266 hindi_question_words = ['क्या', 'कैसे', 'कहाँ', 'क्यों', 'कौन', 'कब']267 marathi_question_words = ['काय', 'कसे', 'कुठे', 'का', 'कोण', 'केव्हा']268 269 for word in hindi_question_words + marathi_question_words:270 if word in query:271 cleaned = query.replace(word, '').strip()272 if cleaned and len(cleaned) > 3:273 query_list.append(cleaned)274 275 # Clean and de-duplicate276 unique_queries = []277 seen_queries = set()278 for q in query_list:279 cleaned_q = re.sub(r"\s+", " ", q).strip()280 if cleaned_q and cleaned_q.lower() not in seen_queries and len(cleaned_q) > 2:281 unique_queries.append(cleaned_q)282 seen_queries.add(cleaned_q.lower())283 284 # Ensure a reasonable limit285 return unique_queries[:8]286 287 def _rerank_retrieved_items(self, query: str, retrieved_items: List[Dict[str, Any]]) -> List[Dict[str, Any]]:288 """289 Uses a Cross-Encoder model to re-score the relevance of each chunk to the query.290 This greatly improves precision.291 """292 if not self.reranker:293 return retrieved_items # Fallback to original score if reranker isn't loaded294 295 if not retrieved_items:296 return []297 298 # Prepare input pairs: [[query, chunk_text], [query, chunk_text], ...]299 sentences_to_rank = [[query, item["chunk"]] for item in retrieved_items]300 301 # Calculate new relevance scores (logits)302 new_scores = self.reranker.predict(sentences_to_rank)303 304 # Update the score and sort305 for i, item in enumerate(retrieved_items):306 # Cross-Encoder scores are high for high relevance307 item["rerank_score"] = float(new_scores[i]) 308 309 # Sort by the new, more accurate rerank_score310 retrieved_items.sort(key=lambda x: x["rerank_score"], reverse=True)311 312 return retrieved_items313 314 315 def _bm25_search(self, query: str, top_k: int = 15) -> List[Dict[str, Any]]:316 """317 Performs BM25 keyword-based search.318 Returns a list of candidate chunks with BM25 scores.319 """320 if self.bm25 is None:321 return []322 323 # Tokenize query324 tokenized_query = query.lower().split()325 326 # Get BM25 scores for all documents327 bm25_scores = self.bm25.get_scores(tokenized_query)328 329 # Get top-k indices330 top_indices = np.argsort(bm25_scores)[::-1][:top_k]331 332 results = []333 for idx in top_indices:334 if idx < len(self.metadatas) and bm25_scores[idx] > 0:335 metadata = self.metadatas[idx]336 results.append({337 "score": float(bm25_scores[idx]),338 "chunk": metadata.get("chunk", ""),339 "source": metadata.get("source"),340 "metadata": metadata341 })342 343 return results344 345 def retrieve(self, query: str, top_k: int = 5, top_k_initial: int = 30, detected_language: str = 'english') -> List[Dict[str, Any]]:346 """347 Retrieves context using hybrid search (BM25 + FAISS) followed by cross-encoder reranking.348 349 top_k: The final number of chunks passed to the LLM.350 top_k_initial: The number of chunks retrieved from FAISS before reranking.351 detected_language: Language of the query for better expansion.352 """353 if not self.index_loaded:354 return []355 356 # 1. Multi-Query Generation with language awareness357 query_list = self._expand_query(query, detected_language)358 359 # 2. Search FAISS index with all queries (High Recall)360 all_retrieved_items: Dict[str, Dict[str, Any]] = {} 361 362 # Distribute the initial search load across all queries (e.g., 30 chunks total)363 search_k_per_query = max(5, top_k_initial // len(query_list)) if len(query_list) > 0 else top_k_initial364 365 for single_query in query_list:366 if not single_query.strip():367 continue368 369 query_embedding = self.model.encode(single_query, convert_to_numpy=True)370 query_embedding = normalize(query_embedding.reshape(1, -1))371 372 # Search FAISS index for a large set of candidates373 distances, indices = self.index.search(query_embedding.astype("float32"), search_k_per_query)374 375 for i, score in zip(indices[0], distances[0]):376 if i >= 0:377 metadata = self.metadatas[i]378 chunk_text = metadata.get("chunk", "")379 380 if chunk_text:381 # De-duplicate chunks, keeping the best original FAISS score382 if chunk_text not in all_retrieved_items or score > all_retrieved_items[chunk_text]["score"]:383 all_retrieved_items[chunk_text] = {384 "score": float(score), # Original FAISS score (used for initial selection)385 "chunk": chunk_text,386 "source": metadata.get("source"),387 "metadata": metadata388 }389 390 # 2.5 BM25 Keyword Search (Hybrid Retrieval)391 bm25_results = self._bm25_search(query, top_k=15)392 for item in bm25_results:393 chunk_text = item["chunk"]394 if chunk_text and chunk_text not in all_retrieved_items:395 all_retrieved_items[chunk_text] = item396 397 initial_candidates = list(all_retrieved_items.values())398 399 # 3. Reranking (High Precision)400 if not initial_candidates:401 return []402 403 # Rerank all candidates to get the true relevance score404 reranked_items = self._rerank_retrieved_items(query, initial_candidates)405 406 # 4. Final Truncation: Return the very best chunks after reranking407 return reranked_items[:top_k]408 409 410 def _get_llm_response(self, query: str, context: str, mode: str = "long", detected_language: str = None) -> str:411 """412 Unified method to get LLM response from either OpenAI or HuggingFace.413 Automatically routes to the correct provider based on LLM_PROVIDER setting.414 """415 if not self.openai_client and not self.hf_client:416 return "Internal error: LLM client is not initialized."417 418 # Detect language and get appropriate prompts419 if detect_language and get_system_prompt and get_user_prompt:420 if detected_language is None:421 detected_language = detect_language(query)422 system_prompt = get_system_prompt(detected_language)423 424 # Add translation context for non-English queries425 enhanced_context = context426 if add_translation_context and detected_language != 'english':427 enhanced_context = add_translation_context(context, detected_language)428 429 prompt = get_user_prompt(query, enhanced_context, detected_language)430 print(f"[LLM] Detected language: {detected_language}")431 else:432 # Fallback to basic English prompts if multilingual module not available433 detected_language = "english"434 system_prompt = (435 "You are the Vipassana Guide AI, a compassionate meditation teacher. "436 "Use ONLY the provided CONTEXT. Do not add external knowledge.\n\n"437 "Format responses with ## headings, **bold**, and *italic* for Pali/Sanskrit terms.\n"438 "Be direct and practical. Include [Source: filename] references."439 )440 prompt = f"""CONTEXT (Vipassana Knowledge Base):441{context}442 443USER'S QUESTION:444{query}445 446Generate the response based strictly on the provided CONTEXT."""447 448 # ========== ROUTE TO CORRECT LLM PROVIDER ==========449 450 if self.llm_provider == "openai":451 # ========== OPENAI IMPLEMENTATION ==========452 try:453 response = self.openai_client.chat.completions.create(454 model=OPENAI_MODEL,455 messages=[456 {"role": "system", "content": system_prompt},457 {"role": "user", "content": prompt}458 ],459 temperature=0.05, # Low temperature ensures the model stays faithful to the context460 max_tokens=1200,461 top_p=0.9,462 timeout=60 # 60 second timeout for this specific request463 )464 return response.choices[0].message.content465 466 except openai.APIConnectionError as e:467 # Network connectivity issues468 error_msg = f"Network connection error to OpenAI API. This may be temporary. Please try again in a moment."469 print(f"[OpenAI API] Connection Error: {e}")470 return error_msg471 except openai.APITimeoutError as e:472 # Request timed out473 error_msg = f"Request to OpenAI timed out. The service may be slow. Please try again."474 print(f"[OpenAI API] Timeout Error: {e}")475 return error_msg476 except openai.AuthenticationError as e:477 # API key issue478 error_msg = f"OpenAI API authentication failed. Please check your API key configuration."479 print(f"[OpenAI API] Authentication Error: {e}")480 return error_msg481 except openai.RateLimitError as e:482 # Rate limit exceeded483 error_msg = f"OpenAI API rate limit exceeded. Please wait a moment and try again."484 print(f"[OpenAI API] Rate Limit Error: {e}")485 return error_msg486 except Exception as e:487 # Generic error handling488 error_msg = f"Error during OpenAI generation: {type(e).__name__} - {str(e)}"489 print(f"[OpenAI API] Generic Error: {e}")490 return error_msg491 492 elif self.llm_provider == "huggingface":493 # ========== HUGGINGFACE IMPLEMENTATION ==========494 try:495 # Use chat completion API for instruction-tuned models like Mistral496 # Format messages similar to OpenAI's chat format497 messages = [498 {"role": "system", "content": system_prompt},499 {"role": "user", "content": prompt}500 ]501 502 # Call HuggingFace Chat Completion API503 response = self.hf_client.chat_completion(504 messages=messages,505 max_tokens=1200,506 temperature=0.05, # Very low temperature for consistency (especially important for multilingual)507 top_p=0.85, # Reduced top_p for more focused output508 )509 510 # Extract the assistant's response511 answer = response.choices[0].message.content512 return answer.strip()513 514 except Exception as e:515 # HuggingFace error handling516 error_msg = f"Error during HuggingFace generation: {type(e).__name__} - {str(e)}"517 print(f"[HF API] Error: {e}")518 print(f"[HF API] Full error details: {e}")519 return error_msg520 521 else:522 return f"Error: Unknown LLM provider: {self.llm_provider}"523 524 def answer(self, query: str, top_k: int = 6, mode: str = "long") -> Tuple[str, List[str]]:525 """Main method to retrieve context and generate an answer with optimized performance."""526 if not self.is_ready:527 return "RAG Agent failed to initialize. Please check the console for errors.", []528 529 # Detect language early for use throughout the pipeline530 detected_language = 'english'531 if detect_language:532 detected_language = detect_language(query)533 print(f"[RAG] Query language detected: {detected_language}")534 535 # 0. Check if this is a course-related query and return direct information536 try:537 from course_extractor import get_course_extractor538 extractor = get_course_extractor()539 540 query_lower = query.lower()541 542 # Detect course-related keywords543 if any(keyword in query_lower for keyword in ['course', 'december', 'november', 'january', 'schedule', 'date', 'centre', 'center', 'when', 'booking', 'register']):544 # Extract month if present545 months = ['january', 'february', 'march', 'april', 'may', 'june', 'july', 'august', 'september', 'october', 'november', 'december']546 month_found = None547 centre_found = None548 549 for month in months:550 if month in query_lower:551 month_found = month552 break553 554 # Check for centre names555 centres = extractor.get_all_centres()556 for centre in centres:557 if centre.lower() in query_lower:558 centre_found = centre559 break560 561 # Try to get course information with language support562 if centre_found and month_found:563 course_info = extractor.get_course_info(centre=centre_found, month=month_found, language=detected_language)564 if course_info:565 return course_info, ["data/knowledge_base/dhamma_org_content.md"]566 elif centre_found:567 course_info = extractor.get_course_info(centre=centre_found, language=detected_language)568 if course_info:569 return course_info, ["data/knowledge_base/dhamma_org_content.md"]570 elif month_found:571 course_info = extractor.get_course_info(month=month_found, language=detected_language)572 if course_info:573 return course_info, ["data/knowledge_base/dhamma_org_content.md"]574 except Exception as e:575 print(f"Course extractor error (non-critical): {e}")576 pass # Fall back to RAG if course extraction fails577 578 # 1. Optimized retrieval with reranking for better performance579 retrieved_items = self.retrieve(query, top_k=top_k, top_k_initial=25, detected_language=detected_language)580 581 if not retrieved_items:582 return "Could not retrieve relevant documents from the knowledge base.", []583 584 # 2. Format context with better organization585 context_parts = []586 sources = []587 588 for i, item in enumerate(retrieved_items):589 chunk = item.get("chunk", "")590 source = item.get("source", "")591 592 if chunk:593 # Add source information to each chunk594 context_parts.append(f"[Source: {source}]\n{chunk}")595 596 if source and source not in sources:597 sources.append(source)598 599 context = "\n\n---\n\n".join(context_parts)600 601 # 3. Generate answer using enhanced context (works with both OpenAI and HuggingFace)602 answer = self._get_llm_response(query, context, mode, detected_language=detected_language)603 604 return answer, sources605 606 def _clean_text(self, text: str) -> str:607 """Basic cleaning for text extracted from PDFs."""608 if not isinstance(text, str):609 return ""610 611 # Standard cleaning logic remains (normalized, cleanup hyphens/spaces, etc.)612 cleaned = unicodedata.normalize('NFKC', text)613 cleaned = re.sub(r"(/c\d+)+", " ", cleaned)614 cleaned = re.sub(r"\[\s*\d+\s*\]", " ", cleaned)615 cleaned = re.sub(r"\(\s*\d+\s*\)", " ", cleaned)616 cleaned = re.sub(r"-\s*\n\s*", "", cleaned)617 cleaned = re.sub(r"\n+", " ", cleaned)618 cleaned = re.sub(r"[\x00-\x1F\x7F-\x9F]", " ", cleaned)619 cleaned = re.sub(r"\s+", " ", cleaned).strip()620 621 return cleaned622 623# NOTE: The initialization logic in app.py will automatically pick up the new 624# RAG agent, but you must ensure that all required dependencies 625# (sentence-transformers[for CrossEncoder], numpy, faiss, openai) are installed 626# in the environment where this code runs.627