Update src/rag_engine.py

refactored to make use of the document loading program

src/rag_engine.py

@@ -1,210 +1,207 @@
 import os
+import shutil
+import time
+from langchain_text_splitters import RecursiveCharacterTextSplitter, TokenTextSplitter
 from langchain_chroma import Chroma
 from langchain_huggingface import HuggingFaceEmbeddings
-from sentence_transformers import CrossEncoder
-from core.ChunkingManager import ChunkingManager, ChunkingStrategy
-import tracker 
+from langchain_community.docstore.document import Document
+from sentence_transformers import CrossEncoder # Re-added for Reranking
+import doc_loader 
 
 # --- CONFIGURATION ---
-UPLOAD_DIR = "/tmp/rag_uploads"
-DB_ROOT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "chroma_db")
-EMBEDDING_MODEL_NAME = "all-MiniLM-L12-v2"
-RERANKER_MODEL_NAME = "BAAI/bge-reranker-base" 
-
-# --- LAZY LOADING SINGLETONS ---
-_embedding_fn = None
-_reranker = None
-_chunk_manager = None
-
-def get_embedding_function():
-    global _embedding_fn
-    if _embedding_fn is None:
-        print("⚙️ Loading Embedding Model...")
-        _embedding_fn = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL_NAME)
-    return _embedding_fn
-
-def get_reranker_model():
-    global _reranker
-    if _reranker is None:
-        print("⚙️ Loading Reranker Model...")
-        _reranker = CrossEncoder(RERANKER_MODEL_NAME)
-    return _reranker
-
-def get_chunk_manager():
-    global _chunk_manager
-    if _chunk_manager is None:
-        print("⚙️ Loading Chunk Manager...")
-        _chunk_manager = ChunkingManager(embedding_model_name=EMBEDDING_MODEL_NAME)
-    return _chunk_manager
-
-# --- DATABASE OPERATIONS ---
-def get_vectorstore(username):
-    safe_username = os.path.basename(username)
-    user_db_path = os.path.join(DB_ROOT, safe_username)
-    
-    if not os.path.exists(user_db_path):
-        os.makedirs(user_db_path, exist_ok=True)
-        
-    return Chroma(
-        persist_directory=user_db_path,
-        embedding_function=get_embedding_function(),
-        collection_name=f"docs_{safe_username}" 
-    )
-
+CHROMA_PATH = "chroma_db"
+UPLOAD_DIR = "temp_ingest" # Re-added directory constant
+EMBED_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
+RERANK_MODEL_NAME = "cross-encoder/ms-marco-MiniLM-L-6-v2" # Re-added model name
+
+# --- LAZY LOADING GLOBALS ---
+# We use a global variable pattern to avoid loading heavy models 
+# until the moment they are actually needed (saves startup RAM).
+_embedding_func = None
+_rerank_model = None
+
+def get_embedding_func():
+    """Lazy loads the embedding model."""
+    global _embedding_func
+    if _embedding_func is None:
+        print(f"⏳ Loading Embedding Model: {EMBED_MODEL_NAME}...")
+        _embedding_func = HuggingFaceEmbeddings(model_name=EMBED_MODEL_NAME)
+        print("✅ Embedding Model Loaded.")
+    return _embedding_func
+
+def get_rerank_model():
+    """Lazy loads the Cross-Encoder model."""
+    global _rerank_model
+    if _rerank_model is None:
+        print(f"⏳ Loading Reranker: {RERANK_MODEL_NAME}...")
+        _rerank_model = CrossEncoder(RERANK_MODEL_NAME)
+        print("✅ Reranker Loaded.")
+    return _rerank_model
+
+# --- FILE OPERATIONS ---
 def save_uploaded_file(uploaded_file):
-    if not os.path.exists(UPLOAD_DIR):
-        os.makedirs(UPLOAD_DIR)
-    safe_filename = os.path.basename(uploaded_file.name)
-    file_path = os.path.join(UPLOAD_DIR, safe_filename)
+    """Saves uploaded file to the temp directory."""
+    os.makedirs(UPLOAD_DIR, exist_ok=True)
+    file_path = os.path.join(UPLOAD_DIR, uploaded_file.name)
+    
     with open(file_path, "wb") as f:
         f.write(uploaded_file.getbuffer())
+        
     return file_path
 
-def process_and_add_document(file_path, username, strategy="paragraph"):
+# --- INGESTION PIPELINE ---
+def process_and_add_document(file_path, username, strategy, use_vision=False, api_key=None):
+    """
+    Ingests a document using the Universal Loader and adds it to the user's vector DB.
+    """
+    user_db_path = os.path.join(CHROMA_PATH, username)
+    
     try:
