alschameri/helping-source
0
1"""2Retrieval-Augmented Generation (RAG) system for Arabic travel agency chatbot.3Handles document indexing, vector search, and context retrieval.4"""5 6import os7import json8import logging9import pickle10import time11import re12from typing import List, Dict, Optional13import numpy as np14from pathlib import Path15 16try:17 import faiss18 FAISS_AVAILABLE = True19except ImportError:20 FAISS_AVAILABLE = False21 22logger = logging.getLogger(__name__)23 24class RAGSystem:25 """RAG system for document indexing and retrieval."""26 27 def __init__(self, gemini_client, data_dir='data', storage_dir='storage', chunk_size=500, overlap=50):28 self.gemini_client = gemini_client29 self.data_dir = Path(data_dir)30 self.storage_dir = Path(storage_dir)31 self.chunk_size = chunk_size32 self.overlap = overlap33 34 # Ensure directories exist35 self.data_dir.mkdir(exist_ok=True)36 self.storage_dir.mkdir(exist_ok=True)37 38 # Initialize index and metadata39 self.index = None40 self.metadata = []41 self.embedding_dim = 768 # Default Gemini embedding dimension42 43 # Conversation memory storage (session_id -> conversation_history)44 self.conversations = {}45 46 # Pattern mappings for incomplete messages47 self.message_patterns = {48 # Price inquiries49 'price_patterns': [50 r'بكم\s*سعر|كم\s*سعر|بكم\s*ثمن|كم\s*ثمن|بكم\s*تكلفة|كم\s*تكلفة',51 r'بكم|كم\s*(.*?)\s*\?*',52 r'السعر|الثمن|التكلفة'53 ],54 # Requirements inquiries 55 'requirements_patterns': [56 r'ماهي\s*متطلبات|ما\s*هي\s*متطلبات|متطلبات',57 r'ماهي\s*الشروط|ما\s*هي\s*الشروط|الشروط',58 r'ماذا\s*احتاج|ما\s*احتاج|احتاج',59 r'الوثائق|المستندات|الأوراق'60 ],61 # Offer/details inquiries62 'details_patterns': [63 r'ماهي\s*عروض|ما\s*هي\s*عروض|عروض',64 r'ماهي\s*تفاصيل|ما\s*هي\s*تفاصيل|تفاصيل',65 r'معلومات\s*عن|معلومات',66 r'اخبرني\s*عن|قل\s*لي\s*عن'67 ],68 # Booking inquiries69 'booking_patterns': [70 r'كيف\s*احجز|كيف\s*الحجز|احجز|حجز',71 r'كيف\s*اسجل|التسجيل|سجل'72 ]73 }74 75 def chunk_text(self, text: str, chunk_size: int = None, overlap: int = None) -> List[Dict]:76 """Split text into overlapping chunks."""77 if chunk_size is None:78 chunk_size = self.chunk_size79 if overlap is None:80 overlap = self.overlap81 82 # Simple word-based chunking83 words = text.split()84 chunks = []85 86 for i in range(0, len(words), chunk_size - overlap):87 chunk_words = words[i:i + chunk_size]88 chunk_text = ' '.join(chunk_words)89 90 chunks.append({91 'text': chunk_text,92 'start_word': i,93 'end_word': min(i + chunk_size, len(words)),94 'word_count': len(chunk_words)95 })96 97 # Break if we've reached the end98 if i + chunk_size >= len(words):99 break100 101 return chunks102 103 def load_documents(self) -> List[Dict]:104 """Load and chunk all text documents from data directory."""105 documents = []106 107 logger.info(f"Loading documents from {self.data_dir}")108 109 for file_path in self.data_dir.glob('*.txt'):110 try:111 with open(file_path, 'r', encoding='utf-8') as f:112 content = f.read().strip()113 114 if not content:115 logger.warning(f"Empty file: {file_path}")116 continue117 118 # Chunk the document119 chunks = self.chunk_text(content)120 121 for i, chunk in enumerate(chunks):122 documents.append({123 'source_file': file_path.name,124 'chunk_id': i,125 'text': chunk['text'],126 'start_word': chunk['start_word'],127 'end_word': chunk['end_word'],128 'word_count': chunk['word_count']129 })130 131 logger.info(f"Loaded {len(chunks)} chunks from {file_path.name}")132 133 except Exception as e:134 logger.error(f"Error loading {file_path}: {e}")135 continue136 137 logger.info(f"Total documents loaded: {len(documents)}")138 return documents139 140 def create_embeddings(self, texts: List[str]) -> np.ndarray:141 """Create embeddings for list of texts using Gemini."""142 try:143 embeddings = self.gemini_client.embed_texts(texts)144 return np.array(embeddings)145 except Exception as e:146 logger.error(f"Error creating embeddings: {e}")147 # Fallback to random embeddings for development148 logger.warning("Using random embeddings as fallback")149 return np.random.rand(len(texts), self.embedding_dim).astype(np.float32)150 151 def build_index(self, embeddings: np.ndarray):152 """Build FAISS index from embeddings."""153 if not FAISS_AVAILABLE:154 logger.warning("FAISS not available, using simple similarity search")155 self.index = embeddings156 return157 158 # Create FAISS index159 dimension = embeddings.shape[1]160 self.index = faiss.IndexFlatIP(dimension) # Inner product for cosine similarity161 162 # Normalize embeddings for cosine similarity163 faiss.normalize_L2(embeddings)164 self.index.add(embeddings)165 166 logger.info(f"Built FAISS index with {self.index.ntotal} vectors")167 168 def save_index(self):169 """Save index and metadata to disk."""170 try:171 # Save metadata172 metadata_path = self.storage_dir / 'metadata.json'173 with open(metadata_path, 'w', encoding='utf-8') as f:174 json.dump(self.metadata, f, ensure_ascii=False, indent=2)175 176 # Save index177 if FAISS_AVAILABLE and hasattr(self.index, 'ntotal'):178 index_path = self.storage_dir / 'index.faiss'179 faiss.write_index(self.index, str(index_path))180 else:181 # Save numpy array as pickle182 index_path = self.storage_dir / 'index.pkl'183 with open(index_path, 'wb') as f:184 pickle.dump(self.index, f)185 186 logger.info("Index and metadata saved successfully")187 188 except Exception as e:189 logger.error(f"Error saving index: {e}")190 raise191 192 def load_index(self) -> bool:193 """Load existing index and metadata from disk."""194 try:195 metadata_path = self.storage_dir / 'metadata.json'196 197 # Check if files exist198 if not metadata_path.exists():199 logger.info("No existing index found")200 return False201 202 # Load metadata203 with open(metadata_path, 'r', encoding='utf-8') as f:204 self.metadata = json.load(f)205 206 # Load index207 if FAISS_AVAILABLE:208 index_path = self.storage_dir / 'index.faiss'209 if index_path.exists():210 self.index = faiss.read_index(str(index_path))211 else:212 return False213 else:214 index_path = self.storage_dir / 'index.pkl'215 if index_path.exists():216 with open(index_path, 'rb') as f:217 self.index = pickle.load(f)218 else:219 return False220 221 logger.info(f"Loaded index with {len(self.metadata)} documents")222 return True223 224 except Exception as e:225 logger.error(f"Error loading index: {e}")226 return False227 228 def create_index(self, force: bool = False):229 """Create new index from documents in data directory."""230 if not force and self.load_index():231 logger.info("Index already exists and force=False")232 return233 234 logger.info("Creating new index...")235 236 # Load documents237 documents = self.load_documents()238 239 if not documents:240 logger.warning("No documents found to index")241 return242 243 # Extract texts for embedding244 texts = [doc['text'] for doc in documents]245 246 # Create embeddings247 logger.info("Creating embeddings...")248 embeddings = self.create_embeddings(texts)249 250 # Build index251 logger.info("Building search index...")252 self.build_index(embeddings)253 254 # Store metadata255 self.metadata = documents256 257 # Save to disk258 self.save_index()259 260 logger.info("Index creation completed")261 262 def simple_similarity_search(self, query_embedding: np.ndarray, k: int = 4) -> List[int]:263 """Simple similarity search when FAISS is not available."""264 if self.index is None:265 return []266 267 # Compute cosine similarity268 similarities = np.dot(self.index, query_embedding.T).flatten()269 270 # Get top k indices271 top_indices = np.argsort(similarities)[::-1][:k]272 return top_indices.tolist()273 274 def search(self, query_embedding: np.ndarray, k: int = 4) -> List[int]:275 """Search for similar documents."""276 if self.index is None:277 logger.warning("No index available for search")278 return []279 280 try:281 if FAISS_AVAILABLE and hasattr(self.index, 'search'):282 # FAISS search283 query_embedding = query_embedding.reshape(1, -1).astype(np.float32)284 faiss.normalize_L2(query_embedding)285 286 scores, indices = self.index.search(query_embedding, k)287 return indices[0].tolist()288 else:289 # Simple similarity search290 return self.simple_similarity_search(query_embedding, k)291 292 except Exception as e:293 logger.error(f"Error during search: {e}")294 return []295 296 def retrieve(self, query: str, k: int = 4) -> List[Dict]:297 """Retrieve relevant documents for a query."""298 try:299 # Create query embedding300 query_embeddings = self.create_embeddings([query])301 query_embedding = query_embeddings[0]302 303 # Search for similar documents304 indices = self.search(query_embedding, k)305 306 # Return relevant documents with metadata307 results = []308 for idx in indices:309 if 0 <= idx < len(self.metadata):310 results.append(self.metadata[idx])311 312 logger.info(f"Retrieved {len(results)} documents for query")313 return results314 315 except Exception as e:316 logger.error(f"Error during retrieval: {e}")317 return []318 319 def process_message_with_context(self, message: str, session_id: str) -> str:320 """Process user message with conversation context and pattern recognition."""321 # Get conversation history322 history = self.get_conversation_history(session_id)323 324 # Check if message is incomplete/abbreviated325 enhanced_message = self._enhance_incomplete_message(message, history)326 327 logger.info(f"Original: '{message}' -> Enhanced: '{enhanced_message}'")328 return enhanced_message329 330 def _enhance_incomplete_message(self, message: str, history: List[Dict]) -> str:331 """Enhance incomplete messages using pattern recognition and context."""332 message_lower = message.lower().strip()333 334 # If message is complete (more than 10 characters and has context), return as-is335 if len(message.strip()) > 10 and any(word in message_lower for word in ['عن', 'في', 'إلى', 'من', 'مع']):336 return message337 338 # Get the last topic from conversation history339 last_topic = self._extract_last_topic(history)340 341 # Pattern matching for incomplete messages342 enhanced_message = message343 344 # Price inquiry patterns345 for pattern in self.message_patterns['price_patterns']:346 if re.search(pattern, message, re.IGNORECASE):347 if last_topic:348 enhanced_message = f"ما هو سعر {last_topic}؟"349 else:350 enhanced_message = "ما هي أسعار العروض المتاحة؟"351 break352 353 # Requirements inquiry patterns 354 for pattern in self.message_patterns['requirements_patterns']:355 if re.search(pattern, message, re.IGNORECASE):356 if last_topic:357 enhanced_message = f"ما هي متطلبات {last_topic}؟"358 else:359 enhanced_message = "ما هي متطلبات السفر والحجز؟"360 break361 362 # Details inquiry patterns363 for pattern in self.message_patterns['details_patterns']:364 if re.search(pattern, message, re.IGNORECASE):365 if last_topic:366 enhanced_message = f"ما هي تفاصيل {last_topic}؟"367 else:368 enhanced_message = "ما هي العروض والخدمات المتاحة؟"369 break370 371 # Booking inquiry patterns372 for pattern in self.message_patterns['booking_patterns']:373 if re.search(pattern, message, re.IGNORECASE):374 if last_topic:375 enhanced_message = f"كيف يمكنني حجز {last_topic}؟"376 else:377 enhanced_message = "كيف يمكنني حجز رحلة؟"378 break379 380 # Handle very short messages like "الخ..." or "وماذا أيضاً"381 if len(message.strip()) < 5 or message.strip() in ['الخ', 'الخ...', 'وماذا', 'وماذا أيضاً', 'ماذا أيضاً']:382 if last_topic:383 enhanced_message = f"أخبرني المزيد عن {last_topic}"384 else:385 enhanced_message = "أخبرني المزيد عن خدماتكم"386 387 return enhanced_message388 389 def _extract_last_topic(self, history: List[Dict]) -> Optional[str]:390 """Extract the main topic from recent conversation history."""391 if not history:392 return None393 394 # Look at the last few messages to find topics395 topics = []396 for entry in history[-3:]: # Check last 3 exchanges397 user_msg = entry.get('user_message', '').lower()398 399 # Common travel topics to look for400 travel_keywords = {401 'تركيا': 'عروض تركيا',402 'دبي': 'عروض دبي', 403 'باريس': 'عروض باريس',404 'اليابان': 'عروض اليابان',405 'المالديف': 'عروض المالديف',406 'القاهرة': 'عروض القاهرة',407 'طوكيو': 'عروض طوكيو',408 'إسطنبول': 'عروض إسطنبول',409 'العمرة': 'رحلات العمرة',410 'الحج': 'رحلات الحج',411 'شهر العسل': 'رحلات شهر العسل'412 }413 414 for keyword, topic in travel_keywords.items():415 if keyword in user_msg:416 topics.append(topic)417 418 # Return the most recent topic419 return topics[-1] if topics else None420 421 def add_to_conversation_history(self, session_id: str, user_message: str, assistant_response: str):422 """Add a conversation turn to history."""423 if session_id not in self.conversations:424 self.conversations[session_id] = []425 426 self.conversations[session_id].append({427 'user_message': user_message,428 'assistant_response': assistant_response,429 'timestamp': time.time()430 })431 432 # Keep only last 10 exchanges to manage memory433 if len(self.conversations[session_id]) > 10:434 self.conversations[session_id] = self.conversations[session_id][-10:]435 436 def get_conversation_history(self, session_id: str) -> List[Dict]:437 """Get conversation history for a session."""438 return self.conversations.get(session_id, [])439 440 def clear_conversation_history(self, session_id: str):441 """Clear conversation history for a session."""442 if session_id in self.conversations:443 del self.conversations[session_id]