Update src/rag_engine.py

src/rag_engine.py

@@ -3,15 +3,17 @@ import shutil
 import logging
 from typing import List, Tuple, Optional
 from huggingface_hub import snapshot_download
-from langchain_community.document_loaders import PyPDFLoader, TextLoader, UnstructuredWordDocumentLoader, UnstructuredPowerPointLoader
-from langchain_text_splitters import RecursiveCharacterTextSplitter
+
+# LangChain Imports
 from langchain_huggingface import HuggingFaceEmbeddings
 from langchain_openai import OpenAIEmbeddings
-from langchain_community.vectorstores import Pinecone as LangchainPinecone
 from langchain_core.documents import Document
+
+# Internal Core Imports (The Good Stuff)
 from core.PineconeManager import PineconeManager
 from core.AcronymManager import AcronymManager
-from flashrank import Ranker, RerankRequest # NEW IMPORT
+from core.ChunkingManager import ChunkingManager, ChunkingStrategy # NEW: The Traffic Cop
+from flashrank import Ranker, RerankRequest
 
 # CONFIGURATION
 PINECONE_KEY = os.getenv("PINECONE_API_KEY")
@@ -19,7 +21,6 @@ UPLOAD_DIR = "source_documents"
 logger = logging.getLogger(__name__)
 
 # 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:
@@ -33,29 +34,19 @@ def get_embedding_func(model_name: str = "sentence-transformers/all-MiniLM-L6-v2
             if not os.getenv("OPENAI_API_KEY"): raise ValueError("OpenAI API Key not found.")
             return OpenAIEmbeddings(model=model_name)
             
-        # CHECK 2: YOUR CUSTOM FINE-TUNE (Updated for Subfolders)
+        # CHECK 2: YOUR CUSTOM FINE-TUNE
         elif "navy-custom-models" in model_name:
              logger.info(f"Downloading custom model from: {model_name}")
-             
-             # 1. Parse the repo and folder from your string
-             # Input: "NavyDevilDoc/navy-custom-models/bge-finetuned"
              parts = model_name.split("/")
-             # Repo ID is the first two parts: "NavyDevilDoc/navy-custom-models"
              repo_id = f"{parts[0]}/{parts[1]}"
-             # Folder is the rest: "bge-finetuned"
              folder_name = parts[2]
              
-             # 2. Download ONLY that folder
              storage_path = snapshot_download(
                  repo_id=repo_id, 
                  repo_type="model", 
                  allow_patterns=f"{folder_name}/*"
              )
-             
-             # 3. Construct the local path to the inner folder
              local_model_path = os.path.join(storage_path, folder_name)
-             
-             # 4. Load from the local path
              return HuggingFaceEmbeddings(model_name=local_model_path)
              
         # CHECK 3: Standard Public Models
@@ -74,88 +65,84 @@ def save_uploaded_file(uploaded_file, username: str) -> str:
         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()]
+# --- CORE LOGIC UPGRADE ---
 
-def process_file(file_path: str, chunking_strategy: str = "paragraph") -> List[Document]:
-    ext = os.path.splitext(file_path)[1].lower()
+def process_file(file_path: str, chunking_strategy: str = "paragraph", embed_model_name: str = "all-mpnet-base-v2") -> List[Document]:
+    """
+    Delegates processing to the robust ChunkingManager in src/core.
+    """
     try:
-        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 []
+        logger.info(f"Initializing ChunkingManager for {file_path} using {chunking_strategy}")
         
-        raw_docs = loader.load()
-        text = "\n\n".join([d.page_content for d in raw_docs])
+        # Initialize the Manager (it handles loading the right tokenizers/embedders)
+        manager = ChunkingManager(embedding_model_name=embed_model_name)
         
