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_arabic(text: str, chunk_size: int = 500, overlap: int = 50) -> List[Dict]:78 """Arabic-aware semantic chunking with proper sentence boundary detection."""79 # Arabic sentence endings80 arabic_sentence_endings = r'[.!؟۔]'81 sentences = re.split(arabic_sentence_endings, text)82 sentences = [s.strip() for s in sentences if s.strip()]83 84 chunks = []85 current_chunk = []86 current_length = 087 88 for sentence in sentences:89 # Count Arabic words properly90 arabic_words = re.findall(r'[\u0600-\u06FF\u0750-\u077F\u08A0-\u08FF\uFB50-\uFDCF\uFDF0-\uFDFF\uFE70-\uFEFF]+', sentence)91 english_words = re.findall(r'[a-zA-Z]+', sentence)92 sentence_length = len(arabic_words) + len(english_words)93 94 if current_length + sentence_length > chunk_size and current_chunk:95 # Create chunk from accumulated sentences96 chunk_text = '. '.join(current_chunk) + '.'97 chunks.append({98 'text': chunk_text,99 'sentence_count': len(current_chunk),100 'word_count': current_length,101 'arabic_words': len(re.findall(r'[\u0600-\u06FF\u0750-\u077F\u08A0-\u08FF\uFB50-\uFDCF\uFDF0-\uFDFF\uFE70-\uFEFF]+', chunk_text)),102 'language': 'mixed' if re.search(r'[a-zA-Z]', chunk_text) and re.search(r'[\u0600-\u06FF]', chunk_text) else 'arabic'103 })104 105 # Start new chunk with overlap106 if overlap > 0 and len(current_chunk) > 1:107 overlap_sentences = current_chunk[-min(overlap // 10, len(current_chunk) - 1):]108 current_chunk = overlap_sentences + [sentence]109 current_length = sum(len(re.findall(r'[\u0600-\u06FF\u0750-\u077F\u08A0-\u08FF\uFB50-\uFDCF\uFDF0-\uFDFF\uFE70-\uFEFF]+', s)) + 110 len(re.findall(r'[a-zA-Z]+', s)) for s in current_chunk)111 else:112 current_chunk = [sentence]113 current_length = sentence_length114 else:115 current_chunk.append(sentence)116 current_length += sentence_length117 118 # Add remaining chunk119 if current_chunk:120 chunk_text = '. '.join(current_chunk) + '.'121 chunks.append({122 'text': chunk_text,123 'sentence_count': len(current_chunk),124 'word_count': current_length,125 'arabic_words': len(re.findall(r'[\u0600-\u06FF\u0750-\u077F\u08A0-\u08FF\uFB50-\uFDCF\uFDF0-\uFDFF\uFE70-\uFEFF]+', chunk_text)),126 'language': 'mixed' if re.search(r'[a-zA-Z]', chunk_text) and re.search(r'[\u0600-\u06FF]', chunk_text) else 'arabic'127 })128 129 return chunks130 131 @staticmethod132 def sliding_window_chunking(text: str, chunk_size: int = 500, overlap: int = 50) -> List[Dict]:133 """Traditional sliding window chunking with character position tracking."""134 words = text.split()135 chunks = []136 137 for i in range(0, len(words), chunk_size - overlap):138 chunk_words = words[i:i + chunk_size]139 chunk_text = ' '.join(chunk_words)140 141 # Calculate character positions142 start_char = len(' '.join(words[:i]))143 end_char = start_char + len(chunk_text)144 if i > 0:145 start_char += 1 # Account for space146 147 chunks.append({148 'text': chunk_text,149 'start_word': i,150 'end_word': min(i + chunk_size, len(words)),151 'start_char': start_char,152 'end_char': end_char,153 'word_count': len(chunk_words),154 'char_count': len(chunk_text)155 })156 157 if i + chunk_size >= len(words):158 break159 160 return chunks161 162class HybridRetriever:163 """Hybrid retrieval combining semantic and lexical search."""164 165 def __init__(self, language='arabic'):166 self.tfidf_vectorizer = None167 self.tfidf_matrix = None168 self.language = language169 170 # Arabic stop words (common words to ignore)171 self.arabic_stop_words = {172 'في', 'من', 'إلى', 'على', 'عن', 'مع', 'أن', 'إن', 'كان', 'لكن',173 'هذا', 'هذه', 'ذلك', 'تلك', 'التي', 'الذي', 'حيث', 'بعد', 'قبل',174 'أو', 'لا', 'نعم', 'كل', 'بعض', 'جميع', 'كذلك', 'أيضا', 'فقط',175 'هي', 'هو', 'أنا', 'أنت', 'نحن', 'أنتم', 'هم', 'هن', 'له', 'لها'176 }177 178 if SKLEARN_AVAILABLE:179 # Configure for Arabic text180 self.tfidf_vectorizer = TfidfVectorizer(181 max_features=10000,182 stop_words=list(self.arabic_stop_words) if language == 'arabic' else 'english',183 ngram_range=(1, 3), # Include trigrams for Arabic184 min_df=2,185 max_df=0.95,186 analyzer='word',187 token_pattern=r'[\u0600-\u06FF\u0750-\u077F\u08A0-\u08FF\uFB50-\uFDCF\uFDF0-\uFDFF\uFE70-\uFEFF\w]+', # Arabic Unicode ranges188 lowercase=True,189 strip_accents='unicode'190 )191 192 def fit_lexical(self, texts: List[str]):193 """Fit TF-IDF vectorizer on document corpus."""194 if not SKLEARN_AVAILABLE:195 logger.warning("scikit-learn not available, lexical search disabled")196 return197 198 try:199 self.tfidf_matrix = self.tfidf_vectorizer.fit_transform(texts)200 logger.info(f"TF-IDF matrix shape: {self.tfidf_matrix.shape}")201 except Exception as e:202 logger.error(f"Error fitting TF-IDF vectorizer: {e}")203 204 def lexical_search(self, query: str, k: int = 10) -> List[Tuple[int, float]]:205 """Perform lexical search using TF-IDF."""206 if not SKLEARN_AVAILABLE or self.tfidf_matrix is None:207 return []208 209 try:210 query_vector = self.tfidf_vectorizer.transform([query])211 similarities = cosine_similarity(query_vector, self.tfidf_matrix).flatten()212 213 # Get top k indices with scores214 top_indices = np.argsort(similarities)[::-1][:k]215 return [(idx, similarities[idx]) for idx in top_indices if similarities[idx] > 0]216 217 except Exception as e:218 logger.error(f"Error in lexical search: {e}")219 return []220 221 def combine_scores(self, semantic_results: List[Tuple[int, float]], 222 lexical_results: List[Tuple[int, float]], 223 semantic_weight: float = 0.7) -> List[Tuple[int, float]]:224 """Combine semantic and lexical scores using weighted fusion."""225 # Normalize scores to [0, 1] range226 def normalize_scores(results):227 if not results:228 return {}229 scores = [score for _, score in results]230 if max(scores) == min(scores):231 return {idx: 1.0 for idx, _ in results}232 233 max_score, min_score = max(scores), min(scores)234 return {idx: (score - min_score) / (max_score - min_score) 235 for idx, score in results}236 237 semantic_scores = normalize_scores(semantic_results)238 lexical_scores = normalize_scores(lexical_results)239 240 # Combine scores241 all_indices = set(semantic_scores.keys()) | set(lexical_scores.keys())242 combined_results = []243 244 for idx in all_indices:245 semantic_score = semantic_scores.get(idx, 0.0)246 lexical_score = lexical_scores.get(idx, 0.0)247 248 # Weighted combination249 hybrid_score = (semantic_weight * semantic_score + 250 (1 - semantic_weight) * lexical_score)251 252 combined_results.append((idx, hybrid_score, semantic_score, lexical_score))253 254 # Sort by hybrid score255 combined_results.sort(key=lambda x: x[1], reverse=True)256 257 return combined_results258 259class RAGSystem:260 """Enhanced RAG system with hybrid retrieval and advanced features."""261 262 def __init__(self, 263 gemini_client,264 data_dir: str = 'data',265 storage_dir: str = 'storage',266 chunk_size: int = 500,267 overlap: int = 50,268 chunking_strategy: str = 'semantic',269 embedding_model: str = 'gemini'):270 """271 Initialize enhanced RAG system.272 273 Args:274 gemini_client: Gemini client for embeddings275 data_dir: Directory containing source documents276 storage_dir: Directory for storing indexes and metadata277 chunk_size: Size of text chunks in words278 overlap: Overlap between chunks in words279 chunking_strategy: 'semantic' or 'sliding_window'280 embedding_model: 'gemini' or 'sentence_transformer'281 """282 self.gemini_client = gemini_client283 self.data_dir = Path(data_dir)284 self.storage_dir = Path(storage_dir)285 self.chunk_size = chunk_size286 self.overlap = overlap287 self.chunking_strategy = chunking_strategy288 self.embedding_model = embedding_model289 290 # Ensure directories exist291 self.data_dir.mkdir(exist_ok=True)292 self.storage_dir.mkdir(exist_ok=True)293 294 # Initialize components295 self.chunks: List[DocumentChunk] = []296 self.embeddings: Optional[np.ndarray] = None297 self.faiss_index = None298 self.hybrid_retriever = HybridRetriever(language='arabic') # Configure for Arabic299 self.chunker = AdvancedChunker()300 301 # Query expansion for Arabic302 self.arabic_synonyms = {303 'سياحة': ['زيارة', 'رحلة', 'سفر', 'جولة'],304 'فندق': ['منتجع', 'نزل', 'سكن', 'إقامة'],305 'طعام': ['مأكل', 'وجبة', 'أكل', 'مطعم'],306 'تسوق': ['شراء', 'بضائع', 'سوق', 'مول'],307 'ترفيه': ['تسلية', 'لعب', 'نشاط', 'متعة'],308 'شاطئ': ['بحر', 'ساحل', 'رمل'],309 'جبل': ['تل', 'مرتفع', 'قمة'],310 'مدينة': ['بلدة', 'عاصمة', 'منطقة']311 }312 313 # Threading lock for concurrent access314 self.lock = threading.Lock()315 316 # Initialize sentence transformer if available317 self.sentence_model = None318 if embedding_model == 'sentence_transformer' and SENTENCE_TRANSFORMERS_AVAILABLE:319 try:320 self.sentence_model = SentenceTransformer('all-MiniLM-L6-v2')321 logger.info("Initialized sentence transformer model")322 except Exception as e:323 logger.error(f"Failed to load sentence transformer: {e}")324 325 # Embedding dimensions326 self.embedding_dim = 768 if embedding_model == 'gemini' else 384327 328 def load_documents(self) -> List[DocumentChunk]:329 """Load and process documents with enhanced chunking."""330 chunks = []331 332 logger.info(f"Loading documents from {self.data_dir}")333 334 # Support multiple file formats335 file_patterns = ['*.txt', '*.md', '*.json']336 all_files = []337 for pattern in file_patterns:338 all_files.extend(self.data_dir.glob(pattern))339 340 for file_path in all_files:341 try:342 with open(file_path, 'r', encoding='utf-8') as f:343 if file_path.suffix == '.json':344 data = json.load(f)345 # Extract text content from JSON346 if isinstance(data, dict):347 content = ' '.join(str(v) for v in data.values() if isinstance(v, str))348 elif isinstance(data, list):349 content = ' '.join(str(item) for item in data if isinstance(item, str))350 else:351 content = str(data)352 else:353 content = f.read().strip()354 355 if not content:356 logger.warning(f"Empty file: {file_path}")357 continue358 359 # Choose chunking strategy360 if self.chunking_strategy == 'semantic':361 raw_chunks = self.chunker.semantic_chunking_arabic(content, self.chunk_size, self.overlap)362 else:363 raw_chunks = self.chunker.sliding_window_chunking(content, self.chunk_size, self.overlap)364 365 # Convert to DocumentChunk objects366 for i, chunk_data in enumerate(raw_chunks):367 chunk = DocumentChunk(368 text=chunk_data['text'],369 source_file=file_path.name,370 chunk_id=i,371 start_char=chunk_data.get('start_char', 0),372 end_char=chunk_data.get('end_char', 0),373 start_word=chunk_data.get('start_word', 0),374 end_word=chunk_data.get('end_word', 0),375 word_count=chunk_data.get('word_count', 0),376 char_count=chunk_data.get('char_count', len(chunk_data['text'])),377 hash_id='',378 metadata={'file_type': file_path.suffix}379 )380 chunks.append(chunk)381 382 logger.info(f"Processed {len(raw_chunks)} chunks from {file_path.name}")383 384 except Exception as e:385 logger.error(f"Error processing {file_path}: {e}")386 continue387 388 logger.info(f"Total chunks loaded: {len(chunks)}")389 return chunks390 391 def create_embeddings_batch(self, texts: List[str], batch_size: int = 32) -> np.ndarray:392 """Create embeddings in batches for better performance."""393 embeddings = []394 395 for i in range(0, len(texts), batch_size):396 batch_texts = texts[i:i + batch_size]397 398 try:399 if self.embedding_model == 'sentence_transformer' and self.sentence_model:400 batch_embeddings = self.sentence_model.encode(batch_texts, convert_to_numpy=True)401 else:402 # Use Gemini embeddings403 batch_embeddings = self.gemini_client.embed_texts(batch_texts)404 batch_embeddings = np.array(batch_embeddings)405 406 embeddings.append(batch_embeddings)407 logger.info(f"Processed embedding batch {i//batch_size + 1}/{(len(texts) + batch_size - 1)//batch_size}")408 409 except Exception as e:410 logger.error(f"Error creating embeddings for batch {i//batch_size + 1}: {e}")411 # Fallback to random embeddings412 fallback_embeddings = np.random.rand(len(batch_texts), self.embedding_dim).astype(np.float32)413 embeddings.append(fallback_embeddings)414 415 return np.vstack(embeddings) if embeddings else np.array([])416 417 def build_faiss_index(self, embeddings: np.ndarray):418 """Build optimized FAISS index."""419 if not FAISS_AVAILABLE:420 logger.warning("FAISS not available")421 return422 423 try:424 dimension = embeddings.shape[1]425 426 # Use more sophisticated index for better retrieval427 if embeddings.shape[0] > 1000:428 # Use HNSW index for larger datasets429 self.faiss_index = faiss.IndexHNSWFlat(dimension, 32)430 self.faiss_index.hnsw.efConstruction = 200431 self.faiss_index.hnsw.efSearch = 50432 else:433 # Use flat index for smaller datasets434 self.faiss_index = faiss.IndexFlatIP(dimension)435 436 # Normalize embeddings for cosine similarity437 embeddings_normalized = embeddings.copy()438 faiss.normalize_L2(embeddings_normalized)439 440 # Add to index441 self.faiss_index.add(embeddings_normalized.astype(np.float32))442 443 logger.info(f"Built FAISS index: {type(self.faiss_index).__name__} with {self.faiss_index.ntotal} vectors")444 445 except Exception as e:446 logger.error(f"Error building FAISS index: {e}")447 self.faiss_index = None448 449 def save_index(self):450 """Save all components to disk."""451 try:452 # Save chunks metadata453 chunks_data = []454 for chunk in self.chunks:455 chunks_data.append({456 'text': chunk.text,457 'source_file': chunk.source_file,458 'chunk_id': chunk.chunk_id,459 'start_char': chunk.start_char,460 'end_char': chunk.end_char,461 'start_word': chunk.start_word,462 'end_word': chunk.end_word,463 'word_count': chunk.word_count,464 'char_count': chunk.char_count,465 'hash_id': chunk.hash_id,466 'metadata': chunk.metadata467 })468 469 with open(self.storage_dir / 'chunks.json', 'w', encoding='utf-8') as f:470 json.dump(chunks_data, f, ensure_ascii=False, indent=2)471 472 # Save embeddings473 if self.embeddings is not None:474 np.save(self.storage_dir / 'embeddings.npy', self.embeddings)475 476 # Save FAISS index477 if self.faiss_index is not None:478 faiss.write_index(self.faiss_index, str(self.storage_dir / 'faiss_index.bin'))479 480 # Save TF-IDF components481 if self.hybrid_retriever.tfidf_vectorizer is not None:482 with open(self.storage_dir / 'tfidf_vectorizer.pkl', 'wb') as f:483 pickle.dump(self.hybrid_retriever.tfidf_vectorizer, f)484 485 with open(self.storage_dir / 'tfidf_matrix.pkl', 'wb') as f:486 pickle.dump(self.hybrid_retriever.tfidf_matrix, f)487 488 logger.info("Index and all components saved successfully")489 490 except Exception as e:491 logger.error(f"Error saving index: {e}")492 raise493 494 def load_index(self) -> bool:495 """Load all components from disk."""496 try:497 chunks_path = self.storage_dir / 'chunks.json'498 embeddings_path = self.storage_dir / 'embeddings.npy'499 500 if not chunks_path.exists():501 logger.info("No existing index found")502 return False503 504 # Load chunks505 with open(chunks_path, 'r', encoding='utf-8') as f:506 chunks_data = json.load(f)507 508 self.chunks = []509 for chunk_data in chunks_data:510 chunk = DocumentChunk(**chunk_data)511 self.chunks.append(chunk)512 513 # Load embeddings514 if embeddings_path.exists():515 self.embeddings = np.load(embeddings_path)516 517 # Load FAISS index518 faiss_path = self.storage_dir / 'faiss_index.bin'519 if FAISS_AVAILABLE and faiss_path.exists():520 self.faiss_index = faiss.read_index(str(faiss_path))521 522 # Load TF-IDF components523 tfidf_vectorizer_path = self.storage_dir / 'tfidf_vectorizer.pkl'524 tfidf_matrix_path = self.storage_dir / 'tfidf_matrix.pkl'525 526 if tfidf_vectorizer_path.exists() and tfidf_matrix_path.exists():527 with open(tfidf_vectorizer_path, 'rb') as f:528 self.hybrid_retriever.tfidf_vectorizer = pickle.load(f)529 530 with open(tfidf_matrix_path, 'rb') as f:531 self.hybrid_retriever.tfidf_matrix = pickle.load(f)532 533 logger.info(f"Loaded index with {len(self.chunks)} chunks")534 return True535 536 except Exception as e:537 logger.error(f"Error loading index: {e}")538 return False539 540 def create_index(self, force: bool = False):541 """Create comprehensive index with all components."""542 if not force and self.load_index():543 logger.info("Index already exists and force=False")544 return545 546 logger.info("Creating enhanced index...")547 548 # Load and process documents549 self.chunks = self.load_documents()550 551 if not self.chunks:552 logger.warning("No documents found to index")553 return554 555 # Extract texts556 texts = [chunk.text for chunk in self.chunks]557 558 # Create embeddings559 logger.info("Creating embeddings...")560 self.embeddings = self.create_embeddings_batch(texts)561 562 # Build FAISS index563 logger.info("Building FAISS index...")564 self.build_faiss_index(self.embeddings)565 566 # Build TF-IDF index567 logger.info("Building TF-IDF index...")568 self.hybrid_retriever.fit_lexical(texts)569 570 # Save all components571 self.save_index()572 573 logger.info("Enhanced index creation completed")574 575 def semantic_search(self, query: str, k: int = 10) -> List[Tuple[int, float]]:576 """Perform semantic search using embeddings."""577 if self.faiss_index is None or self.embeddings is None:578 return []579 580 try:581 # Create query embedding582 if self.embedding_model == 'sentence_transformer' and self.sentence_model:583 query_embedding = self.sentence_model.encode([query], convert_to_numpy=True)584 else:585 query_embeddings = self.gemini_client.embed_texts([query])586 query_embedding = np.array(query_embeddings)587 588 # Normalize query embedding589 faiss.normalize_L2(query_embedding.astype(np.float32))590 591 # Search592 scores, indices = self.faiss_index.search(query_embedding.astype(np.float32), k)593 594 # Return results with scores595 results = []596 for i, (idx, score) in enumerate(zip(indices[0], scores[0])):597 if idx != -1: # Valid index598 results.append((int(idx), float(score)))599 600 return results601 602 except Exception as e:603 logger.error(f"Error in semantic search: {e}")604 return []605 606 def expand_query_arabic(self, query: str) -> str:607 """Expand Arabic query with synonyms for better retrieval."""608 expanded_terms = []609 words = re.findall(r'[\u0600-\u06FF\u0750-\u077F\u08A0-\u08FF\uFB50-\uFDCF\uFDF0-\uFDFF\uFE70-\uFEFF]+', query)610 611 for word in query.split():612 expanded_terms.append(word)613 # Add synonyms if available614 if word in self.arabic_synonyms:615 expanded_terms.extend(self.arabic_synonyms[word])616 617 return ' '.join(expanded_terms)618 619 def hybrid_retrieve(self, query: str, k: int = 5, semantic_weight: float = 0.7, expand_query: bool = True) -> List[RetrievalResult]:620 """Perform hybrid retrieval combining semantic and lexical search."""621 with self.lock:622 try:623 # Expand query for Arabic if enabled624 if expand_query:625 expanded_query = self.expand_query_arabic(query)626 else:627 expanded_query = query628 629 # Perform semantic search with original query630 semantic_results = self.semantic_search(query, k * 2)631 632 # Perform lexical search with expanded query633 lexical_results = self.hybrid_retriever.lexical_search(expanded_query, k * 2)634 635 # Combine results636 combined_results = self.hybrid_retriever.combine_scores(637 semantic_results, lexical_results, semantic_weight638 )639 640 # Create RetrievalResult objects641 results = []642 for rank, (idx, hybrid_score, semantic_score, lexical_score) in enumerate(combined_results[:k]):643 if 0 <= idx < len(self.chunks):644 result = RetrievalResult(645 chunk=self.chunks[idx],646 semantic_score=semantic_score,647 lexical_score=lexical_score,648 hybrid_score=hybrid_score,649 rank=rank650 )651 results.append(result)652 653 logger.info(f"Retrieved {len(results)} results for query: {query}")654 return results655 656 except Exception as e:657 logger.error(f"Error in hybrid retrieval: {e}")658 return []659 660 def retrieve_with_context(self, query: str, k: int = 5, 661 expand_context: bool = True) -> List[Dict]:662 """Retrieve documents with expanded context from adjacent chunks."""663 results = self.hybrid_retrieve(query, k)664 665 if not expand_context:666 return [self._result_to_dict(result) for result in results]667 668 # Group results by source file669 file_groups = defaultdict(list)670 for result in results:671 file_groups[result.chunk.source_file].append(result)672 673 # Expand context for each result674 expanded_results = []675 for result in results:676 context_chunks = [result.chunk]677 678 # Find adjacent chunks679 same_file_chunks = [c for c in self.chunks 680 if c.source_file == result.chunk.source_file]681 same_file_chunks.sort(key=lambda x: x.chunk_id)682 683 current_idx = next((i for i, c in enumerate(same_file_chunks) 684 if c.hash_id == result.chunk.hash_id), -1)685 686 if current_idx != -1:687 # Add previous and next chunks688 if current_idx > 0:689 context_chunks.insert(0, same_file_chunks[current_idx - 1])690 if current_idx < len(same_file_chunks) - 1:691 context_chunks.append(same_file_chunks[current_idx + 1])692 693 # Combine context694 expanded_text = ' '.join(chunk.text for chunk in context_chunks)695 696 result_dict = self._result_to_dict(result)697 result_dict['expanded_text'] = expanded_text698 result_dict['context_chunks'] = len(context_chunks)699 700 expanded_results.append(result_dict)701 702 return expanded_results703 704 def _result_to_dict(self, result: RetrievalResult) -> Dict:705 """Convert RetrievalResult to dictionary."""706 return {707 'text': result.chunk.text,708 'source_file': result.chunk.source_file,709 'chunk_id': result.chunk.chunk_id,710 'semantic_score': result.semantic_score,711 'lexical_score': result.lexical_score,712 'hybrid_score': result.hybrid_score,713 'rank': result.rank,714 'word_count': result.chunk.word_count,715 'char_count': result.chunk.char_count,716 'metadata': result.chunk.metadata717 }718 719 def get_stats(self) -> Dict:720 """Get comprehensive statistics about the RAG system."""721 stats = {722 'total_chunks': len(self.chunks),723 'total_files': len(set(chunk.source_file for chunk in self.chunks)),724 'embedding_model': self.embedding_model,725 'chunking_strategy': self.chunking_strategy,726 'chunk_size': self.chunk_size,727 'overlap': self.overlap,728 'faiss_available': FAISS_AVAILABLE,729 'sklearn_available': SKLEARN_AVAILABLE,730 'sentence_transformers_available': SENTENCE_TRANSFORMERS_AVAILABLE,731 'faiss_index_type': type(self.faiss_index).__name__ if self.faiss_index else None,732 'embedding_dimension': self.embedding_dim,733 'tfidf_features': self.hybrid_retriever.tfidf_matrix.shape[1] if self.hybrid_retriever.tfidf_matrix is not None else 0734 }735 736 if self.chunks:737 word_counts = [chunk.word_count for chunk in self.chunks]738 stats.update({739 'avg_chunk_words': np.mean(word_counts),740 'min_chunk_words': min(word_counts),741 'max_chunk_words': max(word_counts),742 'median_chunk_words': np.median(word_counts)743 })744 745 return stats746 747 748# Example usage and initialization749def initialize_enhanced_rag(gemini_client, **kwargs):750 """Initialize enhanced RAG system with optimal settings."""751 return RAGSystem(752 gemini_client=gemini_client,753 chunking_strategy='semantic', # Use semantic chunking by default754 embedding_model='gemini',755 **kwargs756 )