Update src/rag_engine.py

src/rag_engine.py

@@ -1,204 +1,147 @@
 import os
 import shutil
 import logging
-from typing import List, Literal, Tuple
-
-# --- LANGCHAIN & DB IMPORTS ---
+from typing import List, Tuple, Optional
+from langchain_community.document_loaders import PyPDFLoader, TextLoader, UnstructuredWordDocumentLoader, UnstructuredPowerPointLoader
+from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_huggingface import HuggingFaceEmbeddings
 from langchain_openai import OpenAIEmbeddings
+from langchain_community.vectorstores import Pinecone as LangchainPinecone
 from langchain_core.documents import Document
-from langchain_text_splitters import MarkdownHeaderTextSplitter, RecursiveCharacterTextSplitter
-from sentence_transformers import CrossEncoder
-
-
-# --- CUSTOM CORE IMPORTS ---
 from core.PineconeManager import PineconeManager
-from core.ParagraphChunker import ParagraphChunker
-from core.TokenChunker import TokenChunker
 from core.AcronymManager import AcronymManager
+from flashrank import Ranker, RerankRequest # NEW IMPORT
 
-# --- CONFIGURATION ---
-UPLOAD_DIR = "source_documents"
-EMBED_MODEL_NAME = "sentence-transformers/all-mpnet-base-v2"
-RERANK_MODEL_NAME = "cross-encoder/ms-marco-MiniLM-L-6-v2"
+# CONFIGURATION
 PINECONE_KEY = os.getenv("PINECONE_API_KEY")
-
-# Configure Logging
-logging.basicConfig(level=logging.INFO)
+UPLOAD_DIR = "source_documents"
 logger = logging.getLogger(__name__)
 
-# --- LAZY LOADING GLOBALS ---
-_embedding_func = None
-_rerank_model = None
+# Initialize Reranker (Small, fast CPU model)
+# Only initializes once when the app starts
+try:
+    reranker = Ranker(model_name="ms-marco-TinyBERT-L-2-v2", cache_dir="/tmp/flashrank_cache")
+except Exception as e:
+    logger.warning(f"Reranker failed to load: {e}")
+    reranker = None
 
 def get_embedding_func(model_name: str = "sentence-transformers/all-MiniLM-L6-v2"):
-    """
-    Dynamically loads the correct embedding model based on the selection.
-    """
     try:
-        # 1. OpenAI Models
         if "openai" in model_name.lower():
-            if not os.getenv("OPENAI_API_KEY"):
-                raise ValueError("OpenAI API Key not found.")
-            
-            # Map friendly names to actual API model names if needed
-            # But usually we just pass the exact string like "text-embedding-3-small"
+            if not os.getenv("OPENAI_API_KEY"): raise ValueError("OpenAI API Key not found.")
             return OpenAIEmbeddings(model=model_name)
-            
-        # 2. Hugging Face Models (Local / CPU-friendly)
         else:
-            # Default to all-MiniLM if something weird is passed, or use the specific HF model
             return HuggingFaceEmbeddings(model_name=model_name)
-            
     except Exception as e:
         logger.error(f"Failed to load embedding model '{model_name}': {e}")
-        # Fallback to the safe default if everything explodes
         return HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
 
-def get_rerank_model():
-    """Lazy loads the Cross-Encoder model."""
-    global _rerank_model
-    if _rerank_model is None:
-        logger.info(f"⏳ Loading Reranker: {RERANK_MODEL_NAME}...")
-        _rerank_model = CrossEncoder(RERANK_MODEL_NAME)
-        logger.info("✅ Reranker Loaded.")
-    return _rerank_model
-
-# --- PART 1: CHUNKING LOGIC (The New System) ---
-
-def _process_markdown(file_path: str, chunk_size: int = 1000, chunk_overlap: int = 100) -> List[Document]:
-    """Internal helper to process Markdown files using Header Semantic Splitting."""
+def save_uploaded_file(uploaded_file, username: str) -> str:
+    user_dir = os.path.join(UPLOAD_DIR, username)
+    os.makedirs(user_dir, exist_ok=True)
+    file_path = os.path.join(user_dir, uploaded_file.name)
+    with open(file_path, "wb") as f:
+        f.write(uploaded_file.getbuffer())
+    return file_path
+
+class ParagraphChunker:
+    def split_text(self, text):
+        return [p.strip() for p in text.split('\n\n') if p.strip()]
+
+def process_file(file_path: str, chunking_strategy: str = "paragraph") -> List[Document]:
+    ext = os.path.splitext(file_path)[1].lower()
     try:
