rag/chain.py

import os
from typing import List, Dict, Any, Tuple
from langchain_core.documents import Document
from models.llm import GroqLLM
from rag.retriever import AdvancedRetriever
from rag.memory import ConversationMemory
from rag.prompt import QA_PROMPT, REWRITE_PROMPT, MAP_PROMPT, REDUCE_PROMPT
from loaders.pdf_loader import load_multiple_pdfs
from chunking.splitter import split_documents
from utils.file_registry import (
    compute_file_hash,
    load_registry,
    register_file,
    is_duplicate,
    clear_registry,
)
import time

class RAGChain:
    """
    Main orchestration class for the RAG pipeline.
    Handles history-aware retrieval, generation, map-reduce summarization,
    and document indexing to unify the backend logic for Gradio and FastAPI.
    """
    def __init__(self, retriever: AdvancedRetriever, memory: ConversationMemory):
        # We lazy-load the LLM so the Gradio UI can boot up quickly.
        self.llm_wrapper = None
        self.langchain_llm = None
        
        self.retriever = retriever
        self.memory = memory
        self.indexed_documents: List[Document] = []
        
    def ensure_llm(self):
        """Lazy loader for the LLM."""
        if self.llm_wrapper is None:
            print("Initializing Groq Llama 3.1 8B...")
            self.llm_wrapper = GroqLLM()
            self.langchain_llm = self.llm_wrapper.get_llm()
            print("Groq Llama 3.1 8B initialized.")

    def index_documents(self, filepaths: List[str]) -> Dict[str, Any]:
        """
        Hash-guarded ingestion pipeline.

        For every supplied filepath:
        1. Compute SHA-256 hash.
        2. Look up hash in the on-disk registry.
        3. Skip the file (no chunking, no embedding, no FAISS insertion) if
           it has already been indexed.
        4. Otherwise: load → chunk → embed → insert into FAISS → register.

        Returns a summary dict with per-file status so callers can build
        informative UI messages.
        """
        registry = load_registry()

        skipped: List[str] = []
        indexed: List[str] = []
        total_new_chunks = 0

        vectorstore = self.retriever.vectorstore

        for filepath in filepaths:
            filename = os.path.basename(filepath)
            file_hash = compute_file_hash(filepath)

            if is_duplicate(registry, file_hash):
                skipped.append(filename)
                continue

            # --- New file: run the full ingestion pipeline ---
            docs = load_multiple_pdfs([filepath])
            if not docs:
                # Non-extractable PDF; skip silently rather than aborting the
                # whole batch.
                skipped.append(filename)
                continue

            self.indexed_documents.extend(docs)
            chunks = split_documents(docs)

            if vectorstore.vectorstore is None:
                vectorstore.create_index(chunks)
            else:
                vectorstore.vectorstore.add_documents(chunks)

            vectorstore.save_index()

            # Persist hash → metadata so subsequent uploads are detected.
            register_file(registry, file_hash, filename, len(chunks))

            indexed.append(filename)
            total_new_chunks += len(chunks)

        # --- Build a human-readable summary message ---
        parts: List[str] = []
        if indexed:
            parts.append(
                f"Indexed {len(indexed)} new file(s) into {total_new_chunks} chunk(s): "
                + ", ".join(indexed)
            )
        if skipped:
            parts.append(
                f"Skipped {len(skipped)} duplicate file(s): " + ", ".join(skipped)
            )
        if not parts:
            parts.append("No files were processed.")

        return {
            "message": " | ".join(parts),
            "num_new_chunks": total_new_chunks,
            "indexed": indexed,
            "skipped": skipped,
        }

    def rewrite_query(self, query: str) -> str:
        """Rewrites the query into a standalone question using chat history."""
        self.ensure_llm()
        history_str = self.memory.get_history_string()
        if history_str == "No previous history.":
            return query # No rewrite needed for the first turn
            
        prompt_text = REWRITE_PROMPT.format(chat_history=history_str, question=query)
        rewritten = self.langchain_llm.invoke(prompt_text)
        # Clean up output if the model rambles
        content = rewritten.content if hasattr(rewritten, 'content') else str(rewritten)
        return content.strip().split("\n")[0]

    def ask(self, query: str, stream: bool = True) -> Dict[str, Any]:
        """
        Main Q&A function.
        1. Rewrite query if there's history.
        2. Retrieve and rerank context.
        3. Generate answer (streaming or synchronous).
        """
        self.ensure_llm()
        start_time = time.time()
        
        # 1. Query Rewrite
        standalone_query = self.rewrite_query(query)
        
        # 2. Retrieval
        docs_with_scores = self.retriever.retrieve(standalone_query)
        context_str = self.retriever.format_docs(docs_with_scores) if docs_with_scores else ""
        
        # 3. Prompt Construction
        prompt_text = QA_PROMPT.format(context=context_str, question=standalone_query)
        
        response_data = {
            "query": query,
            "standalone_query": standalone_query,
            "context_docs": docs_with_scores,
            "time_taken": 0.0
        }
        
        # 4. Generation
        if stream:
            # Note: For streaming, we return the streamer object. The caller needs to iterate over it
            # and update the memory after streaming finishes.
            streamer = self.llm_wrapper.generate_stream(prompt_text)
            response_data["streamer"] = streamer
        else:
            answer = self.langchain_llm.invoke(prompt_text)
            answer_text = answer.content if hasattr(answer, 'content') else str(answer)
            response_data["answer"] = answer_text.strip()
            self.memory.add_user_message(query)
            self.memory.add_assistant_message(response_data["answer"])
            
        response_data["time_taken"] = time.time() - start_time
        return response_data
        
    def summarize(self) -> str:
        """
        Map-Reduce Summarization:
        1. Map: Summarize each chunk individually.
        2. Reduce: Combine chunk summaries into a final structured summary.
        """
        if not self.indexed_documents:
            raise ValueError("No documents indexed to summarize.")
            
        self.ensure_llm()
        
        # 1. Map phase
        intermediate_summaries = []
        for doc in self.indexed_documents:
            map_prompt_text = MAP_PROMPT.format(text=doc.page_content)
            chunk_summary = self.langchain_llm.invoke(map_prompt_text)
            chunk_text = chunk_summary.content if hasattr(chunk_summary, 'content') else str(chunk_summary)
            intermediate_summaries.append(chunk_text.strip())
            
        # 2. Reduce phase
        combined_summaries = "\n\n".join(intermediate_summaries)
        reduce_prompt_text = REDUCE_PROMPT.format(text=combined_summaries)
        
        final_summary = self.langchain_llm.invoke(reduce_prompt_text)
        final_text = final_summary.content if hasattr(final_summary, 'content') else str(final_summary)
        return final_text.strip()


# ---------------------------------------------------------------------------
# Module-level utility — does NOT touch LLM, memory, or retriever wiring
# ---------------------------------------------------------------------------

def reset_vector_store(rag_chain: "RAGChain") -> None:
    """
    Hard-reset the vector database.

    Actions performed (in order):
    1. Delete every file inside the FAISS index directory.
    2. Remove the processed-files registry (processed_files.json).
    3. Set the in-memory vectorstore to None.
    4. Clear the cached list of indexed documents.

    The Groq LLM, conversation memory, embeddings model, reranker, and all
    UI components are left entirely unaffected.

    Args:
        rag_chain: The live RAGChain instance to reset.
    """
    vs_wrapper = rag_chain.retriever.vectorstore
    index_path = vs_wrapper.index_path

    # 1. Delete FAISS index files from disk
    if os.path.isdir(index_path):
        for fname in os.listdir(index_path):
            fpath = os.path.join(index_path, fname)
            if os.path.isfile(fpath):
                os.remove(fpath)

    # 2. Delete the hash registry
    clear_registry()

    # 3. Reset in-memory state
    vs_wrapper.vectorstore = None
    rag_chain.indexed_documents = []