rag_system.py

"""
RAG (Retrieval Augmented Generation) System for Math Tutor Bot
Allows the tutor to answer questions based on uploaded documents
"""

import os
import numpy as np
from sentence_transformers import SentenceTransformer
import faiss
from PyPDF2 import PdfReader
from typing import List, Dict
import pickle

class RAGSystem:
    def __init__(self, documents_dir='uploaded_docs', index_dir='rag_index'):
        """
        Initialize RAG system with document and index directories
        """
        self.documents_dir = documents_dir
        self.index_dir = index_dir

        # Create directories if they don't exist
        os.makedirs(documents_dir, exist_ok=True)
        os.makedirs(index_dir, exist_ok=True)

        # Load embedding model (lightweight, runs on CPU)
        print("Loading embedding model...")
        self.embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
        print("✓ Embedding model loaded!")

        # Initialize or load FAISS index
        self.index = None
        self.chunks = []
        self.chunk_metadata = []
        self.load_index()

    def chunk_text(self, text: str, chunk_size=500, overlap=50) -> List[str]:
        """
        Split text into overlapping chunks
        """
        words = text.split()
        chunks = []

        for i in range(0, len(words), chunk_size - overlap):
            chunk = ' '.join(words[i:i + chunk_size])
            chunks.append(chunk)

        return chunks

    def process_pdf(self, pdf_path: str) -> str:
        """
        Extract text from PDF file
        """
        try:
            reader = PdfReader(pdf_path)
            text = ""
            for page in reader.pages:
                text += page.extract_text() + "\n"
            return text
        except Exception as e:
            print(f"Error reading PDF {pdf_path}: {e}")
            return ""

    def process_txt(self, txt_path: str) -> str:
        """
        Read text from TXT file
        """
        try:
            with open(txt_path, 'r', encoding='utf-8') as f:
                return f.read()
        except Exception as e:
            print(f"Error reading TXT {txt_path}: {e}")
            return ""

    def add_document(self, file_path: str, filename: str):
        """
        Add a document to the RAG system
        """
        print(f"Processing document: {filename}")

        # Check if document already exists
        existing_filenames = [m['filename'] for m in self.chunk_metadata]
        if filename in existing_filenames:
            print(f"⚠️  Document '{filename}' already exists in the system!")
            print(f"Skipping upload to prevent duplicates.")
            return False

        # Extract text based on file type
        if filename.lower().endswith('.pdf'):
            text = self.process_pdf(file_path)
        elif filename.lower().endswith('.txt'):
            text = self.process_txt(file_path)
        else:
            print(f"Unsupported file type: {filename}")
            return False

        if not text.strip():
            print(f"No text extracted from {filename}")
            return False

        # Chunk the text
        new_chunks = self.chunk_text(text)
        print(f"Created {len(new_chunks)} chunks from {filename}")

        # Create embeddings
        print("Creating embeddings...")
        embeddings = self.embedding_model.encode(new_chunks, show_progress_bar=True)

        # Add to index
        if self.index is None:
            # Create new index
            dimension = embeddings.shape[1]
            self.index = faiss.IndexFlatL2(dimension)

        self.index.add(embeddings.astype('float32'))

        # Store chunks and metadata
        for i, chunk in enumerate(new_chunks):
            self.chunks.append(chunk)
            self.chunk_metadata.append({
                'filename': filename,
                'chunk_id': i,
                'total_chunks': len(new_chunks)
            })

        # Save index
        self.save_index()
        print(f"✓ Document {filename} added successfully!")
        return True

    def retrieve(self, query: str, top_k=3, relevance_threshold=1.5) -> List[Dict]:
        """
        Retrieve most relevant chunks for a query

        Args:
            query: The query string
            top_k: Number of chunks to retrieve
            relevance_threshold: Maximum L2 distance to consider relevant (default: 1.5)
                                Lower = more strict, higher = more lenient
        """
        if self.index is None or self.index.ntotal == 0:
            return []

        # Encode query
        query_embedding = self.embedding_model.encode([query])

        # Search in FAISS index
        distances, indices = self.index.search(query_embedding.astype('float32'), top_k)

        # Prepare results - only include if below relevance threshold
        results = []
        for i, idx in enumerate(indices[0]):
            distance = float(distances[0][i])
            if idx < len(self.chunks) and distance < relevance_threshold:
                results.append({
                    'text': self.chunks[idx],
                    'metadata': self.chunk_metadata[idx],
                    'score': distance
                })
                print(f"[RAG Retrieval] Chunk {idx} - distance: {distance:.3f} (relevant)")
            elif idx < len(self.chunks):
                print(f"[RAG Retrieval] Chunk {idx} - distance: {distance:.3f} (skipped - not relevant)")

        return results

    def get_context_for_query(self, query: str, top_k=3, relevance_threshold=1.5) -> str:
        """
        Get formatted context string for LLM from retrieved chunks

        Args:
            query: The query string
            top_k: Number of chunks to retrieve
            relevance_threshold: Maximum L2 distance (default: 1.5, use 999 to get all chunks)
        """
        results = self.retrieve(query, top_k, relevance_threshold)

        if not results:
            print("[RAG] No relevant context found - using base knowledge only")
            return ""

        print(f"[RAG] Found {len(results)} relevant chunk(s) - injecting context")
        context = "**Relevant information from uploaded documents:**\n\n"
        for i, result in enumerate(results, 1):
            filename = result['metadata']['filename']
            text = result['text']
            context += f"[Source: {filename}]\n{text}\n\n"

        return context

    def save_index(self):
        """
        Save FAISS index and metadata to disk
        """
        if self.index is not None:
            faiss.write_index(self.index, os.path.join(self.index_dir, 'faiss.index'))

        # Save chunks and metadata
        with open(os.path.join(self.index_dir, 'chunks.pkl'), 'wb') as f:
            pickle.dump({
                'chunks': self.chunks,
                'metadata': self.chunk_metadata
            }, f)

        print("✓ Index saved successfully!")

    def load_index(self):
        """
        Load existing FAISS index and metadata from disk
        """
        index_path = os.path.join(self.index_dir, 'faiss.index')
        chunks_path = os.path.join(self.index_dir, 'chunks.pkl')

        if os.path.exists(index_path) and os.path.exists(chunks_path):
            print("Loading existing RAG index...")
            self.index = faiss.read_index(index_path)

            with open(chunks_path, 'rb') as f:
                data = pickle.load(f)
                self.chunks = data['chunks']
                self.chunk_metadata = data['metadata']

            print(f"✓ Loaded {len(self.chunks)} chunks from {len(set(m['filename'] for m in self.chunk_metadata))} documents")
        else:
            print("No existing index found. Starting fresh.")

    def get_stats(self) -> Dict:
        """
        Get statistics about the RAG system
        """
        if not self.chunks:
            return {
                'total_chunks': 0,
                'total_documents': 0,
                'documents': []
            }

        documents = {}
        for metadata in self.chunk_metadata:
            filename = metadata['filename']
            if filename not in documents:
                documents[filename] = 0
            documents[filename] += 1

        return {
            'total_chunks': len(self.chunks),
            'total_documents': len(documents),
            'documents': [{'name': k, 'chunks': v} for k, v in documents.items()]
        }

    def clear_all(self):
        """
        Clear all documents and reset the index
        """
        self.index = None
        self.chunks = []
        self.chunk_metadata = []

        # Remove saved files
        index_path = os.path.join(self.index_dir, 'faiss.index')
        chunks_path = os.path.join(self.index_dir, 'chunks.pkl')

        if os.path.exists(index_path):
            os.remove(index_path)
        if os.path.exists(chunks_path):
            os.remove(chunks_path)

        print("✓ All documents cleared!")


if __name__ == "__main__":
    # Test the RAG system
    print("RAG System Test")
    print("=" * 60)

    rag = RAGSystem()
    print("\n" + "=" * 60)
    print("RAG System ready!")
    print("=" * 60)