-        with open(file_path, 'r', encoding='utf-8') as f:
-            markdown_text = f.read()
-
-        headers_to_split_on = [
-            ("#", "Header 1"),
-            ("##", "Header 2"),
-            ("###", "Header 3"),
-        ]
-
-        markdown_splitter = MarkdownHeaderTextSplitter(headers_to_split_on=headers_to_split_on)
-        md_header_splits = markdown_splitter.split_text(markdown_text)
-
-        text_splitter = RecursiveCharacterTextSplitter(
-            chunk_size=chunk_size, 
-            chunk_overlap=chunk_overlap
-        )
-        final_docs = text_splitter.split_documents(md_header_splits)
+        if ext == ".pdf": loader = PyPDFLoader(file_path)
+        elif ext == ".txt": loader = TextLoader(file_path, encoding='utf-8')
+        elif ext == ".docx": loader = UnstructuredWordDocumentLoader(file_path)
+        elif ext == ".pptx": loader = UnstructuredPowerPointLoader(file_path)
+        elif ext == ".md": loader = TextLoader(file_path, encoding='utf-8')
+        else: return []
+        
+        raw_docs = loader.load()
+        text = "\n\n".join([d.page_content for d in raw_docs])
         
-        for doc in final_docs:
-            doc.metadata['source'] = os.path.basename(file_path)
-            doc.metadata['file_type'] = 'md'
-            doc.metadata['strategy'] = 'markdown_header'
-
-        return final_docs
-    except Exception as e:
-        logger.error(f"Error processing Markdown file {file_path}: {e}")
-        return []
-
-def process_file(
-    file_path: str, 
-    chunking_strategy: Literal["paragraph", "token"] = "paragraph",
-    chunk_size: int = 512,
-    chunk_overlap: int = 100,
-    model_name: str = "gpt-4o"
-) -> List[Document]:
-    """
-    Main chunking engine. Routes file to specific chunkers based on type/strategy.
-    """
-    if not os.path.exists(file_path):
-        logger.error(f"File not found: {file_path}")
-        return []
-
-    file_extension = os.path.splitext(file_path)[1].lower()
-    file_name = os.path.basename(file_path)
-    logger.info(f"Processing {file_name} using strategy: {chunking_strategy}")
-
-    # 1. Handle Markdown
-    if file_extension == ".md":
-        return _process_markdown(file_path, chunk_size, chunk_overlap)
-
-    # 2. Handle PDF and TXT
-    elif file_extension in [".pdf", ".txt"]:
         if chunking_strategy == "token":
-            chunker = TokenChunker(
-                model_name=model_name,
-                chunk_size=chunk_size,
-                chunk_overlap=chunk_overlap
-            )
+            splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
+            chunks = splitter.create_documents([text])
         else:
-            chunker = ParagraphChunker(model_name=model_name)
-
-        try:
-            if file_extension == ".pdf":
-                docs = chunker.process_document(file_path)
-            elif file_extension == ".txt":
-                docs = chunker.process_text_file(file_path)
+            chunker = ParagraphChunker()
+            texts = chunker.split_text(text)
+            chunks = [Document(page_content=t) for t in texts]
             
-            # Ensure metadata consistency
-            for doc in docs:
-                doc.metadata["source"] = file_name
-                doc.metadata["strategy"] = chunking_strategy
-            
-            return docs
+        # Add metadata
+        filename = os.path.basename(file_path)
+        for doc in chunks:
+            doc.metadata["source"] = filename
+            doc.metadata["strategy"] = chunking_strategy
             
-        except Exception as e:
-            logger.error(f"Error using {chunking_strategy} chunker on {file_name}: {e}")
-            return []
-    else:
-        logger.warning(f"Unsupported file extension: {file_extension}")
+        return chunks
+    except Exception as e:
+        logger.error(f"Error processing {file_path}: {e}")
         return []
 
-# --- PART 2: DATABASE & FILE MANAGEMENT (Pinecone Version) ---
-
-def save_uploaded_file(uploaded_file, username: str = "default") -> str:
-    """Saves a StreamlitUploadedFile to disk so the loaders can read it."""
+def search_knowledge_base(query: str, username: str, index_name: str, embed_model_name: str, k: int = 5, final_k: int = 5):
+    """
+    Searches Pinecone with Reranking.
+    1. Fetches 3x candidates (Top 15).
+    2. Reranks using TinyBERT.
+    3. Returns Top 5.
+    """
+    if not PINECONE_KEY or not index_name: return []
+    
     try:
