""" Enhanced Retrieval-Augmented Generation (RAG) system for Arabic travel agency chatbot. Implements hybrid retrieval, advanced chunking strategies, and multiple similarity methods. """ import os import json import logging import pickle import re import hashlib from typing import List, Dict, Optional, Tuple, Union import numpy as np from pathlib import Path from dataclasses import dataclass from collections import defaultdict import threading from concurrent.futures import ThreadPoolExecutor, as_completed try: import faiss FAISS_AVAILABLE = True except ImportError: FAISS_AVAILABLE = False try: from sklearn.feature_extraction.text import TfidfVectorizer from sklearn.metrics.pairwise import cosine_similarity SKLEARN_AVAILABLE = True except ImportError: SKLEARN_AVAILABLE = False try: from sentence_transformers import SentenceTransformer SENTENCE_TRANSFORMERS_AVAILABLE = True except ImportError: SENTENCE_TRANSFORMERS_AVAILABLE = False logger = logging.getLogger(__name__) @dataclass class DocumentChunk: """Enhanced document chunk with metadata.""" text: str source_file: str chunk_id: int start_char: int end_char: int start_word: int end_word: int word_count: int char_count: int hash_id: str metadata: Dict = None def __post_init__(self): if self.metadata is None: self.metadata = {} if not self.hash_id: self.hash_id = hashlib.md5( f"{self.source_file}_{self.chunk_id}_{self.text[:100]}".encode() ).hexdigest() @dataclass class RetrievalResult: """Enhanced retrieval result with scoring.""" chunk: DocumentChunk semantic_score: float lexical_score: float hybrid_score: float rank: int class AdvancedChunker: """Advanced text chunking with multiple strategies.""" @staticmethod def semantic_chunking(text: str, chunk_size: int = 500, overlap: int = 50) -> List[Dict]: """Semantic-aware chunking that preserves sentence boundaries.""" sentences = re.split(r'[.!?]+', text) sentences = [s.strip() for s in sentences if s.strip()] chunks = [] current_chunk = [] current_length = 0 for sentence in sentences: sentence_length = len(sentence.split()) if current_length + sentence_length > chunk_size and current_chunk: # Create chunk from accumulated sentences chunk_text = '. '.join(current_chunk) + '.' chunks.append({ 'text': chunk_text, 'sentence_count': len(current_chunk), 'word_count': current_length }) # Start new chunk with overlap if overlap > 0 and len(current_chunk) > 1: overlap_sentences = current_chunk[-min(overlap // 10, len(current_chunk) - 1):] current_chunk = overlap_sentences + [sentence] current_length = sum(len(s.split()) for s in current_chunk) else: current_chunk = [sentence] current_length = sentence_length else: current_chunk.append(sentence) current_length += sentence_length # Add remaining chunk if current_chunk: chunk_text = '. '.join(current_chunk) + '.' chunks.append({ 'text': chunk_text, 'sentence_count': len(current_chunk), 'word_count': current_length }) return chunks @staticmethod def sliding_window_chunking(text: str, chunk_size: int = 500, overlap: int = 50) -> List[Dict]: """Traditional sliding window chunking with character position tracking.""" words = text.split() chunks = [] for i in range(0, len(words), chunk_size - overlap): chunk_words = words[i:i + chunk_size] chunk_text = ' '.join(chunk_words) # Calculate character positions start_char = len(' '.join(words[:i])) end_char = start_char + len(chunk_text) if i > 0: start_char += 1 # Account for space chunks.append({ 'text': chunk_text, 'start_word': i, 'end_word': min(i + chunk_size, len(words)), 'start_char': start_char, 'end_char': end_char, 'word_count': len(chunk_words), 'char_count': len(chunk_text) }) if i + chunk_size >= len(words): break return chunks class HybridRetriever: """Hybrid retrieval combining semantic and lexical search.""" def __init__(self): self.tfidf_vectorizer = None self.tfidf_matrix = None if SKLEARN_AVAILABLE: self.tfidf_vectorizer = TfidfVectorizer( max_features=5000, stop_words='english', # Add Arabic stop words if needed ngram_range=(1, 2), min_df=2, max_df=0.95 ) def fit_lexical(self, texts: List[str]): """Fit TF-IDF vectorizer on document corpus.""" if not SKLEARN_AVAILABLE: logger.warning("scikit-learn not available, lexical search disabled") return try: self.tfidf_matrix = self.tfidf_vectorizer.fit_transform(texts) logger.info(f"TF-IDF matrix shape: {self.tfidf_matrix.shape}") except Exception as e: logger.error(f"Error fitting TF-IDF vectorizer: {e}") def lexical_search(self, query: str, k: int = 10) -> List[Tuple[int, float]]: """Perform lexical search using TF-IDF.""" if not SKLEARN_AVAILABLE or self.tfidf_matrix is None: return [] try: query_vector = self.tfidf_vectorizer.transform([query]) similarities = cosine_similarity(query_vector, self.tfidf_matrix).flatten() # Get top k indices with scores top_indices = np.argsort(similarities)[::-1][:k] return [(idx, similarities[idx]) for idx in top_indices if similarities[idx] > 0] except Exception as e: logger.error(f"Error in lexical search: {e}") return [] def combine_scores(self, semantic_results: List[Tuple[int, float]], lexical_results: List[Tuple[int, float]], semantic_weight: float = 0.7) -> List[Tuple[int, float]]: """Combine semantic and lexical scores using weighted fusion.""" # Normalize scores to [0, 1] range def normalize_scores(results): if not results: return {} scores = [score for _, score in results] if max(scores) == min(scores): return {idx: 1.0 for idx, _ in results} max_score, min_score = max(scores), min(scores) return {idx: (score - min_score) / (max_score - min_score) for idx, score in results} semantic_scores = normalize_scores(semantic_results) lexical_scores = normalize_scores(lexical_results) # Combine scores all_indices = set(semantic_scores.keys()) | set(lexical_scores.keys()) combined_results = [] for idx in all_indices: semantic_score = semantic_scores.get(idx, 0.0) lexical_score = lexical_scores.get(idx, 0.0) # Weighted combination hybrid_score = (semantic_weight * semantic_score + (1 - semantic_weight) * lexical_score) combined_results.append((idx, hybrid_score, semantic_score, lexical_score)) # Sort by hybrid score combined_results.sort(key=lambda x: x[1], reverse=True) return combined_results class RAGSystem: """Enhanced RAG system with hybrid retrieval and advanced features.""" def __init__(self, gemini_client, data_dir: str = 'data', storage_dir: str = 'storage', chunk_size: int = 1000, overlap: int = 200, chunking_strategy: str = 'semantic', embedding_model: str = 'gemini'): """ Initialize enhanced RAG system. Args: gemini_client: Gemini client for embeddings data_dir: Directory containing source documents storage_dir: Directory for storing indexes and metadata chunk_size: Size of text chunks in words overlap: Overlap between chunks in words chunking_strategy: 'semantic' or 'sliding_window' embedding_model: 'gemini' or 'sentence_transformer' """ self.gemini_client = gemini_client self.data_dir = Path(data_dir) self.storage_dir = Path(storage_dir) self.chunk_size = chunk_size self.overlap = overlap self.chunking_strategy = chunking_strategy self.embedding_model = embedding_model # Ensure directories exist self.data_dir.mkdir(exist_ok=True) self.storage_dir.mkdir(exist_ok=True) # Initialize components self.chunks: List[DocumentChunk] = [] self.embeddings: Optional[np.ndarray] = None self.faiss_index = None self.hybrid_retriever = HybridRetriever() self.chunker = AdvancedChunker() # Threading lock for concurrent access self.lock = threading.Lock() # Initialize sentence transformer if available self.sentence_model = None if embedding_model == 'sentence_transformer' and SENTENCE_TRANSFORMERS_AVAILABLE: try: self.sentence_model = SentenceTransformer('all-MiniLM-L6-v2') logger.info("Initialized sentence transformer model") except Exception as e: logger.error(f"Failed to load sentence transformer: {e}") # Embedding dimensions self.embedding_dim = 768 if embedding_model == 'gemini' else 384 def load_documents(self) -> List[DocumentChunk]: """Load and process documents with enhanced chunking.""" chunks = [] logger.info(f"Loading documents from {self.data_dir}") # Support multiple file formats file_patterns = ['*.txt', '*.md', '*.json'] all_files = [] for pattern in file_patterns: all_files.extend(self.data_dir.glob(pattern)) for file_path in all_files: try: with open(file_path, 'r', encoding='utf-8') as f: if file_path.suffix == '.json': data = json.load(f) # Extract text content from JSON if isinstance(data, dict): content = ' '.join(str(v) for v in data.values() if isinstance(v, str)) elif isinstance(data, list): content = ' '.join(str(item) for item in data if isinstance(item, str)) else: content = str(data) else: content = f.read().strip() if not content: logger.warning(f"Empty file: {file_path}") continue # Choose chunking strategy if self.chunking_strategy == 'semantic': raw_chunks = self.chunker.semantic_chunking(content, self.chunk_size, self.overlap) else: raw_chunks = self.chunker.sliding_window_chunking(content, self.chunk_size, self.overlap) # Convert to DocumentChunk objects for i, chunk_data in enumerate(raw_chunks): chunk = DocumentChunk( text=chunk_data['text'], source_file=file_path.name, chunk_id=i, start_char=chunk_data.get('start_char', 0), end_char=chunk_data.get('end_char', 0), start_word=chunk_data.get('start_word', 0), end_word=chunk_data.get('end_word', 0), word_count=chunk_data.get('word_count', 0), char_count=chunk_data.get('char_count', len(chunk_data['text'])), hash_id='', metadata={'file_type': file_path.suffix} ) chunks.append(chunk) logger.info(f"Processed {len(raw_chunks)} chunks from {file_path.name}") except Exception as e: logger.error(f"Error processing {file_path}: {e}") continue logger.info(f"Total chunks loaded: {len(chunks)}") return chunks def create_embeddings_batch(self, texts: List[str], batch_size: int = 32) -> np.ndarray: """Create embeddings in batches for better performance.""" embeddings = [] for i in range(0, len(texts), batch_size): batch_texts = texts[i:i + batch_size] try: if self.embedding_model == 'sentence_transformer' and self.sentence_model: batch_embeddings = self.sentence_model.encode(batch_texts, convert_to_numpy=True) else: # Use Gemini embeddings batch_embeddings = self.gemini_client.embed_texts(batch_texts) batch_embeddings = np.array(batch_embeddings) embeddings.append(batch_embeddings) logger.info(f"Processed embedding batch {i//batch_size + 1}/{(len(texts) + batch_size - 1)//batch_size}") except Exception as e: logger.error(f"Error creating embeddings for batch {i//batch_size + 1}: {e}") # Fallback to random embeddings fallback_embeddings = np.random.rand(len(batch_texts), self.embedding_dim).astype(np.float32) embeddings.append(fallback_embeddings) return np.vstack(embeddings) if embeddings else np.array([]) def build_faiss_index(self, embeddings: np.ndarray): """Build optimized FAISS index.""" if not FAISS_AVAILABLE: logger.warning("FAISS not available") return try: dimension = embeddings.shape[1] # Use more sophisticated index for better retrieval if embeddings.shape[0] > 1000: # Use HNSW index for larger datasets self.faiss_index = faiss.IndexHNSWFlat(dimension, 32) self.faiss_index.hnsw.efConstruction = 200 self.faiss_index.hnsw.efSearch = 50 else: # Use flat index for smaller datasets self.faiss_index = faiss.IndexFlatIP(dimension) # Normalize embeddings for cosine similarity embeddings_normalized = embeddings.copy() faiss.normalize_L2(embeddings_normalized) # Add to index self.faiss_index.add(embeddings_normalized.astype(np.float32)) logger.info(f"Built FAISS index: {type(self.faiss_index).__name__} with {self.faiss_index.ntotal} vectors") except Exception as e: logger.error(f"Error building FAISS index: {e}") self.faiss_index = None def save_index(self): """Save all components to disk.""" try: # Save chunks metadata chunks_data = [] for chunk in self.chunks: chunks_data.append({ 'text': chunk.text, 'source_file': chunk.source_file, 'chunk_id': chunk.chunk_id, 'start_char': chunk.start_char, 'end_char': chunk.end_char, 'start_word': chunk.start_word, 'end_word': chunk.end_word, 'word_count': chunk.word_count, 'char_count': chunk.char_count, 'hash_id': chunk.hash_id, 'metadata': chunk.metadata }) with open(self.storage_dir / 'chunks.json', 'w', encoding='utf-8') as f: json.dump(chunks_data, f, ensure_ascii=False, indent=2) # Save embeddings if self.embeddings is not None: np.save(self.storage_dir / 'embeddings.npy', self.embeddings) # Save FAISS index if self.faiss_index is not None: faiss.write_index(self.faiss_index, str(self.storage_dir / 'faiss_index.bin')) # Save TF-IDF components if self.hybrid_retriever.tfidf_vectorizer is not None: with open(self.storage_dir / 'tfidf_vectorizer.pkl', 'wb') as f: pickle.dump(self.hybrid_retriever.tfidf_vectorizer, f) with open(self.storage_dir / 'tfidf_matrix.pkl', 'wb') as f: pickle.dump(self.hybrid_retriever.tfidf_matrix, f) logger.info("Index and all components saved successfully") except Exception as e: logger.error(f"Error saving index: {e}") raise def load_index(self) -> bool: """Load all components from disk.""" try: chunks_path = self.storage_dir / 'chunks.json' embeddings_path = self.storage_dir / 'embeddings.npy' if not chunks_path.exists(): logger.info("No existing index found") return False # Load chunks with open(chunks_path, 'r', encoding='utf-8') as f: chunks_data = json.load(f) self.chunks = [] for chunk_data in chunks_data: chunk = DocumentChunk(**chunk_data) self.chunks.append(chunk) # Load embeddings if embeddings_path.exists(): self.embeddings = np.load(embeddings_path) # Load FAISS index faiss_path = self.storage_dir / 'faiss_index.bin' if FAISS_AVAILABLE and faiss_path.exists(): self.faiss_index = faiss.read_index(str(faiss_path)) # Load TF-IDF components tfidf_vectorizer_path = self.storage_dir / 'tfidf_vectorizer.pkl' tfidf_matrix_path = self.storage_dir / 'tfidf_matrix.pkl' if tfidf_vectorizer_path.exists() and tfidf_matrix_path.exists(): with open(tfidf_vectorizer_path, 'rb') as f: self.hybrid_retriever.tfidf_vectorizer = pickle.load(f) with open(tfidf_matrix_path, 'rb') as f: self.hybrid_retriever.tfidf_matrix = pickle.load(f) logger.info(f"Loaded index with {len(self.chunks)} chunks") return True except Exception as e: logger.error(f"Error loading index: {e}") return False def create_index(self, force: bool = False): """Create comprehensive index with all components.""" if not force and self.load_index(): logger.info("Index already exists and force=False") return logger.info("Creating enhanced index...") # Load and process documents self.chunks = self.load_documents() if not self.chunks: logger.warning("No documents found to index") return # Extract texts texts = [chunk.text for chunk in self.chunks] # Create embeddings logger.info("Creating embeddings...") self.embeddings = self.create_embeddings_batch(texts) # Build FAISS index logger.info("Building FAISS index...") self.build_faiss_index(self.embeddings) # Build TF-IDF index logger.info("Building TF-IDF index...") self.hybrid_retriever.fit_lexical(texts) # Save all components self.save_index() logger.info("Enhanced index creation completed") def semantic_search(self, query: str, k: int = 10) -> List[Tuple[int, float]]: """Perform semantic search using embeddings.""" if self.faiss_index is None or self.embeddings is None: return [] try: # Create query embedding if self.embedding_model == 'sentence_transformer' and self.sentence_model: query_embedding = self.sentence_model.encode([query], convert_to_numpy=True) else: query_embeddings = self.gemini_client.embed_texts([query]) query_embedding = np.array(query_embeddings) # Normalize query embedding faiss.normalize_L2(query_embedding.astype(np.float32)) # Search scores, indices = self.faiss_index.search(query_embedding.astype(np.float32), k) # Return results with scores results = [] for i, (idx, score) in enumerate(zip(indices[0], scores[0])): if idx != -1: # Valid index results.append((int(idx), float(score))) return results except Exception as e: logger.error(f"Error in semantic search: {e}") return [] def hybrid_retrieve(self, query: str, k: int = 5, semantic_weight: float = 0.7) -> List[RetrievalResult]: """Perform hybrid retrieval combining semantic and lexical search.""" with self.lock: try: # Perform semantic search semantic_results = self.semantic_search(query, k * 2) # Perform lexical search lexical_results = self.hybrid_retriever.lexical_search(query, k * 2) # Combine results combined_results = self.hybrid_retriever.combine_scores( semantic_results, lexical_results, semantic_weight ) # Create RetrievalResult objects results = [] for rank, (idx, hybrid_score, semantic_score, lexical_score) in enumerate(combined_results[:k]): if 0 <= idx < len(self.chunks): result = RetrievalResult( chunk=self.chunks[idx], semantic_score=semantic_score, lexical_score=lexical_score, hybrid_score=hybrid_score, rank=rank ) results.append(result) logger.info(f"Retrieved {len(results)} results for query") return results except Exception as e: logger.error(f"Error in hybrid retrieval: {e}") return [] def retrieve_with_context(self, query: str, k: int = 5, expand_context: bool = True) -> List[Dict]: """Retrieve documents with expanded context from adjacent chunks.""" results = self.hybrid_retrieve(query, k) if not expand_context: return [self._result_to_dict(result) for result in results] # Group results by source file file_groups = defaultdict(list) for result in results: file_groups[result.chunk.source_file].append(result) # Expand context for each result expanded_results = [] for result in results: context_chunks = [result.chunk] # Find adjacent chunks same_file_chunks = [c for c in self.chunks if c.source_file == result.chunk.source_file] same_file_chunks.sort(key=lambda x: x.chunk_id) current_idx = next((i for i, c in enumerate(same_file_chunks) if c.hash_id == result.chunk.hash_id), -1) if current_idx != -1: # Add previous and next chunks if current_idx > 0: context_chunks.insert(0, same_file_chunks[current_idx - 1]) if current_idx < len(same_file_chunks) - 1: context_chunks.append(same_file_chunks[current_idx + 1]) # Combine context expanded_text = ' '.join(chunk.text for chunk in context_chunks) result_dict = self._result_to_dict(result) result_dict['expanded_text'] = expanded_text result_dict['context_chunks'] = len(context_chunks) expanded_results.append(result_dict) return expanded_results def _result_to_dict(self, result: RetrievalResult) -> Dict: """Convert RetrievalResult to dictionary.""" return { 'text': result.chunk.text, 'source_file': result.chunk.source_file, 'chunk_id': result.chunk.chunk_id, 'semantic_score': result.semantic_score, 'lexical_score': result.lexical_score, 'hybrid_score': result.hybrid_score, 'rank': result.rank, 'word_count': result.chunk.word_count, 'char_count': result.chunk.char_count, 'metadata': result.chunk.metadata } def get_stats(self) -> Dict: """Get comprehensive statistics about the RAG system.""" stats = { 'total_chunks': len(self.chunks), 'total_files': len(set(chunk.source_file for chunk in self.chunks)), 'embedding_model': self.embedding_model, 'chunking_strategy': self.chunking_strategy, 'chunk_size': self.chunk_size, 'overlap': self.overlap, 'faiss_available': FAISS_AVAILABLE, 'sklearn_available': SKLEARN_AVAILABLE, 'sentence_transformers_available': SENTENCE_TRANSFORMERS_AVAILABLE, 'faiss_index_type': type(self.faiss_index).__name__ if self.faiss_index else None, 'embedding_dimension': self.embedding_dim, 'tfidf_features': self.hybrid_retriever.tfidf_matrix.shape[1] if self.hybrid_retriever.tfidf_matrix is not None else 0 } if self.chunks: word_counts = [chunk.word_count for chunk in self.chunks] stats.update({ 'avg_chunk_words': np.mean(word_counts), 'min_chunk_words': min(word_counts), 'max_chunk_words': max(word_counts), 'median_chunk_words': np.median(word_counts) }) return stats # Example usage and initialization def initialize_enhanced_rag(gemini_client, **kwargs): """Initialize enhanced RAG system with optimal settings.""" return RAGSystem( gemini_client=gemini_client, chunking_strategy='semantic', # Use semantic chunking by default embedding_model='gemini', **kwargs )