-        print(f"🧠 Chunking {os.path.basename(file_path)}...")
-        
-        strat_map = {
-            "token": ChunkingStrategy.TOKEN,
-            "paragraph": ChunkingStrategy.PARAGRAPH,
-            "page": ChunkingStrategy.PAGE
-        }
-        selected_strategy = strat_map.get(strategy, ChunkingStrategy.PARAGRAPH)
-        
-        manager = get_chunk_manager()
-        chunks = manager.process_document(
-            file_path=file_path,
-            strategy=selected_strategy,
-            preprocess=True
-        )
-        
-        if not chunks:
-            return False, "No text extracted. Is the file empty/scanned?"
-
-        # FIX #1: Tag every chunk with the strategy used
-        for chunk in chunks:
-            chunk.metadata["strategy"] = strategy
-
-        print(f"💾 Indexing {len(chunks)} chunks into Vector DB...")
-        db = get_vectorstore(username)
-        db.add_documents(chunks)
-        
-        tracker.upload_user_db(username)
-        
-        if os.path.exists(file_path):
-            os.remove(file_path)
+        # 1. EXTRACT TEXT (Using doc_loader)
+        # We need a pseudo-object because doc_loader expects a Streamlit object,
+        # but we are reading from disk.
+        with open(file_path, "rb") as f:
+            class FileObj:
+                def __init__(self, f, name):
+                    self.f = f
+                    self.name = name
+                def read(self): return self.f.read()
             
-        return True, f"Successfully added {len(chunks)} chunks to Knowledge Base."
-        
+            file_obj = FileObj(f, os.path.basename(file_path))
+            raw_text = doc_loader.extract_text_from_file(file_obj, use_vision=use_vision, api_key=api_key)
+
+        if not raw_text or not raw_text.strip():
+            return False, "Document appears empty or could not be read."
+
+        # 2. CHUNK TEXT
+        chunks = []
+        if strategy == "paragraph":
+            splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
+            chunks = splitter.split_text(raw_text)
+        elif strategy == "token":
+            splitter = TokenTextSplitter(chunk_size=512, chunk_overlap=50)
+            chunks = splitter.split_text(raw_text)
+        elif strategy == "page":
+            splitter = RecursiveCharacterTextSplitter(chunk_size=2000, chunk_overlap=200)
+            chunks = splitter.split_text(raw_text)
+            
+        # 3. CREATE DOCUMENTS
+        docs = [
+            Document(
+                page_content=chunk, 
+                metadata={"source": os.path.basename(file_path), "strategy": strategy}
+            ) 
+            for chunk in chunks
+        ]
+        
+        # 4. INDEX TO CHROMA
+        if docs:
+            # Use the getter function (Lazy Load)
+            emb_fn = get_embedding_func()
+            db = Chroma(persist_directory=user_db_path, embedding_function=emb_fn)
+            db.add_documents(docs)
+            return True, f"Successfully indexed {len(docs)} chunks from {os.path.basename(file_path)}."
+        else:
+            return False, "No chunks created."
+
     except Exception as e:
-        print(f"❌ Processing Error: {e}")
-        return False, str(e)
+        return False, f"Error processing document: {e}"
 
-# --- RETRIEVAL ENGINE ---
-def search_knowledge_base(query, username, k=6):
+# --- SEARCH PIPELINE (Now with Reranking!) ---
+def search_knowledge_base(query, username, k=10, final_k=4):
     """
-    Two-Stage Retrieval System (RAG):
-    1. Retrieval: Get 10 candidates via fast Vector Search.
-    2. Reranking: Sort them via Cross-Encoder (Slow/Precise).
-    3. Return top k.
+    Retrieves top K chunks, then uses a Cross-Encoder to re-rank them
+    and returns the top final_k most relevant chunks.
     """
+    user_db_path = os.path.join(CHROMA_PATH, username)
+    if not os.path.exists(user_db_path):
+        return []
+        
     try:
-        db = get_vectorstore(username)
-        if db._collection.count() == 0:
+        # 1. INITIAL RETRIEVAL (Vector Similarity)
+        emb_fn = get_embedding_func()
+        db = Chroma(persist_directory=user_db_path, embedding_function=emb_fn)
+        # Fetch more candidates (k=10) to give the reranker options
+        results = db.similarity_search_with_relevance_scores(query, k=k)
+        
+        if not results:
             return []
 
-        reranker = get_reranker_model()
-        
-        # 1. Vector Search (Broad Net)
-        vector_results = db.similarity_search(query, k=25)
+        # 2. RERANKING
+        # Extract just the text for the cross-encoder
+        candidate_docs = [doc for doc, _ in results]
+        candidate_texts = [doc.page_content for doc in candidate_docs]
         
-        # 2. "Poor Man's" Keyword Search (The Safety Net)
-        # We perform a basic text search for unique terms in the query
-        # This catches acronyms like "C&D" if we normalize them
+        # Form pairs: (Query, Document Text)
+        pairs = [[query, text] for text in candidate_texts]
         
-        # Normalize query acronyms (e.g., "C&D" -> "C D")
-        normalized_query = query.replace("&", " ")
-        keyword_results = []
+        # Score pairs
+        reranker = get_rerank_model()
+        scores = reranker.predict(pairs)
         
