alschameri/helping-source

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 pickle10from typing import List, Dict, Optional11import numpy as np12from pathlib import Path13 14try:15    import faiss16    FAISS_AVAILABLE = True17except ImportError:18    FAISS_AVAILABLE = False19 20logger = logging.getLogger(__name__)21 22class RAGSystem:23    """RAG system for document indexing and retrieval."""24    25    def __init__(self, gemini_client, data_dir='data', storage_dir='storage', chunk_size=500, overlap=50):26        self.gemini_client = gemini_client27        self.data_dir = Path(data_dir)28        self.storage_dir = Path(storage_dir)29        self.chunk_size = chunk_size30        self.overlap = overlap31        32        # Ensure directories exist33        self.data_dir.mkdir(exist_ok=True)34        self.storage_dir.mkdir(exist_ok=True)35        36        # Initialize index and metadata37        self.index = None38        self.metadata = []39        self.embedding_dim = 768  # Default Gemini embedding dimension40        41    def chunk_text(self, text: str, chunk_size: int = None, overlap: int = None) -> List[Dict]:42        """Split text into overlapping chunks."""43        if chunk_size is None:44            chunk_size = self.chunk_size45        if overlap is None:46            overlap = self.overlap47            48        # Simple word-based chunking49        words = text.split()50        chunks = []51        52        for i in range(0, len(words), chunk_size - overlap):53            chunk_words = words[i:i + chunk_size]54            chunk_text = ' '.join(chunk_words)55            56            chunks.append({57                'text': chunk_text,58                'start_word': i,59                'end_word': min(i + chunk_size, len(words)),60                'word_count': len(chunk_words)61            })62            63            # Break if we've reached the end64            if i + chunk_size >= len(words):65                break66                67        return chunks68    69    def load_documents(self) -> List[Dict]:70        """Load and chunk all text documents from data directory."""71        documents = []72        73        logger.info(f"Loading documents from {self.data_dir}")74        75        for file_path in self.data_dir.glob('*.txt'):76            try:77                with open(file_path, 'r', encoding='utf-8') as f:78                    content = f.read().strip()79                80                if not content:81                    logger.warning(f"Empty file: {file_path}")82                    continue83                84                # Chunk the document85                chunks = self.chunk_text(content)86                87                for i, chunk in enumerate(chunks):88                    documents.append({89                        'source_file': file_path.name,90                        'chunk_id': i,91                        'text': chunk['text'],92                        'start_word': chunk['start_word'],93                        'end_word': chunk['end_word'],94                        'word_count': chunk['word_count']95                    })96                97                logger.info(f"Loaded {len(chunks)} chunks from {file_path.name}")98                99            except Exception as e:100                logger.error(f"Error loading {file_path}: {e}")101                continue102        103        logger.info(f"Total documents loaded: {len(documents)}")104        return documents105    106    def create_embeddings(self, texts: List[str]) -> np.ndarray:107        """Create embeddings for list of texts using Gemini."""108        try:109            embeddings = self.gemini_client.embed_texts(texts)110            return np.array(embeddings)111        except Exception as e:112            logger.error(f"Error creating embeddings: {e}")113            # Fallback to random embeddings for development114            logger.warning("Using random embeddings as fallback")115            return np.random.rand(len(texts), self.embedding_dim).astype(np.float32)116    117    def build_index(self, embeddings: np.ndarray):118        """Build FAISS index from embeddings."""119        if not FAISS_AVAILABLE:120            logger.warning("FAISS not available, using simple similarity search")121            self.index = embeddings122            return123        124        # Create FAISS index125        dimension = embeddings.shape[1]126        self.index = faiss.IndexFlatIP(dimension)  # Inner product for cosine similarity127        128        # Normalize embeddings for cosine similarity129        faiss.normalize_L2(embeddings)130        self.index.add(embeddings)131        132        logger.info(f"Built FAISS index with {self.index.ntotal} vectors")133    134    def save_index(self):135        """Save index and metadata to disk."""136        try:137            # Save metadata138            metadata_path = self.storage_dir / 'metadata.json'139            with open(metadata_path, 'w', encoding='utf-8') as f:140                json.dump(self.metadata, f, ensure_ascii=False, indent=2)141            142            # Save index143            if FAISS_AVAILABLE and hasattr(self.index, 'ntotal'):144                index_path = self.storage_dir / 'index.faiss'145                faiss.write_index(self.index, str(index_path))146            else:147                # Save numpy array as pickle148                index_path = self.storage_dir / 'index.pkl'149                with open(index_path, 'wb') as f:150                    pickle.dump(self.index, f)151            152            logger.info("Index and metadata saved successfully")153            154        except Exception as e:155            logger.error(f"Error saving index: {e}")156            raise157    158    def load_index(self) -> bool:159        """Load existing index and metadata from disk."""160        try:161            metadata_path = self.storage_dir / 'metadata.json'162            163            # Check if files exist164            if not metadata_path.exists():165                logger.info("No existing index found")166                return False167            168            # Load metadata169            with open(metadata_path, 'r', encoding='utf-8') as f:170                self.metadata = json.load(f)171            172            # Load index173            if FAISS_AVAILABLE:174                index_path = self.storage_dir / 'index.faiss'175                if index_path.exists():176                    self.index = faiss.read_index(str(index_path))177                else:178                    return False179            else:180                index_path = self.storage_dir / 'index.pkl'181                if index_path.exists():182                    with open(index_path, 'rb') as f:183                        self.index = pickle.load(f)184                else:185                    return False186            187            logger.info(f"Loaded index with {len(self.metadata)} documents")188            return True189            190        except Exception as e:191            logger.error(f"Error loading index: {e}")192            return False193    194    def create_index(self, force: bool = False):195        """Create new index from documents in data directory."""196        if not force and self.load_index():197            logger.info("Index already exists and force=False")198            return199        200        logger.info("Creating new index...")201        202        # Load documents203        documents = self.load_documents()204        205        if not documents:206            logger.warning("No documents found to index")207            return208        209        # Extract texts for embedding210        texts = [doc['text'] for doc in documents]211        212        # Create embeddings213        logger.info("Creating embeddings...")214        embeddings = self.create_embeddings(texts)215        216        # Build index217        logger.info("Building search index...")218        self.build_index(embeddings)219        220        # Store metadata221        self.metadata = documents222        223        # Save to disk224        self.save_index()225        226        logger.info("Index creation completed")227    228    def simple_similarity_search(self, query_embedding: np.ndarray, k: int = 4) -> List[int]:229        """Simple similarity search when FAISS is not available."""230        if self.index is None:231            return []232        233        # Compute cosine similarity234        similarities = np.dot(self.index, query_embedding.T).flatten()235        236        # Get top k indices237        top_indices = np.argsort(similarities)[::-1][:k]238        return top_indices.tolist()239    240    def search(self, query_embedding: np.ndarray, k: int = 4) -> List[int]:241        """Search for similar documents."""242        if self.index is None:243            logger.warning("No index available for search")244            return []245        246        try:247            if FAISS_AVAILABLE and hasattr(self.index, 'search'):248                # FAISS search249                query_embedding = query_embedding.reshape(1, -1).astype(np.float32)250                faiss.normalize_L2(query_embedding)251                252                scores, indices = self.index.search(query_embedding, k)253                return indices[0].tolist()254            else:255                # Simple similarity search256                return self.simple_similarity_search(query_embedding, k)257                258        except Exception as e:259            logger.error(f"Error during search: {e}")260            return []261    262    def retrieve(self, query: str, k: int = 4) -> List[Dict]:263        """Retrieve relevant documents for a query."""264        try:265            # Create query embedding266            query_embeddings = self.create_embeddings([query])267            query_embedding = query_embeddings[0]268            269            # Search for similar documents270            indices = self.search(query_embedding, k)271            272            # Return relevant documents with metadata273            results = []274            for idx in indices:275                if 0 <= idx < len(self.metadata):276                    results.append(self.metadata[idx])277            278            logger.info(f"Retrieved {len(results)} documents for query")279            return results280            281        except Exception as e:282            logger.error(f"Error during retrieval: {e}")283            return []284 285# # Initialize services when module is imported286# initialize_services()