-        if chunking_strategy == "token":
-            splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
-            chunks = splitter.create_documents([text])
-        else:
-            chunker = ParagraphChunker()
-            texts = chunker.split_text(text)
-            chunks = [Document(page_content=t) for t in texts]
-            
-        # Add metadata
-        filename = os.path.basename(file_path)
-        for doc in chunks:
-            doc.metadata["source"] = filename
-            doc.metadata["strategy"] = chunking_strategy
+        # Execute the robust processing pipeline
+        # This now uses your ParagraphChunker.py or TokenChunker.py logic!
+        chunks = manager.process_document(
+            file_path=file_path,
+            strategy=chunking_strategy,
+            preprocess=True
+        )
+        
+        # Handle case where process_document returns a Dict (e.g. hierarchical) vs List
+        if isinstance(chunks, dict):
+            # For now, flatten dictionary returns if any (future proofing)
+            flat_chunks = []
+            for key, val in chunks.items():
+                if isinstance(val, list): flat_chunks.extend(val)
+            return flat_chunks
             
         return chunks
+
     except Exception as e:
         logger.error(f"Error processing {file_path}: {e}")
         return []
 
-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 []
+def ingest_file(file_path: str, username: str, index_name: str, embed_model_name: str = "sentence-transformers/all-MiniLM-L6-v2", strategy: str = "paragraph") -> Tuple[bool, str]:
+    if not PINECONE_KEY or not index_name: return False, "Pinecone Configuration Missing."
     
     try:
+        # 1. Chunking (Delegated to Core)
+        # Note: We pass the embedding model name so the chunker can initialize correctly
+        docs = process_file(file_path, chunking_strategy=strategy, embed_model_name=embed_model_name)
+        
+        if not docs: return False, "No valid chunks generated."
+
+        # 2. Acronym Learning
+        acronym_mgr = AcronymManager()
+        for doc in docs:
+            acronym_mgr.scan_text_for_acronyms(doc.page_content)
+
+        # 3. Pinecone Manager
         pm = PineconeManager(PINECONE_KEY)
+        
+        # 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
+        filename = os.path.basename(file_path)
+        pm.delete_file(index_name, filename, namespace=username)
+
+        # 6. UPLOAD
         vstore = pm.get_vectorstore(index_name, emb_fn, namespace=username)
+        # Generate stable IDs using the metadata source or filename
+        custom_ids = [f"{doc.metadata.get('source', filename)}_{i}" for i, doc in enumerate(docs)]
+        vstore.add_documents(docs, ids=custom_ids)
         
-        # 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)
-        
-        # 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
+        return True, f"Successfully updated {filename} ({len(docs)} chunks)."
 
     except Exception as e:
-        logger.error(f"Search failed: {e}")
-        return []
+        logger.error(f"Ingestion failed: {e}")
+        return False, str(e)
 
 def process_and_add_text(text: str, source_name: str, username: str, index_name: str, embed_model_name: str = None) -> Tuple[bool, str]:
     """
-    Ingests raw text with a specific embedding model.
+    Ingests raw text (e.g. flattened summaries) using the ChunkingManager.
     """
     if not PINECONE_KEY or not index_name: return False, "Pinecone Configuration Missing."
     
@@ -165,7 +152,7 @@ def process_and_add_text(text: str, source_name: str, username: str, index_name:
         # 1. PRE-EMPTIVE DELETE
         pm.delete_file(index_name, source_name, namespace=username)
         
-        # 2. SAVE BACKUP
+        # 2. SAVE BACKUP (Critical: ChunkingManager works best with files)
         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)
