Update src/rag_engine.py

increased k from 3 retrieved documents to 10

src/rag_engine.py

@@ -1,198 +1,198 @@
-import os
-from langchain_chroma import Chroma
-from langchain_huggingface import HuggingFaceEmbeddings
-from sentence_transformers import CrossEncoder
-from core.ChunkingManager import ChunkingManager, ChunkingStrategy
-import tracker # To trigger syncs
-
-# --- 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-L6-v2"
-RERANKER_MODEL_NAME = "cross-encoder/ms-marco-MiniLM-L-6-v2" 
-
-# --- LAZY LOADING SINGLETONS ---
-# We use these globals to store the models once loaded, so we don't reload them 
-# every time a function is called, but we also don't load them on import.
-_embedding_fn = None
-_reranker = None
-_chunk_manager = None
-
-def get_embedding_function():
-    """Lazy loads the embedding model only when needed."""
-    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():
-    """Lazy loads the CrossEncoder only when needed."""
-    global _reranker
-    if _reranker is None:
-        print("⚙️ Loading Reranker Model...")
-        _reranker = CrossEncoder(RERANKER_MODEL_NAME)
-    return _reranker
-
-def get_chunk_manager():
-    """Lazy loads the Chunking 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):
-    """Returns the persistent ChromaDB for a SPECIFIC USER."""
-    # Safety: Ensure username doesn't contain path traversal characters
-    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}" 
-    )
-
-def save_uploaded_file(uploaded_file):
-    """Saves upload to temp, sanitizing the filename."""
-    if not os.path.exists(UPLOAD_DIR):
-        os.makedirs(UPLOAD_DIR)
-        
-    # SECURITY FIX: Sanitize filename to prevent directory traversal
-    safe_filename = os.path.basename(uploaded_file.name)
-    file_path = os.path.join(UPLOAD_DIR, safe_filename)
-    
-    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"):
-    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)
-        
-        # Use the lazy-loaded chunk manager
-        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?"
-
-        print(f"💾 Indexing {len(chunks)} chunks into Vector DB...")
-        db = get_vectorstore(username)
-        db.add_documents(chunks)
-        
-        # Sync immediately
-        tracker.upload_user_db(username)
-        
-        if os.path.exists(file_path):
-            os.remove(file_path)
-            
-        return True, f"Successfully added {len(chunks)} chunks to Knowledge Base."
-        
-    except Exception as e:
-        print(f"❌ Processing Error: {e}")
-        return False, str(e)
-
-# --- RETRIEVAL ENGINE ---
-def search_knowledge_base(query, username, k=3):
-    """
-    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.
-    """
-    db = get_vectorstore(username)
-    reranker = get_reranker_model()
-    
-    # 1. Broad Search (Retrieve more than needed to filter later)
-    results = db.similarity_search(query, k=10)
-    
-    if not results:
-        return []
-
-    # 2. Reranking
-    # Prepare pairs: [[Query, Text1], [Query, Text2]...]
-    passages = [doc.page_content for doc in results]
-    ranks = reranker.rank(query, passages)
-    
-    # 3. Sort and Filter
-    # Reranker returns list of dicts: {'corpus_id': 0, 'score': 0.9}
-    top_results = []
-    
-    # Sort ranks by score descending just to be safe (though .rank() usually sorts)
-    sorted_ranks = sorted(ranks, key=lambda x: x['score'], reverse=True)
-    
-    for rank in sorted_ranks[:k]:
-        doc_index = rank['corpus_id']
-        doc = results[doc_index]
-        # Append score for transparency
-        doc.metadata["relevance_score"] = round(rank['score'], 4)
-        top_results.append(doc)
-        
-    return top_results
-
-def list_documents(username):
-    """
-    Returns a list of unique files currently in the user's database.
-    WARNING: This pulls all metadata. Performance degrades >10k chunks.
-    """
-    try:
-        db = get_vectorstore(username)
-        data = db.get() 
-        metadatas = data['metadatas']
-        
-        file_stats = {}
-        
-        for meta in metadatas:
-            src = meta.get('source', 'unknown')
-            filename = os.path.basename(src)
-            
-            if src not in file_stats:
-                file_stats[src] = {'source': src, 'filename': filename, 'chunks': 0}
-            file_stats[src]['chunks'] += 1
-            
-        return list(file_stats.values())
-        
-    except Exception as e:
-        print(f"❌ Error listing docs: {e}")
-        return []
-
-def delete_document(username, source_path):
-    """Removes all chunks associated with a specific source file."""
-    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)}"
-        
-    except Exception as e:
-        return False, str(e)
-
-def reset_knowledge_base(username):
-    """Nuke option: Clears the entire database for the user."""
-    try:
-        db = get_vectorstore(username)
-        db.delete_collection()
-        tracker.upload_user_db(username)
-        return True, "Knowledge Base completely reset."
-    except Exception as e:
+import os
+from langchain_chroma import Chroma
+from langchain_huggingface import HuggingFaceEmbeddings
+from sentence_transformers import CrossEncoder
+from core.ChunkingManager import ChunkingManager, ChunkingStrategy
+import tracker # To trigger syncs
+
+# --- 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-L6-v2"
+RERANKER_MODEL_NAME = "cross-encoder/ms-marco-MiniLM-L-6-v2" 
+
+# --- LAZY LOADING SINGLETONS ---
+# We use these globals to store the models once loaded, so we don't reload them 
+# every time a function is called, but we also don't load them on import.
+_embedding_fn = None
+_reranker = None
+_chunk_manager = None
+
+def get_embedding_function():
+    """Lazy loads the embedding model only when needed."""
+    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():
+    """Lazy loads the CrossEncoder only when needed."""
+    global _reranker
+    if _reranker is None:
+        print("⚙️ Loading Reranker Model...")
+        _reranker = CrossEncoder(RERANKER_MODEL_NAME)
+    return _reranker
+
+def get_chunk_manager():
+    """Lazy loads the Chunking 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):
+    """Returns the persistent ChromaDB for a SPECIFIC USER."""
+    # Safety: Ensure username doesn't contain path traversal characters
+    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}" 
+    )
+
+def save_uploaded_file(uploaded_file):
+    """Saves upload to temp, sanitizing the filename."""
+    if not os.path.exists(UPLOAD_DIR):
+        os.makedirs(UPLOAD_DIR)
+        
+    # SECURITY FIX: Sanitize filename to prevent directory traversal
+    safe_filename = os.path.basename(uploaded_file.name)
+    file_path = os.path.join(UPLOAD_DIR, safe_filename)
+    
+    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"):
+    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)
+        
+        # Use the lazy-loaded chunk manager
+        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?"
+
+        print(f"💾 Indexing {len(chunks)} chunks into Vector DB...")
+        db = get_vectorstore(username)
+        db.add_documents(chunks)
+        
+        # Sync immediately
+        tracker.upload_user_db(username)
+        
+        if os.path.exists(file_path):
+            os.remove(file_path)
+            
+        return True, f"Successfully added {len(chunks)} chunks to Knowledge Base."
+        
+    except Exception as e:
+        print(f"❌ Processing Error: {e}")
+        return False, str(e)
+
+# --- RETRIEVAL ENGINE ---
+def search_knowledge_base(query, username, k=10):
+    """
+    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.
+    """
+    db = get_vectorstore(username)
+    reranker = get_reranker_model()
+    
+    # 1. Broad Search (Retrieve more than needed to filter later)
+    results = db.similarity_search(query, k=10)
+    
+    if not results:
+        return []
+
+    # 2. Reranking
+    # Prepare pairs: [[Query, Text1], [Query, Text2]...]
+    passages = [doc.page_content for doc in results]
+    ranks = reranker.rank(query, passages)
+    
+    # 3. Sort and Filter
+    # Reranker returns list of dicts: {'corpus_id': 0, 'score': 0.9}
+    top_results = []
+    
+    # Sort ranks by score descending just to be safe (though .rank() usually sorts)
+    sorted_ranks = sorted(ranks, key=lambda x: x['score'], reverse=True)
+    
+    for rank in sorted_ranks[:k]:
+        doc_index = rank['corpus_id']
+        doc = results[doc_index]
+        # Append score for transparency
+        doc.metadata["relevance_score"] = round(rank['score'], 4)
+        top_results.append(doc)
+        
+    return top_results
+
+def list_documents(username):
+    """
+    Returns a list of unique files currently in the user's database.
+    WARNING: This pulls all metadata. Performance degrades >10k chunks.
+    """
+    try:
+        db = get_vectorstore(username)
+        data = db.get() 
+        metadatas = data['metadatas']
+        
+        file_stats = {}
+        
+        for meta in metadatas:
+            src = meta.get('source', 'unknown')
+            filename = os.path.basename(src)
+            
+            if src not in file_stats:
+                file_stats[src] = {'source': src, 'filename': filename, 'chunks': 0}
+            file_stats[src]['chunks'] += 1
+            
+        return list(file_stats.values())
+        
+    except Exception as e:
+        print(f"❌ Error listing docs: {e}")
+        return []
+
+def delete_document(username, source_path):
+    """Removes all chunks associated with a specific source file."""
+    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)}"
+        
+    except Exception as e:
+        return False, str(e)
+
+def reset_knowledge_base(username):
+    """Nuke option: Clears the entire database for the user."""
+    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)