-        user_dir = os.path.join(UPLOAD_DIR, username)
-        os.makedirs(user_dir, exist_ok=True)
-        file_path = os.path.join(user_dir, uploaded_file.name)
+        pm = PineconeManager(PINECONE_KEY)
+        emb_fn = get_embedding_func(embed_model_name)
+        vstore = pm.get_vectorstore(index_name, emb_fn, namespace=username)
+        
+        # 1. RETRIEVE BROAD (Fetch 3x what we need)
+        broad_k = final_k * 3
+        initial_docs = vstore.similarity_search(query, k=broad_k)
+        
+        if not initial_docs or not reranker: 
+            return initial_docs[:final_k]
+        
+        # 2. RERANK (The Brain Upgrade)
+        passages = [
+            {"id": str(i), "text": doc.page_content, "meta": doc.metadata} 
+            for i, doc in enumerate(initial_docs)
+        ]
+        
+        rerank_request = RerankRequest(query=query, passages=passages)
+        ranked_results = reranker.rerank(rerank_request)
         
-        with open(file_path, "wb") as f:
-            f.write(uploaded_file.getbuffer())
-        return file_path
+        # 3. SELECT TOP K
+        final_docs = []
+        for res in ranked_results[:final_k]:
+            meta = res.get("meta", {})
+            meta["rerank_score"] = res.get("score") # Useful for debugging
+            final_docs.append(Document(page_content=res["text"], metadata=meta))
+            
+        return final_docs
+
     except Exception as e:
-        logger.error(f"Error saving file: {e}")
-        return None
+        logger.error(f"Search failed: {e}")
+        return []
 
 def process_and_add_text(text: str, source_name: str, username: str, index_name: str) -> Tuple[bool, str]:
-    """
-    Ingests raw text.
-    UPGRADE: Performs 'Clean Replace' - deletes old version of this source before adding new.
-    """
     if not PINECONE_KEY or not index_name: return False, "Pinecone Configuration Missing."
-    
     try:
         pm = PineconeManager(PINECONE_KEY)
         
-        # 1. PRE-EMPTIVE DELETE (The Fix)
-        # We wipe any existing vectors with this source name to prevent duplicates.
-        # This effectively makes this an "Update/Replace" operation.
+        # 1. PRE-EMPTIVE DELETE
         pm.delete_file(index_name, source_name, namespace=username)
         
-        # 2. SAVE PHYSICAL BACKUP (For Quiz Engine)
+        # 2. SAVE BACKUP
         user_docs_dir = os.path.join(UPLOAD_DIR, username)
         os.makedirs(user_docs_dir, exist_ok=True)
         backup_path = os.path.join(user_docs_dir, source_name)
-        
         with open(backup_path, "w", encoding='utf-8') as f:
             f.write(text)
 
-        # 3. UPLOAD TO PINECONE
-        emb_fn = get_embedding_func() # Uses default or last active model logic internally
-        
-        doc = Document(
-            page_content=text,
-            metadata={"source": source_name, "strategy": "flattened", "file_type": "generated"}
-        )
-        
+        # 3. UPLOAD
+        emb_fn = get_embedding_func()
+        doc = Document(page_content=text, metadata={"source": source_name, "strategy": "flattened", "file_type": "generated"})
         vstore = pm.get_vectorstore(index_name, emb_fn, namespace=username)
-        # Custom ID isn't strictly necessary for single-doc flattened text, but good for consistency
         vstore.add_documents([doc], ids=[f"{source_name}_0"])
         
         return True, f"Updated: {source_name}"
