Update src/rag_engine.py

src/rag_engine.py

@@ -9,10 +9,10 @@ from langchain_huggingface import HuggingFaceEmbeddings
 from langchain_openai import OpenAIEmbeddings
 from langchain_core.documents import Document
 
-# Internal Core Imports (The Good Stuff)
+# Internal Core Imports
 from core.PineconeManager import PineconeManager
 from core.AcronymManager import AcronymManager
-from core.ChunkingManager import ChunkingManager, ChunkingStrategy # NEW: The Traffic Cop
+from core.ChunkingManager import ChunkingManager
 from flashrank import Ranker, RerankRequest
 
 # CONFIGURATION
@@ -20,7 +20,7 @@ PINECONE_KEY = os.getenv("PINECONE_API_KEY")
 UPLOAD_DIR = "source_documents"
 logger = logging.getLogger(__name__)
 
-# Initialize Reranker (Small, fast CPU model)
+# Initialize Reranker
 try:
     reranker = Ranker(model_name="ms-marco-TinyBERT-L-2-v2", cache_dir="/tmp/flashrank_cache")
 except Exception as e:
@@ -65,36 +65,20 @@ def save_uploaded_file(uploaded_file, username: str) -> str:
         f.write(uploaded_file.getbuffer())
     return file_path
 
-# --- CORE LOGIC UPGRADE ---
-
 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.
-    """
+    """Delegates to ChunkingManager."""
     try:
         logger.info(f"Initializing ChunkingManager for {file_path} using {chunking_strategy}")
-        
-        # Initialize the Manager (it handles loading the right tokenizers/embedders)
         manager = ChunkingManager(embedding_model_name=embed_model_name)
+        chunks = manager.process_document(file_path=file_path, strategy=chunking_strategy, preprocess=True)
         
-        # 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 []
@@ -103,84 +87,83 @@ def ingest_file(file_path: str, username: str, index_name: str, embed_model_name
     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
+        # 1. Chunking
         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
+        # 2. METADATA SANITIZATION (The Fix for Pinecone IDs)
+        # We enforce that 'source' is just the filename, stripping the path.
+        clean_filename = os.path.basename(file_path)
+        for doc in docs:
+            doc.metadata["source"] = clean_filename
+            # Remove any absolute paths that might have leaked into metadata
+            if "file_path" in doc.metadata: del doc.metadata["file_path"]
+
+        # 3. Acronym Learning
         acronym_mgr = AcronymManager()
         for doc in docs:
             acronym_mgr.scan_text_for_acronyms(doc.page_content)
 
-        # 3. Pinecone Manager
+        # 4. Pinecone Manager
         pm = PineconeManager(PINECONE_KEY)
         
-        # 4. SAFETY CHECK
+        # 5. 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. PRE-EMPTIVE DELETE
+        pm.delete_file(index_name, clean_filename, namespace=username)
 
-        # 6. UPLOAD
+        # 7. 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)]
+        # Now IDs will be "filename.txt_0", "filename.txt_1" etc.
+        custom_ids = [f"{clean_filename}_{i}" for i, doc in enumerate(docs)]
         vstore.add_documents(docs, ids=custom_ids)
         
-        return True, f"Successfully updated {filename} ({len(docs)} chunks)."
+        return True, f"Successfully updated {clean_filename} ({len(docs)} chunks)."
 
     except Exception as e:
         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 (e.g. flattened summaries) using the ChunkingManager.
-    """
     if not PINECONE_KEY or not index_name: return False, "Pinecone Configuration Missing."
     
     try:
         pm = PineconeManager(PINECONE_KEY)
+        clean_source = os.path.basename(source_name)
         
-        # 1. PRE-EMPTIVE DELETE
-        pm.delete_file(index_name, source_name, namespace=username)
+        # 1. DELETE OLD
+        pm.delete_file(index_name, clean_source, namespace=username)
         
-        # 2. SAVE BACKUP (Critical: ChunkingManager works best with files)
+        # 2. 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)
+        backup_path = os.path.join(user_docs_dir, clean_source)
         
         with open(backup_path, "w", encoding='utf-8') as f:
             f.write(text)
 
-        # 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}")
+        # 3. CHUNK
         manager = ChunkingManager(embedding_model_name=embed_model_name)
-        
-        # 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
+        # 4. SANITIZE METADATA
         for doc in docs:
+            doc.metadata["source"] = clean_source
             doc.metadata["file_type"] = "generated"
             doc.metadata["strategy"] = "flattened"
 
-        # 4. UPLOAD
+        # 5. UPLOAD
         emb_fn = get_embedding_func(embed_model_name) 
         vstore = pm.get_vectorstore(index_name, emb_fn, namespace=username)
-        
-        custom_ids = [f"{source_name}_{i}" for i in enumerate(docs)]
+        custom_ids = [f"{clean_source}_{i}" for i, _ in enumerate(docs)]
         vstore.add_documents(docs, ids=custom_ids)
         
-        return True, f"Updated: {source_name} ({len(docs)} chunks)"
+        return True, f"Updated: {clean_source} ({len(docs)} chunks)"
         
     except Exception as e:
         logger.error(f"Error indexing text: {e}")
@@ -188,37 +171,27 @@ def process_and_add_text(text: str, source_name: str, username: str, index_name:
 
 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:
         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
         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)
-        ]
-        
+        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"Search failed: {e}")
         return []
@@ -227,7 +200,6 @@ 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)
-    
     if PINECONE_KEY and index_name:
         try:
             pm = PineconeManager(PINECONE_KEY)
@@ -247,8 +219,8 @@ 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)
+        # We wipe it clean first
         if os.path.exists(user_dir): shutil.rmtree(user_dir)
         os.makedirs(user_dir, exist_ok=True)
 
@@ -260,7 +232,10 @@ def rebuild_cache_from_pinecone(username: str, index_name: str) -> Tuple[bool, s
             vectors = response.vectors 
             for vec_id, vec_data in vectors.items():
                 meta = vec_data.metadata or {}
-                source = meta.get('source', 'unknown.txt')
+                # THE RESYNC FIX: Force basename to avoid "dir/dir/file" bugs
+                raw_source = meta.get('source', 'unknown.txt')
+                source = os.path.basename(raw_source)
+                
                 text = meta.get('text') or meta.get('page_content') or ''
                 try:
                     if "_" in vec_id: chunk_index = int(vec_id.rsplit('_', 1)[-1])
@@ -276,7 +251,7 @@ def rebuild_cache_from_pinecone(username: str, index_name: str) -> Tuple[bool, s
             file_path = os.path.join(user_dir, filename)
             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!"
+        return True, f"Restored {count} files from Pinecone!"
     except Exception as e:
         logger.error(f"Cache rebuild failed: {e}")
         return False, str(e)