-        # (Optional: In a production DB like Pinecone/Weaviate, this is built-in.
-        # For Chroma local, we rely on the vector net mostly, but we can 
-        # extend k significantly to catch edge cases).
+        # Attach scores to documents and sort
+        scored_docs = list(zip(candidate_docs, scores))
+        # Sort by score descending (High score = Better match)
+        scored_docs.sort(key=lambda x: x[1], reverse=True)
         
-        # STRATEGY: Just widen the net significantly.
-        # Vector models often hide the match at rank 30 or 40 if the spelling differs.
-        results = db.similarity_search(query, k=50) # Widen from 25 to 50
-        
-        if not results:
-            return []
+        # 3. RETURN TOP N
+        # Return only the document objects of the top final_k
+        final_docs = [doc for doc, score in scored_docs[:final_k]]
+        return final_docs
 
-        # 2. Reranking
-        passages = [doc.page_content for doc in results]
-        ranks = reranker.rank(query, passages)
-        
-        top_results = []
-        sorted_ranks = sorted(ranks, key=lambda x: x['score'], reverse=True)
-
-        # Return the top k results
-        for rank in sorted_ranks[:k]:
-            doc_index = rank['corpus_id']
-            doc = results[doc_index]
-            doc.metadata["relevance_score"] = round(rank['score'], 4)
-            top_results.append(doc)
-            
-        return top_results
-        
     except Exception as e:
-        print(f"⚠️ Search Error (likely empty DB): {e}")
+        print(f"RAG Error: {e}")
         return []
 
+# --- MANAGEMENT UTILS ---
 def list_documents(username):
+    """Returns a list of unique sources in the user's DB."""
+    user_db_path = os.path.join(CHROMA_PATH, username)
+    if not os.path.exists(user_db_path):
+        return []
+        
     try:
-        db = get_vectorstore(username)
-        # Check if empty before fetching to prevent errors
-        if db._collection.count() == 0:
-            return []
-            
-        data = db.get() 
+        emb_fn = get_embedding_func()
+        db = Chroma(persist_directory=user_db_path, embedding_function=emb_fn)
+        data = db.get()
         metadatas = data['metadatas']
         
-        file_stats = {}
-        
-        for meta in metadatas:
-            src = meta.get('source', 'unknown')
-            filename = os.path.basename(src)
-            # FIX #2: Retrieve the strategy (Default to 'unknown' for old docs)
-            strat = meta.get('strategy', 'unknown')
-            
-            if src not in file_stats:
-                file_stats[src] = {
-                    'source': src, 
-                    'filename': filename, 
-                    'chunks': 0, 
-                    'strategy': strat
-                }
-            file_stats[src]['chunks'] += 1
+        inventory = {}
+        for m in metadatas:
+            src = m.get('source', 'Unknown')
+            if src not in inventory:
+                inventory[src] = {"chunks": 0, "strategy": m.get('strategy', 'Unknown')}
+            inventory[src]["chunks"] += 1
             
-        return list(file_stats.values())
-        
-    except Exception as e:
-        print(f"❌ Error listing docs: {e}")
+        return [{"filename": k, "chunks": v["chunks"], "strategy": v["strategy"], "source": k} for k, v in inventory.items()]
+    except:
         return []
 
-def delete_document(username, source_path):
+def delete_document(username, source_name):
+    """Removes all chunks associated with a specific source file."""
+    user_db_path = os.path.join(CHROMA_PATH, username)
     try:
-        print(f"🗑️ Deleting {source_path} for {username}...")
-        db = get_vectorstore(username)
-        db.delete(where={"source": source_path})
-        tracker.upload_user_db(username)
-        return True, f"Deleted {os.path.basename(source_path)}"
+        emb_fn = get_embedding_func()
+        db = Chroma(persist_directory=user_db_path, embedding_function=emb_fn)
+        
+        data = db.get()
+        ids_to_delete = []
+        for i, meta in enumerate(data['metadatas']):
+            if meta.get('source') == source_name:
+                ids_to_delete.append(data['ids'][i])
+                
+        if ids_to_delete:
+            db.delete(ids=ids_to_delete)
+            return True, f"Deleted {source_name}."
+        else:
+            return False, "File not found in index."
     except Exception as e:
-        return False, str(e)
+        return False, f"Delete failed: {e}"
 
 def reset_knowledge_base(username):
-    try:
-        db = get_vectorstore(username)
-        db.delete_collection()
-        tracker.upload_user_db(username)
-        return True, "Knowledge Base completely reset."
-    except Exception as e:
-        return False, str(e)
+    """Wipes the entire user database."""
+    user_db_path = os.path.join(CHROMA_PATH, username)
+    if os.path.exists(user_db_path):
+        shutil.rmtree(user_db_path)
+        return True, "Database Reset."
+    return False, "Database already empty."