alschameri/helping-source
0
1"""2Enhanced Retrieval-Augmented Generation (RAG) system for Arabic travel agency chatbot.3Implements hybrid retrieval, advanced chunking strategies, and multiple similarity methods.4"""5 6import os7import json8import logging9import pickle10import re11import hashlib12from typing import List, Dict, Optional, Tuple, Union13import numpy as np14from pathlib import Path15from dataclasses import dataclass16from collections import defaultdict17import threading18from concurrent.futures import ThreadPoolExecutor, as_completed19 20try:21 import faiss22 FAISS_AVAILABLE = True23except ImportError:24 FAISS_AVAILABLE = False25 26try:27 from sklearn.feature_extraction.text import TfidfVectorizer28 from sklearn.metrics.pairwise import cosine_similarity29 SKLEARN_AVAILABLE = True30except ImportError:31 SKLEARN_AVAILABLE = False32 33try:34 from sentence_transformers import SentenceTransformer35 SENTENCE_TRANSFORMERS_AVAILABLE = True36except ImportError:37 SENTENCE_TRANSFORMERS_AVAILABLE = False38 39logger = logging.getLogger(__name__)40 41@dataclass42class DocumentChunk:43 """Enhanced document chunk with metadata."""44 text: str45 source_file: str46 chunk_id: int47 start_char: int48 end_char: int49 start_word: int50 end_word: int51 word_count: int52 char_count: int53 hash_id: str54 metadata: Dict = None55 56 def __post_init__(self):57 if self.metadata is None:58 self.metadata = {}59 if not self.hash_id:60 self.hash_id = hashlib.md5(61 f"{self.source_file}_{self.chunk_id}_{self.text[:100]}".encode()62 ).hexdigest()63 64@dataclass65class RetrievalResult:66 """Enhanced retrieval result with scoring."""67 chunk: DocumentChunk68 semantic_score: float69 lexical_score: float70 hybrid_score: float71 rank: int72 73class AdvancedChunker:74 """Advanced text chunking with multiple strategies."""75 76 @staticmethod77 def semantic_chunking(text: str, chunk_size: int = 500, overlap: int = 50) -> List[Dict]:78 """Semantic-aware chunking that preserves sentence boundaries."""79 sentences = re.split(r'[.!?]+', text)80 sentences = [s.strip() for s in sentences if s.strip()]81 82 chunks = []83 current_chunk = []84 current_length = 085 86 for sentence in sentences:87 sentence_length = len(sentence.split())88 89 if current_length + sentence_length > chunk_size and current_chunk:90 # Create chunk from accumulated sentences91 chunk_text = '. '.join(current_chunk) + '.'92 chunks.append({93 'text': chunk_text,94 'sentence_count': len(current_chunk),95 'word_count': current_length96 })97 98 # Start new chunk with overlap99 if overlap > 0 and len(current_chunk) > 1:100 overlap_sentences = current_chunk[-min(overlap // 10, len(current_chunk) - 1):]101 current_chunk = overlap_sentences + [sentence]102 current_length = sum(len(s.split()) for s in current_chunk)103 else:104 current_chunk = [sentence]105 current_length = sentence_length106 else:107 current_chunk.append(sentence)108 current_length += sentence_length109 110 # Add remaining chunk111 if current_chunk:112 chunk_text = '. '.join(current_chunk) + '.'113 chunks.append({114 'text': chunk_text,115 'sentence_count': len(current_chunk),116 'word_count': current_length117 })118 119 return chunks120 121 @staticmethod122 def sliding_window_chunking(text: str, chunk_size: int = 500, overlap: int = 50) -> List[Dict]:123 """Traditional sliding window chunking with character position tracking."""124 words = text.split()125 chunks = []126 127 for i in range(0, len(words), chunk_size - overlap):128 chunk_words = words[i:i + chunk_size]129 chunk_text = ' '.join(chunk_words)130 131 # Calculate character positions132 start_char = len(' '.join(words[:i]))133 end_char = start_char + len(chunk_text)134 if i > 0:135 start_char += 1 # Account for space136 137 chunks.append({138 'text': chunk_text,139 'start_word': i,140 'end_word': min(i + chunk_size, len(words)),141 'start_char': start_char,142 'end_char': end_char,143 'word_count': len(chunk_words),144 'char_count': len(chunk_text)145 })146 147 if i + chunk_size >= len(words):148 break149 150 return chunks151 152class HybridRetriever:153 """Hybrid retrieval combining semantic and lexical search."""154 155 def __init__(self):156 self.tfidf_vectorizer = None157 self.tfidf_matrix = None158 if SKLEARN_AVAILABLE:159 self.tfidf_vectorizer = TfidfVectorizer(160 max_features=5000,161 stop_words='english', # Add Arabic stop words if needed162 ngram_range=(1, 2),163 min_df=2,164 max_df=0.95165 )166 167 def fit_lexical(self, texts: List[str]):168 """Fit TF-IDF vectorizer on document corpus."""169 if not SKLEARN_AVAILABLE:170 logger.warning("scikit-learn not available, lexical search disabled")171 return172 173 try:174 self.tfidf_matrix = self.tfidf_vectorizer.fit_transform(texts)175 logger.info(f"TF-IDF matrix shape: {self.tfidf_matrix.shape}")176 except Exception as e:177 logger.error(f"Error fitting TF-IDF vectorizer: {e}")178 179 def lexical_search(self, query: str, k: int = 10) -> List[Tuple[int, float]]:180 """Perform lexical search using TF-IDF."""181 if not SKLEARN_AVAILABLE or self.tfidf_matrix is None:182 return []183 184 try:185 query_vector = self.tfidf_vectorizer.transform([query])186 similarities = cosine_similarity(query_vector, self.tfidf_matrix).flatten()187 188 # Get top k indices with scores189 top_indices = np.argsort(similarities)[::-1][:k]190 return [(idx, similarities[idx]) for idx in top_indices if similarities[idx] > 0]191 192 except Exception as e:193 logger.error(f"Error in lexical search: {e}")194 return []195 196 def combine_scores(self, semantic_results: List[Tuple[int, float]], 197 lexical_results: List[Tuple[int, float]], 198 semantic_weight: float = 0.7) -> List[Tuple[int, float]]:199 """Combine semantic and lexical scores using weighted fusion."""200 # Normalize scores to [0, 1] range201 def normalize_scores(results):202 if not results:203 return {}204 scores = [score for _, score in results]205 if max(scores) == min(scores):206 return {idx: 1.0 for idx, _ in results}207 208 max_score, min_score = max(scores), min(scores)209 return {idx: (score - min_score) / (max_score - min_score) 210 for idx, score in results}211 212 semantic_scores = normalize_scores(semantic_results)213 lexical_scores = normalize_scores(lexical_results)214 215 # Combine scores216 all_indices = set(semantic_scores.keys()) | set(lexical_scores.keys())217 combined_results = []218 219 for idx in all_indices:220 semantic_score = semantic_scores.get(idx, 0.0)221 lexical_score = lexical_scores.get(idx, 0.0)222 223 # Weighted combination224 hybrid_score = (semantic_weight * semantic_score + 225 (1 - semantic_weight) * lexical_score)226 227 combined_results.append((idx, hybrid_score, semantic_score, lexical_score))228 229 # Sort by hybrid score230 combined_results.sort(key=lambda x: x[1], reverse=True)231 232 return combined_results233 234class RAGSystem:235 """Enhanced RAG system with hybrid retrieval and advanced features."""236 237 def __init__(self, 238 gemini_client,239 data_dir: str = 'data',240 storage_dir: str = 'storage',241 chunk_size: int = 1000,242 overlap: int = 200,243 chunking_strategy: str = 'semantic',244 embedding_model: str = 'gemini'):245 """246 Initialize enhanced RAG system.247 248 Args:249 gemini_client: Gemini client for embeddings250 data_dir: Directory containing source documents251 storage_dir: Directory for storing indexes and metadata252 chunk_size: Size of text chunks in words253 overlap: Overlap between chunks in words254 chunking_strategy: 'semantic' or 'sliding_window'255 embedding_model: 'gemini' or 'sentence_transformer'256 """257 self.gemini_client = gemini_client258 self.data_dir = Path(data_dir)259 self.storage_dir = Path(storage_dir)260 self.chunk_size = chunk_size261 self.overlap = overlap262 self.chunking_strategy = chunking_strategy263 self.embedding_model = embedding_model264 265 # Ensure directories exist266 self.data_dir.mkdir(exist_ok=True)267 self.storage_dir.mkdir(exist_ok=True)268 269 # Initialize components270 self.chunks: List[DocumentChunk] = []271 self.embeddings: Optional[np.ndarray] = None272 self.faiss_index = None273 self.hybrid_retriever = HybridRetriever()274 self.chunker = AdvancedChunker()275 276 # Threading lock for concurrent access277 self.lock = threading.Lock()278 279 # Initialize sentence transformer if available280 self.sentence_model = None281 if embedding_model == 'sentence_transformer' and SENTENCE_TRANSFORMERS_AVAILABLE:282 try:283 self.sentence_model = SentenceTransformer('all-MiniLM-L6-v2')284 logger.info("Initialized sentence transformer model")285 except Exception as e:286 logger.error(f"Failed to load sentence transformer: {e}")287 288 # Embedding dimensions289 self.embedding_dim = 768 if embedding_model == 'gemini' else 384290 291 def load_documents(self) -> List[DocumentChunk]:292 """Load and process documents with enhanced chunking."""293 chunks = []294 295 logger.info(f"Loading documents from {self.data_dir}")296 297 # Support multiple file formats298 file_patterns = ['*.txt', '*.md', '*.json']299 all_files = []300 for pattern in file_patterns:301 all_files.extend(self.data_dir.glob(pattern))302 303 for file_path in all_files:304 try:305 with open(file_path, 'r', encoding='utf-8') as f:306 if file_path.suffix == '.json':307 data = json.load(f)308 # Extract text content from JSON309 if isinstance(data, dict):310 content = ' '.join(str(v) for v in data.values() if isinstance(v, str))311 elif isinstance(data, list):312 content = ' '.join(str(item) for item in data if isinstance(item, str))313 else:314 content = str(data)315 else:316 content = f.read().strip()317 318 if not content:319 logger.warning(f"Empty file: {file_path}")320 continue321 322 # Choose chunking strategy323 if self.chunking_strategy == 'semantic':324 raw_chunks = self.chunker.semantic_chunking(content, self.chunk_size, self.overlap)325 else:326 raw_chunks = self.chunker.sliding_window_chunking(content, self.chunk_size, self.overlap)327 328 # Convert to DocumentChunk objects329 for i, chunk_data in enumerate(raw_chunks):330 chunk = DocumentChunk(331 text=chunk_data['text'],332 source_file=file_path.name,333 chunk_id=i,334 start_char=chunk_data.get('start_char', 0),335 end_char=chunk_data.get('end_char', 0),336 start_word=chunk_data.get('start_word', 0),337 end_word=chunk_data.get('end_word', 0),338 word_count=chunk_data.get('word_count', 0),339 char_count=chunk_data.get('char_count', len(chunk_data['text'])),340 hash_id='',341 metadata={'file_type': file_path.suffix}342 )343 chunks.append(chunk)344 345 logger.info(f"Processed {len(raw_chunks)} chunks from {file_path.name}")346 347 except Exception as e:348 logger.error(f"Error processing {file_path}: {e}")349 continue350 351 logger.info(f"Total chunks loaded: {len(chunks)}")352 return chunks353 354 def create_embeddings_batch(self, texts: List[str], batch_size: int = 32) -> np.ndarray:355 """Create embeddings in batches for better performance."""356 embeddings = []357 358 for i in range(0, len(texts), batch_size):359 batch_texts = texts[i:i + batch_size]360 361 try:362 if self.embedding_model == 'sentence_transformer' and self.sentence_model:363 batch_embeddings = self.sentence_model.encode(batch_texts, convert_to_numpy=True)364 else:365 # Use Gemini embeddings366 batch_embeddings = self.gemini_client.embed_texts(batch_texts)367 batch_embeddings = np.array(batch_embeddings)368 369 embeddings.append(batch_embeddings)370 logger.info(f"Processed embedding batch {i//batch_size + 1}/{(len(texts) + batch_size - 1)//batch_size}")371 372 except Exception as e:373 logger.error(f"Error creating embeddings for batch {i//batch_size + 1}: {e}")374 # Fallback to random embeddings375 fallback_embeddings = np.random.rand(len(batch_texts), self.embedding_dim).astype(np.float32)376 embeddings.append(fallback_embeddings)377 378 return np.vstack(embeddings) if embeddings else np.array([])379 380 def build_faiss_index(self, embeddings: np.ndarray):381 """Build optimized FAISS index."""382 if not FAISS_AVAILABLE:383 logger.warning("FAISS not available")384 return385 386 try:387 dimension = embeddings.shape[1]388 389 # Use more sophisticated index for better retrieval390 if embeddings.shape[0] > 1000:391 # Use HNSW index for larger datasets392 self.faiss_index = faiss.IndexHNSWFlat(dimension, 32)393 self.faiss_index.hnsw.efConstruction = 200394 self.faiss_index.hnsw.efSearch = 50395 else:396 # Use flat index for smaller datasets397 self.faiss_index = faiss.IndexFlatIP(dimension)398 399 # Normalize embeddings for cosine similarity400 embeddings_normalized = embeddings.copy()401 faiss.normalize_L2(embeddings_normalized)402 403 # Add to index404 self.faiss_index.add(embeddings_normalized.astype(np.float32))405 406 logger.info(f"Built FAISS index: {type(self.faiss_index).__name__} with {self.faiss_index.ntotal} vectors")407 408 except Exception as e:409 logger.error(f"Error building FAISS index: {e}")410 self.faiss_index = None411 412 def save_index(self):413 """Save all components to disk."""414 try:415 # Save chunks metadata416 chunks_data = []417 for chunk in self.chunks:418 chunks_data.append({419 'text': chunk.text,420 'source_file': chunk.source_file,421 'chunk_id': chunk.chunk_id,422 'start_char': chunk.start_char,423 'end_char': chunk.end_char,424 'start_word': chunk.start_word,425 'end_word': chunk.end_word,426 'word_count': chunk.word_count,427 'char_count': chunk.char_count,428 'hash_id': chunk.hash_id,429 'metadata': chunk.metadata430 })431 432 with open(self.storage_dir / 'chunks.json', 'w', encoding='utf-8') as f:433 json.dump(chunks_data, f, ensure_ascii=False, indent=2)434 435 # Save embeddings436 if self.embeddings is not None:437 np.save(self.storage_dir / 'embeddings.npy', self.embeddings)438 439 # Save FAISS index440 if self.faiss_index is not None:441 faiss.write_index(self.faiss_index, str(self.storage_dir / 'faiss_index.bin'))442 443 # Save TF-IDF components444 if self.hybrid_retriever.tfidf_vectorizer is not None:445 with open(self.storage_dir / 'tfidf_vectorizer.pkl', 'wb') as f:446 pickle.dump(self.hybrid_retriever.tfidf_vectorizer, f)447 448 with open(self.storage_dir / 'tfidf_matrix.pkl', 'wb') as f:449 pickle.dump(self.hybrid_retriever.tfidf_matrix, f)450 451 logger.info("Index and all components saved successfully")452 453 except Exception as e:454 logger.error(f"Error saving index: {e}")455 raise456 457 def load_index(self) -> bool:458 """Load all components from disk."""459 try:460 chunks_path = self.storage_dir / 'chunks.json'461 embeddings_path = self.storage_dir / 'embeddings.npy'462 463 if not chunks_path.exists():464 logger.info("No existing index found")465 return False466 467 # Load chunks468 with open(chunks_path, 'r', encoding='utf-8') as f:469 chunks_data = json.load(f)470 471 self.chunks = []472 for chunk_data in chunks_data:473 chunk = DocumentChunk(**chunk_data)474 self.chunks.append(chunk)475 476 # Load embeddings477 if embeddings_path.exists():478 self.embeddings = np.load(embeddings_path)479 480 # Load FAISS index481 faiss_path = self.storage_dir / 'faiss_index.bin'482 if FAISS_AVAILABLE and faiss_path.exists():483 self.faiss_index = faiss.read_index(str(faiss_path))484 485 # Load TF-IDF components486 tfidf_vectorizer_path = self.storage_dir / 'tfidf_vectorizer.pkl'487 tfidf_matrix_path = self.storage_dir / 'tfidf_matrix.pkl'488 489 if tfidf_vectorizer_path.exists() and tfidf_matrix_path.exists():490 with open(tfidf_vectorizer_path, 'rb') as f:491 self.hybrid_retriever.tfidf_vectorizer = pickle.load(f)492 493 with open(tfidf_matrix_path, 'rb') as f:494 self.hybrid_retriever.tfidf_matrix = pickle.load(f)495 496 logger.info(f"Loaded index with {len(self.chunks)} chunks")497 return True498 499 except Exception as e:500 logger.error(f"Error loading index: {e}")501 return False502 503 def create_index(self, force: bool = False):504 """Create comprehensive index with all components."""505 if not force and self.load_index():506 logger.info("Index already exists and force=False")507 return508 509 logger.info("Creating enhanced index...")510 511 # Load and process documents512 self.chunks = self.load_documents()513 514 if not self.chunks:515 logger.warning("No documents found to index")516 return517 518 # Extract texts519 texts = [chunk.text for chunk in self.chunks]520 521 # Create embeddings522 logger.info("Creating embeddings...")523 self.embeddings = self.create_embeddings_batch(texts)524 525 # Build FAISS index526 logger.info("Building FAISS index...")527 self.build_faiss_index(self.embeddings)528 529 # Build TF-IDF index530 logger.info("Building TF-IDF index...")531 self.hybrid_retriever.fit_lexical(texts)532 533 # Save all components534 self.save_index()535 536 logger.info("Enhanced index creation completed")537 538 def semantic_search(self, query: str, k: int = 10) -> List[Tuple[int, float]]:539 """Perform semantic search using embeddings."""540 if self.faiss_index is None or self.embeddings is None:541 return []542 543 try:544 # Create query embedding545 if self.embedding_model == 'sentence_transformer' and self.sentence_model:546 query_embedding = self.sentence_model.encode([query], convert_to_numpy=True)547 else:548 query_embeddings = self.gemini_client.embed_texts([query])549 query_embedding = np.array(query_embeddings)550 551 # Normalize query embedding552 faiss.normalize_L2(query_embedding.astype(np.float32))553 554 # Search555 scores, indices = self.faiss_index.search(query_embedding.astype(np.float32), k)556 557 # Return results with scores558 results = []559 for i, (idx, score) in enumerate(zip(indices[0], scores[0])):560 if idx != -1: # Valid index561 results.append((int(idx), float(score)))562 563 return results564 565 except Exception as e:566 logger.error(f"Error in semantic search: {e}")567 return []568 569 def hybrid_retrieve(self, query: str, k: int = 5, semantic_weight: float = 0.7) -> List[RetrievalResult]:570 """Perform hybrid retrieval combining semantic and lexical search."""571 with self.lock:572 try:573 # Perform semantic search574 semantic_results = self.semantic_search(query, k * 2)575 576 # Perform lexical search577 lexical_results = self.hybrid_retriever.lexical_search(query, k * 2)578 579 # Combine results580 combined_results = self.hybrid_retriever.combine_scores(581 semantic_results, lexical_results, semantic_weight582 )583 584 # Create RetrievalResult objects585 results = []586 for rank, (idx, hybrid_score, semantic_score, lexical_score) in enumerate(combined_results[:k]):587 if 0 <= idx < len(self.chunks):588 result = RetrievalResult(589 chunk=self.chunks[idx],590 semantic_score=semantic_score,591 lexical_score=lexical_score,592 hybrid_score=hybrid_score,593 rank=rank594 )595 results.append(result)596 597 logger.info(f"Retrieved {len(results)} results for query")598 return results599 600 except Exception as e:601 logger.error(f"Error in hybrid retrieval: {e}")602 return []603 604 def retrieve_with_context(self, query: str, k: int = 5, 605 expand_context: bool = True) -> List[Dict]:606 """Retrieve documents with expanded context from adjacent chunks."""607 results = self.hybrid_retrieve(query, k)608 609 if not expand_context:610 return [self._result_to_dict(result) for result in results]611 612 # Group results by source file613 file_groups = defaultdict(list)614 for result in results:615 file_groups[result.chunk.source_file].append(result)616 617 # Expand context for each result618 expanded_results = []619 for result in results:620 context_chunks = [result.chunk]621 622 # Find adjacent chunks623 same_file_chunks = [c for c in self.chunks 624 if c.source_file == result.chunk.source_file]625 same_file_chunks.sort(key=lambda x: x.chunk_id)626 627 current_idx = next((i for i, c in enumerate(same_file_chunks) 628 if c.hash_id == result.chunk.hash_id), -1)629 630 if current_idx != -1:631 # Add previous and next chunks632 if current_idx > 0:633 context_chunks.insert(0, same_file_chunks[current_idx - 1])634 if current_idx < len(same_file_chunks) - 1:635 context_chunks.append(same_file_chunks[current_idx + 1])636 637 # Combine context638 expanded_text = ' '.join(chunk.text for chunk in context_chunks)639 640 result_dict = self._result_to_dict(result)641 result_dict['expanded_text'] = expanded_text642 result_dict['context_chunks'] = len(context_chunks)643 644 expanded_results.append(result_dict)645 646 return expanded_results647 648 def _result_to_dict(self, result: RetrievalResult) -> Dict:649 """Convert RetrievalResult to dictionary."""650 return {651 'text': result.chunk.text,652 'source_file': result.chunk.source_file,653 'chunk_id': result.chunk.chunk_id,654 'semantic_score': result.semantic_score,655 'lexical_score': result.lexical_score,656 'hybrid_score': result.hybrid_score,657 'rank': result.rank,658 'word_count': result.chunk.word_count,659 'char_count': result.chunk.char_count,660 'metadata': result.chunk.metadata661 }662 663 def get_stats(self) -> Dict:664 """Get comprehensive statistics about the RAG system."""665 stats = {666 'total_chunks': len(self.chunks),667 'total_files': len(set(chunk.source_file for chunk in self.chunks)),668 'embedding_model': self.embedding_model,669 'chunking_strategy': self.chunking_strategy,670 'chunk_size': self.chunk_size,671 'overlap': self.overlap,672 'faiss_available': FAISS_AVAILABLE,673 'sklearn_available': SKLEARN_AVAILABLE,674 'sentence_transformers_available': SENTENCE_TRANSFORMERS_AVAILABLE,675 'faiss_index_type': type(self.faiss_index).__name__ if self.faiss_index else None,676 'embedding_dimension': self.embedding_dim,677 'tfidf_features': self.hybrid_retriever.tfidf_matrix.shape[1] if self.hybrid_retriever.tfidf_matrix is not None else 0678 }679 680 if self.chunks:681 word_counts = [chunk.word_count for chunk in self.chunks]682 stats.update({683 'avg_chunk_words': np.mean(word_counts),684 'min_chunk_words': min(word_counts),685 'max_chunk_words': max(word_counts),686 'median_chunk_words': np.median(word_counts)687 })688 689 return stats690 691 692# Example usage and initialization693def initialize_enhanced_rag(gemini_client, **kwargs):694 """Initialize enhanced RAG system with optimal settings."""695 return RAGSystem(696 gemini_client=gemini_client,697 chunking_strategy='semantic', # Use semantic chunking by default698 embedding_model='gemini',699 **kwargs700 )