@@ -207,12 +150,7 @@ def process_and_add_text(text: str, source_name: str, username: str, index_name:
         return False, str(e)
 
 def ingest_file(file_path: str, username: str, index_name: str, embed_model_name: str = None, strategy: str = "paragraph") -> Tuple[bool, str]:
-    """
-    Chunks and uploads file.
-    UPGRADE: Performs 'Clean Replace' - deletes old chunks before uploading new ones.
-    """
     if not PINECONE_KEY or not index_name: return False, "Pinecone Configuration Missing."
-
     try:
         # 1. Chunking
         docs = process_file(file_path, chunking_strategy=strategy)
@@ -226,26 +164,20 @@ def ingest_file(file_path: str, username: str, index_name: str, embed_model_name
         # 3. Pinecone Manager
         pm = PineconeManager(PINECONE_KEY)
         
-        # 4. SAFETY CHECK (Dimensions)
+        # 4. SAFETY CHECK
         emb_fn = get_embedding_func(embed_model_name)
         test_vec = emb_fn.embed_query("test")
         model_dim = len(test_vec)
-        
         if not pm.check_dimension_compatibility(index_name, model_dim):
             return False, f"Dimension Mismatch! Index '{index_name}' expects {model_dim}d vectors."
 
-        # 5. PRE-EMPTIVE DELETE (The Fix)
-        # Wipe the slate clean for this specific filename
+        # 5. PRE-EMPTIVE DELETE
         filename = os.path.basename(file_path)
         pm.delete_file(index_name, filename, namespace=username)
 
-        # 6. UPLOAD NEW CHUNKS
+        # 6. UPLOAD
         vstore = pm.get_vectorstore(index_name, emb_fn, namespace=username)
-        
-        # Generate readable IDs: "filename_0", "filename_1"
-        # This helps with the 'Frankenstein' sorting fix we added earlier
         custom_ids = [f"{doc.metadata.get('source', filename)}_{i}" for i, doc in enumerate(docs)]
-        
         vstore.add_documents(docs, ids=custom_ids)
         
         return True, f"Successfully updated {filename} ({len(docs)} chunks)."
@@ -254,147 +186,57 @@ def ingest_file(file_path: str, username: str, index_name: str, embed_model_name
         logger.error(f"Ingestion failed: {e}")
         return False, str(e)
 
-def search_knowledge_base(query: str, username: str, index_name: str, embed_model_name: str, k: int = 10, final_k: int = 4) -> List[Document]:
-    """Retrieves from Pinecone -> Reranks."""
-    if not PINECONE_KEY or not index_name: return []
+def delete_document(username: str, filename: str, index_name: str):
+    user_dir = os.path.join(UPLOAD_DIR, username)
+    file_path = os.path.join(user_dir, filename)
+    if os.path.exists(file_path): os.remove(file_path)
     
-    try:
-        # 1. Expand Query (Acronyms)
-        acronym_mgr = AcronymManager()
-        expanded_query = acronym_mgr.expand_query(query)
-
-        # 2. Vector Search
-        pm = PineconeManager(PINECONE_KEY)
-        emb_fn = get_embedding_func(embed_model_name)
-        vstore = pm.get_vectorstore(index_name, emb_fn, namespace=username)
-        
-        results = vstore.similarity_search(expanded_query, k=k)
-        if not results: return []
-
-        # 3. Reranking
-        candidate_docs = results
-        candidate_texts = [doc.page_content for doc in candidate_docs]
-        pairs = [[expanded_query, text] for text in candidate_texts]
-        
-        reranker = get_rerank_model()
-        scores = reranker.predict(pairs)
-        
-        # Sort
-        scored_docs = list(zip(candidate_docs, scores))
-        scored_docs.sort(key=lambda x: x[1], reverse=True)
-        
-        return [doc for doc, score in scored_docs[:final_k]]
-
-    except Exception as e:
-        logger.error(f"Search Error: {e}")
-        return []
+    if PINECONE_KEY and index_name:
+        try:
+            pm = PineconeManager(PINECONE_KEY)
+            pm.delete_file(index_name, filename, namespace=username)
+        except Exception as e:
+            logger.error(f"Pinecone delete failed: {e}")
 
 def list_documents(username: str) -> List[dict]:
-    """
-    NOTE: Pinecone does not support easy listing of all unique files.
-    We return the Local Cache (source_documents) as a proxy for what is 
-    available for the Quiz Engine.
-    """
     user_dir = os.path.join(UPLOAD_DIR, username)
     if not os.path.exists(user_dir): return []
-    
-    files = []
-    for f in os.listdir(user_dir):
-        if f.lower().endswith(('.pdf', '.txt', '.md')):
-            files.append({"filename": f, "source": f, "strategy": "local_cache"})
-    return files
-
-def delete_document(username: str, filename: str, index_name: str) -> Tuple[bool, str]:
-    """Deletes from Pinecone AND Local Disk."""
-    if not PINECONE_KEY or not index_name: return False, "Config Missing."
-    
-    try:
-        # 1. Delete from Pinecone
-        pm = PineconeManager(PINECONE_KEY)
-        pm.delete_file(index_name, filename, namespace=username)
-        
-        # 2. Delete from Disk (Clean up Quiz Cache)
-        local_path = os.path.join(UPLOAD_DIR, username, filename)
-        if os.path.exists(local_path):
-            os.remove(local_path)
-            
-        return True, f"Deleted {filename} from Index and Disk."
-    except Exception as e:
-        return False, str(e)
-
-def reset_knowledge_base(username: str) -> Tuple[bool, str]:
-    """
-    WARNING: This deletes the USER NAMESPACE in Pinecone, not the whole Index.
-    """
-    # Pinecone delete_all is index-wide usually. 
-    # For safety in namespace-based multi-tenancy, we usually skip this 
-    # or implement a delete_all(delete_all=True, namespace=username)
-    return False, "Resetting entire DB via API is disabled for safety. Use Delete."
+    return [{"filename": f, "source": f} for f in os.listdir(user_dir) if f.lower().endswith(('.txt', '.md', '.pdf', '.docx'))]
 
 def rebuild_cache_from_pinecone(username: str, index_name: str) -> Tuple[bool, str]:
-    """
-    Downloads text from Pinecone and reconstructs local source files.
-    FIX: Sorts chunks numerically (_0, _1, _2) to prevent 'Frankenstein' files.
-    """
-    if not PINECONE_KEY or not index_name:
-        return False, "Pinecone config missing."
-    
+    if not PINECONE_KEY or not index_name: return False, "Pinecone config missing."
     try:
         pm = PineconeManager(PINECONE_KEY)
-        
-        # 1. Get all Vector IDs
         ids = pm.get_all_ids(index_name, username)
         if not ids: return False, "No data found in Pinecone."
             
-        # 2. Fetch content
         batch_size = 100
-        reconstructed_files = {} # { "filename.txt": [ (index, text), (index, text) ] }
-        
+        reconstructed_files = {} 
         for i in range(0, len(ids), batch_size):
             batch_ids = ids[i : i + batch_size]
             response = pm.fetch_vectors(index_name, batch_ids, username)
             vectors = response.vectors 
-            
             for vec_id, vec_data in vectors.items():
                 meta = vec_data.metadata or {}
                 source = meta.get('source', 'unknown.txt')
-                # Try to get text from 'text' (langchain default) or 'page_content' (our backup)
                 text = meta.get('text') or meta.get('page_content') or ''
-                
-                # EXTRACT CHUNK INDEX FROM ID (e.g., "doc.txt_12" -> 12)
                 try:
-                    # Assumes ID format "filename_index" from our new ingestion logic
-                    if "_" in vec_id:
-                        chunk_index = int(vec_id.rsplit('_', 1)[-1])
-                    else:
-                        chunk_index = 0
-                except ValueError:
-                    chunk_index = 0 # Fallback
-                
-                if source not in reconstructed_files:
-                    reconstructed_files[source] = []
+                    if "_" in vec_id: chunk_index = int(vec_id.rsplit('_', 1)[-1])
+                    else: chunk_index = 0
+                except ValueError: chunk_index = 0
+                if source not in reconstructed_files: reconstructed_files[source] = []
                 reconstructed_files[source].append((chunk_index, text))
         
-        # 3. Write to Disk (Sorted)
         user_dir = os.path.join(UPLOAD_DIR, username)
         os.makedirs(user_dir, exist_ok=True)
-        
         count = 0
         for filename, chunks in reconstructed_files.items():
-            # SORT BY INDEX (The Fix)
-            # This ensures Paragraph 1 comes before Paragraph 2
-            chunks.sort(key=lambda x: x[0])
-            
-            # Join text only
+            chunks.sort(key=lambda x: x[0]) # SORTING FIX
             full_text = "\n\n".join([c[1] for c in chunks])
-            
             file_path = os.path.join(user_dir, filename)
-            with open(file_path, "w", encoding="utf-8") as f:
-                f.write(full_text)
+            with open(file_path, "w", encoding="utf-8") as f: f.write(full_text)
             count += 1
-            
         return True, f"Restored {count} files (Sorted) from Pinecone!"
-        
     except Exception as e:
         logger.error(f"Cache rebuild failed: {e}")
         return False, str(e)