""" RAG Engine Module ================= Handles all RAG pipeline operations: - PDF text extraction - Text chunking with overlap - Embedding generation using SentenceTransformers - FAISS vector storage and retrieval - Metadata and document registry management - Persistence of embeddings and metadata """ import os import json import hashlib from datetime import datetime from typing import List, Dict, Tuple, Optional import numpy as np import faiss from sentence_transformers import SentenceTransformer import PyPDF2 import google.generativeai as genai from PIL import Image import io # OCR imports (optional) try: import pytesseract OCR_AVAILABLE = True except ImportError: OCR_AVAILABLE = False print("Warning: pytesseract not installed. OCR functionality will be disabled.") # ============================================ # CONFIGURATION # ============================================ # Chunking parameters DEFAULT_CHUNK_SIZE = 200 # words per chunk DEFAULT_OVERLAP_SIZE = 50 # overlapping words # Retrieval parameters DEFAULT_TOP_K = 5 # number of chunks to retrieve # Embedding model EMBEDDING_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2" EMBEDDING_DIMENSION = 384 class RAGEngine: """ Main RAG Engine class that handles: - Document processing and embedding - FAISS index management - Query processing and answer generation - Persistence of all data """ def __init__(self, gemini_api_key: str, storage_dir: Optional[str] = None): """ Initialize the RAG Engine. Args: gemini_api_key: API key for Google Gemini storage_dir: Optional custom storage directory for per-user isolation """ # Set storage paths if storage_dir is None: storage_dir = os.path.join(os.path.dirname(__file__), "storage") self.storage_dir = storage_dir self.faiss_index_path = os.path.join(storage_dir, "faiss.index") self.metadata_path = os.path.join(storage_dir, "metadata.json") self.documents_path = os.path.join(storage_dir, "documents.json") # Ensure storage directory exists os.makedirs(storage_dir, exist_ok=True) # Initialize embedding model print("Loading embedding model...") self.embed_model = SentenceTransformer(EMBEDDING_MODEL_NAME) # Initialize Gemini genai.configure(api_key=gemini_api_key) self.gemini_model = genai.GenerativeModel("gemini-2.5-flash") # Initialize or load FAISS index self.index: Optional[faiss.IndexFlatL2] = None self.metadata: List[Dict] = [] # Stores chunk text, source, page self.documents: Dict[str, Dict] = {} # Document registry # Load existing data if available self._load_persistent_data() print(f"RAG Engine initialized. Documents: {len(self.documents)}, Chunks: {len(self.metadata)}") # ============================================ # PERSISTENCE METHODS # ============================================ def _load_persistent_data(self): """Load FAISS index, metadata, and document registry from disk.""" # Load document registry if os.path.exists(self.documents_path): with open(self.documents_path, "r", encoding="utf-8") as f: self.documents = json.load(f) print(f"Loaded {len(self.documents)} documents from registry") # Load metadata if os.path.exists(self.metadata_path): with open(self.metadata_path, "r", encoding="utf-8") as f: self.metadata = json.load(f) print(f"Loaded {len(self.metadata)} chunks metadata") # Load FAISS index if os.path.exists(self.faiss_index_path) and len(self.metadata) > 0: self.index = faiss.read_index(self.faiss_index_path) print(f"Loaded FAISS index with {self.index.ntotal} vectors") else: # Create new empty index self.index = faiss.IndexFlatL2(EMBEDDING_DIMENSION) print("Created new FAISS index") def _save_persistent_data(self): """Save FAISS index, metadata, and document registry to disk.""" # Save document registry with open(self.documents_path, "w", encoding="utf-8") as f: json.dump(self.documents, f, indent=2, ensure_ascii=False) # Save metadata with open(self.metadata_path, "w", encoding="utf-8") as f: json.dump(self.metadata, f, indent=2, ensure_ascii=False) # Save FAISS index if self.index is not None and self.index.ntotal > 0: faiss.write_index(self.index, self.faiss_index_path) print("Persistent data saved successfully") # ============================================ # DOCUMENT PROCESSING METHODS # ============================================ @staticmethod def compute_file_hash(file_content: bytes) -> str: """ Compute SHA-256 hash of file content. Args: file_content: Raw bytes of the file Returns: Hexadecimal hash string """ return hashlib.sha256(file_content).hexdigest() @staticmethod def chunk_text_with_overlap(text: str, chunk_size: int = DEFAULT_CHUNK_SIZE, overlap_size: int = DEFAULT_OVERLAP_SIZE) -> List[str]: """ Split text into overlapping chunks. Args: text: Input text to chunk chunk_size: Number of words per chunk overlap_size: Number of overlapping words between chunks Returns: List of text chunks """ words = text.split() chunks = [] start = 0 while start < len(words): end = start + chunk_size chunk = " ".join(words[start:end]) if chunk.strip(): # Only add non-empty chunks chunks.append(chunk) start += chunk_size - overlap_size return chunks @staticmethod def extract_text_from_image(image: Image.Image) -> str: """ Extract text from an image using OCR. Args: image: PIL Image object Returns: Extracted text string """ if not OCR_AVAILABLE: return "" try: # Convert to RGB if needed if image.mode != 'RGB': image = image.convert('RGB') # Run OCR text = pytesseract.image_to_string(image, lang='eng') return text.strip() except Exception as e: print(f"OCR error: {e}") return "" def extract_text_from_pdf(self, pdf_content: bytes) -> List[Dict]: """ Extract text from PDF page by page, including OCR for images. Args: pdf_content: Raw bytes of PDF file Returns: List of dicts with page_num, text, and ocr_text """ pages = [] try: reader = PyPDF2.PdfReader(io.BytesIO(pdf_content)) for page_num, page in enumerate(reader.pages): # Extract regular text text = page.extract_text() ocr_text = "" # Extract images and apply OCR if OCR_AVAILABLE: try: # Get images from page if '/XObject' in page['/Resources']: xObject = page['/Resources']['/XObject'].get_object() for obj in xObject: if xObject[obj]['/Subtype'] == '/Image': try: # Extract image data size = (xObject[obj]['/Width'], xObject[obj]['/Height']) data = xObject[obj].get_data() # Try to create image if xObject[obj]['/ColorSpace'] == '/DeviceRGB': mode = "RGB" elif xObject[obj]['/ColorSpace'] == '/DeviceGray': mode = "L" else: mode = "RGB" # Default try: image = Image.frombytes(mode, size, data) # Apply OCR img_text = self.extract_text_from_image(image) if img_text: ocr_text += img_text + "\n" except Exception as img_error: # Try with PIL's open if frombytes fails try: image = Image.open(io.BytesIO(data)) img_text = self.extract_text_from_image(image) if img_text: ocr_text += img_text + "\n" except: pass except Exception as e: # Skip this image if extraction fails continue except Exception as e: print(f"Error extracting images from page {page_num + 1}: {e}") # Combine regular text and OCR text combined_text = "" if text and text.strip(): combined_text += text.strip() if ocr_text.strip(): if combined_text: combined_text += "\n\n[Text from images:]\n" + ocr_text.strip() else: combined_text = ocr_text.strip() if combined_text: pages.append({ "page_num": page_num + 1, "text": combined_text, "has_ocr": bool(ocr_text.strip()) }) except Exception as e: print(f"Error extracting PDF text: {e}") raise return pages def process_pdf(self, filename: str, file_content: bytes, chunk_size: int = DEFAULT_CHUNK_SIZE, overlap_size: int = DEFAULT_OVERLAP_SIZE) -> List[Dict]: """ Process a PDF: extract text (including OCR), chunk it, and prepare metadata. Args: filename: Original filename file_content: Raw bytes of PDF chunk_size: Words per chunk overlap_size: Overlap between chunks Returns: List of chunk metadata dicts """ # Extract pages pages = self.extract_text_from_pdf(file_content) # Chunk each page chunks_metadata = [] for page_info in pages: page_chunks = self.chunk_text_with_overlap( page_info["text"], chunk_size, overlap_size ) for chunk_text in page_chunks: chunks_metadata.append({ "text": chunk_text, "source": filename, "page": page_info["page_num"], "has_ocr": page_info.get("has_ocr", False) }) return chunks_metadata # ============================================ # DUPLICATE DETECTION METHODS # ============================================ def check_duplicate(self, file_hash: str) -> Optional[Dict]: """ Check if a document with the same hash already exists. Args: file_hash: SHA-256 hash of the file Returns: Document info if duplicate found, None otherwise """ for doc_id, doc_info in self.documents.items(): if doc_info.get("hash") == file_hash: return {"doc_id": doc_id, **doc_info} return None def get_document_by_filename(self, filename: str) -> Optional[Dict]: """ Get document info by filename. Args: filename: Original filename Returns: Document info if found, None otherwise """ for doc_id, doc_info in self.documents.items(): if doc_info.get("filename") == filename: return {"doc_id": doc_id, **doc_info} return None # ============================================ # EMBEDDING AND INDEXING METHODS # ============================================ def generate_embeddings(self, texts: List[str]) -> np.ndarray: """ Generate embeddings for a list of texts. Args: texts: List of text strings Returns: Numpy array of embeddings """ embeddings = self.embed_model.encode(texts) return np.array(embeddings).astype("float32") def add_to_index(self, chunks_metadata: List[Dict]) -> int: """ Add new chunks to FAISS index and metadata. Args: chunks_metadata: List of chunk dicts with text, source, page Returns: Number of chunks added """ if not chunks_metadata: return 0 # Extract texts for embedding texts = [c["text"] for c in chunks_metadata] # Generate embeddings embeddings = self.generate_embeddings(texts) # Add to FAISS index self.index.add(embeddings) # Add to metadata self.metadata.extend(chunks_metadata) return len(chunks_metadata) def remove_document_from_index(self, filename: str): """ Remove all chunks of a document from the index. Note: FAISS IndexFlatL2 doesn't support removal, so we rebuild. Args: filename: Filename of document to remove """ # Filter out chunks from this document remaining_metadata = [ m for m in self.metadata if m["source"] != filename ] if len(remaining_metadata) == len(self.metadata): return # Nothing to remove # Rebuild index with remaining chunks self.metadata = remaining_metadata if self.metadata: texts = [m["text"] for m in self.metadata] embeddings = self.generate_embeddings(texts) self.index = faiss.IndexFlatL2(EMBEDDING_DIMENSION) self.index.add(embeddings) else: self.index = faiss.IndexFlatL2(EMBEDDING_DIMENSION) print(f"Removed document '{filename}' from index") # ============================================ # DOCUMENT UPLOAD METHODS # ============================================ def upload_document(self, filename: str, file_content: bytes, action: str = "auto") -> Dict: """ Upload and process a document. Args: filename: Original filename file_content: Raw bytes of PDF action: "auto", "use_existing", "replace", or "cancel" Returns: Result dict with status and info """ # Compute hash file_hash = self.compute_file_hash(file_content) # Check for duplicate existing_doc = self.check_duplicate(file_hash) if existing_doc: if action == "auto": # Return duplicate warning return { "status": "duplicate", "filename": filename, "existing_filename": existing_doc["filename"], "hash": file_hash, "message": f"Document already exists as '{existing_doc['filename']}'", "options": ["use_existing", "replace", "cancel"] } elif action == "use_existing": return { "status": "success", "filename": existing_doc["filename"], "message": "Using existing document embeddings", "chunks": 0, "reused": True } elif action == "cancel": return { "status": "cancelled", "filename": filename, "message": "Upload cancelled" } elif action == "replace": # Remove old document and continue with upload self.remove_document_from_index(existing_doc["filename"]) del self.documents[existing_doc["doc_id"]] # Process new document try: chunks_metadata = self.process_pdf(filename, file_content) if not chunks_metadata: return { "status": "error", "filename": filename, "message": "No text could be extracted from PDF" } # Add to index num_chunks = self.add_to_index(chunks_metadata) # Register document doc_id = f"doc_{len(self.documents) + 1}_{int(datetime.now().timestamp())}" self.documents[doc_id] = { "filename": filename, "hash": file_hash, "upload_timestamp": datetime.now().isoformat(), "num_chunks": num_chunks, "num_pages": max(c["page"] for c in chunks_metadata) } # Persist changes self._save_persistent_data() return { "status": "success", "filename": filename, "message": f"Document processed successfully", "chunks": num_chunks, "pages": self.documents[doc_id]["num_pages"] } except Exception as e: return { "status": "error", "filename": filename, "message": f"Error processing document: {str(e)}" } # ============================================ # QUERY AND RETRIEVAL METHODS # ============================================ def retrieve_relevant_chunks(self, query: str, top_k: int = DEFAULT_TOP_K, doc_id: str = None) -> List[Dict]: """ Retrieve most relevant chunks for a query. Args: query: User's question top_k: Number of chunks to retrieve Returns: List of relevant chunks with metadata """ if self.index is None or self.index.ntotal == 0: return [] # Limit top_k to available chunks top_k = min(top_k, self.index.ntotal) # Embed query query_embedding = self.embed_model.encode([query]).astype("float32") # Search FAISS (request more if scoping to a single doc might filter results) k_search = max(top_k, min(50, self.index.ntotal)) distances, indices = self.index.search(query_embedding, k=k_search) # If doc_id provided, determine filename to filter by filename_filter = None if doc_id and doc_id in self.documents: filename_filter = self.documents[doc_id].get('filename') # Gather results and apply optional filename filter results = [] for i, idx in enumerate(indices[0]): if idx < len(self.metadata): meta = self.metadata[idx] if filename_filter and meta.get('source') != filename_filter: continue results.append({ **meta, "distance": float(distances[0][i]), "relevance_rank": len(results) + 1 }) if len(results) >= top_k: break return results def generate_answer(self, query: str, context_chunks: List[Dict]) -> str: """ Generate answer using Gemini with retrieved context. Args: query: User's question context_chunks: Retrieved relevant chunks Returns: Generated answer string """ if not context_chunks: return "I don't have enough information to answer this question. Please upload relevant documents first." # Build context string context_parts = [] for chunk in context_chunks: context_parts.append( f"[Source: {chunk['source']}, Page {chunk['page']}]\n{chunk['text']}" ) context = "\n\n".join(context_parts) # Create prompt prompt = f"""You are a helpful assistant that answers questions based ONLY on the provided context. Do NOT make up information that is not in the context. If the context doesn't contain enough information to answer, say so clearly. You may summarize, combine, or rephrase information from the context to make your answer clear and helpful. CONTEXT: {context} QUESTION: {query} ANSWER:""" try: response = self.gemini_model.generate_content(prompt) return response.text except Exception as e: return f"Error generating answer: {str(e)}" def verify_sources(self, query: str, answer: str, context_chunks: List[Dict]) -> List[int]: """ Verify which chunks actually support the generated answer. Args: query: User's question answer: Generated answer context_chunks: All retrieved chunks Returns: List of indices of chunks that support the answer """ if not context_chunks: return [] # Build context with numbered chunks context_parts = [] for i, chunk in enumerate(context_chunks): context_parts.append( f"[{i}] Source: {chunk['source']}, Page {chunk['page']}\n{chunk['text']}" ) context = "\n\n".join(context_parts) # Create verification prompt prompt = f"""You are a citation verification assistant. Given a question, an answer, and numbered source chunks, identify which chunks were actually used to generate the answer. Return ONLY a comma-separated list of chunk numbers that directly support the answer (e.g., "0,2,3"). If no chunks support the answer, return "NONE". Do not include explanations or any other text. QUESTION: {query} ANSWER: {answer} NUMBERED CHUNKS: {context} CHUNK NUMBERS THAT SUPPORT THE ANSWER:""" try: response = self.gemini_model.generate_content(prompt) result = response.text.strip() # Parse the response if result.upper() == "NONE": return [] # Extract numbers used_indices = [] for part in result.split(","): try: idx = int(part.strip()) if 0 <= idx < len(context_chunks): used_indices.append(idx) except ValueError: continue return used_indices except Exception as e: print(f"Error verifying sources: {e}") # Fallback: return all chunks if verification fails return list(range(len(context_chunks))) def ask(self, query: str, top_k: int = DEFAULT_TOP_K, doc_id: str = None) -> Dict: """ Main query method: retrieve context, generate answer, and filter sources. Args: query: User's question top_k: Number of chunks to retrieve Returns: Dict with answer and verified sources """ # Retrieve relevant chunks (optionally scoped to a document) relevant_chunks = self.retrieve_relevant_chunks(query, top_k, doc_id=doc_id) # Generate answer answer = self.generate_answer(query, relevant_chunks) # Verify which chunks actually support the answer used_indices = self.verify_sources(query, answer, relevant_chunks) # Filter sources to only those that support the answer sources = [] seen = set() for idx in used_indices: if idx < len(relevant_chunks): chunk = relevant_chunks[idx] source_key = f"{chunk['source']}_{chunk['page']}" if source_key not in seen: sources.append({ "file": chunk["source"], "page": chunk["page"] }) seen.add(source_key) return { "answer": answer, "sources": sources, "num_chunks_used": len(sources), "num_chunks_retrieved": len(relevant_chunks) } # ============================================ # DOCUMENT MANAGEMENT METHODS # ============================================ def get_all_documents(self) -> List[Dict]: """ Get list of all uploaded documents. Returns: List of document info dicts """ return [ {"doc_id": doc_id, **info} for doc_id, info in self.documents.items() ] def delete_document(self, doc_id: str) -> Dict: """ Delete a document and its embeddings. Args: doc_id: Document ID to delete Returns: Result dict """ if doc_id not in self.documents: return { "status": "error", "message": f"Document {doc_id} not found" } filename = self.documents[doc_id]["filename"] # Remove from index self.remove_document_from_index(filename) # Remove from registry del self.documents[doc_id] # Persist changes self._save_persistent_data() return { "status": "success", "message": f"Document '{filename}' deleted successfully" } def get_stats(self) -> Dict: """ Get system statistics. Returns: Dict with stats """ return { "total_documents": len(self.documents), "total_chunks": len(self.metadata), "index_size": self.index.ntotal if self.index else 0, "embedding_model": EMBEDDING_MODEL_NAME, "embedding_dimension": EMBEDDING_DIMENSION }