@@ -173,58 +160,68 @@ def process_and_add_text(text: str, source_name: str, username: str, index_name:
         with open(backup_path, "w", encoding='utf-8') as f:
             f.write(text)
 
-        # 3. UPLOAD (The Fix: Use the passed model name)
-        emb_fn = get_embedding_func(embed_model_name) 
+        # 3. CHUNK & PROCESS (Using the Manager on the backup file)
+        # This ensures flattened text gets the same robust metadata/cleaning as regular files
+        logger.info(f"Processing flattened text via ChunkingManager: {source_name}")
+        manager = ChunkingManager(embedding_model_name=embed_model_name)
         
-        doc = Document(
-            page_content=text,
-            metadata={"source": source_name, "strategy": "flattened", "file_type": "generated"}
-        )
+        # We use 'token' strategy for summaries usually, or 'paragraph' if preferred
+        docs = manager.process_document(backup_path, strategy="paragraph", preprocess=True)
         
+        # Override metadata to ensure it's marked as generated
+        for doc in docs:
+            doc.metadata["file_type"] = "generated"
+            doc.metadata["strategy"] = "flattened"
+
+        # 4. UPLOAD
+        emb_fn = get_embedding_func(embed_model_name) 
         vstore = pm.get_vectorstore(index_name, emb_fn, namespace=username)
-        vstore.add_documents([doc], ids=[f"{source_name}_0"])
         
-        return True, f"Updated: {source_name}"
+        custom_ids = [f"{source_name}_{i}" for i in enumerate(docs)]
+        vstore.add_documents(docs, ids=custom_ids)
+        
+        return True, f"Updated: {source_name} ({len(docs)} chunks)"
+        
     except Exception as e:
         logger.error(f"Error indexing text: {e}")
         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]:
-    if not PINECONE_KEY or not index_name: return False, "Pinecone Configuration Missing."
+def search_knowledge_base(query: str, username: str, index_name: str, embed_model_name: str, k: int = 5, final_k: int = 5):
+    if not PINECONE_KEY or not index_name: return []
+    
     try:
-        # 1. Chunking
-        docs = process_file(file_path, chunking_strategy=strategy)
-        if not docs: return False, "No valid chunks generated."
-
-        # 2. Acronym Learning
-        acronym_mgr = AcronymManager()
-        for doc in docs:
-            acronym_mgr.scan_text_for_acronyms(doc.page_content)
-
-        # 3. Pinecone Manager
         pm = PineconeManager(PINECONE_KEY)
-        
-        # 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
-        filename = os.path.basename(file_path)
-        pm.delete_file(index_name, filename, namespace=username)
-
-        # 6. UPLOAD
         vstore = pm.get_vectorstore(index_name, emb_fn, namespace=username)
-        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)."
+        # 1. RETRIEVE BROAD
+        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
+        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)
+        
+        # 3. SELECT TOP K
+        final_docs = []
+        for res in ranked_results[:final_k]:
+            meta = res.get("meta", {})
+            meta["rerank_score"] = res.get("score") 
+            final_docs.append(Document(page_content=res["text"], metadata=meta))
+            
+        return final_docs
 
     except Exception as e:
-        logger.error(f"Ingestion failed: {e}")
-        return False, str(e)
+        logger.error(f"Search failed: {e}")
+        return []
 
 def delete_document(username: str, filename: str, index_name: str):
     user_dir = os.path.join(UPLOAD_DIR, username)
@@ -250,6 +247,11 @@ def rebuild_cache_from_pinecone(username: str, index_name: str) -> Tuple[bool, s
         ids = pm.get_all_ids(index_name, username)
         if not ids: return False, "No data found in Pinecone."
             
+        # Nuke local folder first to handle deletions
+        user_dir = os.path.join(UPLOAD_DIR, username)
+        if os.path.exists(user_dir): shutil.rmtree(user_dir)
+        os.makedirs(user_dir, exist_ok=True)
+
         batch_size = 100
         reconstructed_files = {} 
         for i in range(0, len(ids), batch_size):
@@ -267,11 +269,9 @@ def rebuild_cache_from_pinecone(username: str, index_name: str) -> Tuple[bool, s
                 if source not in reconstructed_files: reconstructed_files[source] = []
                 reconstructed_files[source].append((chunk_index, text))
         
-        user_dir = os.path.join(UPLOAD_DIR, username)
-        os.makedirs(user_dir, exist_ok=True)
         count = 0
         for filename, chunks in reconstructed_files.items():
-            chunks.sort(key=lambda x: x[0]) # SORTING FIX
+            chunks.sort(key=lambda x: x[0]) 
             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)