retrive_document_both.py

import os
import datetime
import time
import gradio as gr

from langchain_community.document_loaders import PyPDFLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_community.retrievers import BM25Retriever
from langchain_community.llms import LlamaCpp
from langchain.chains import create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_core.prompts import ChatPromptTemplate

# -------------------------------------------------------------------------
# 1. Global Setup (Models Loaded Once)
# -------------------------------------------------------------------------
print("Initializing Embedding Engine & LLM Runtime...")
# Note: Embeddings are initialized globally but only utilized if Vector mode is active
embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")

MODEL_PATH = "model/gemma-4-E2B-it-Q8_0.gguf"  # Ensure this path is correct
llm = LlamaCpp(
    model_path=MODEL_PATH,
    n_ctx=4096,
    temperature=0.1,
    max_tokens=512,  
    verbose=False
)

# Build the rigid system prompt structure
system_prompt = (
    "You are a strict Document Retrieval analysis assistant.\n"
    "Answer the user's question using ONLY the following pieces of retrieved financial context. "
    "If you do not know the answer or if it is not explicitly stated in the context, "
    "state clearly that the book does not provide this information. Do not make up facts.\n\n"
    "Context:\n{context}"
)
prompt = ChatPromptTemplate.from_messages([
    ("system", system_prompt),
    ("human", "{input}"),
])
question_answer_chain = create_stuff_documents_chain(llm, prompt)

# -------------------------------------------------------------------------
# 2. Core Processing Functions
# -------------------------------------------------------------------------
def process_pdf(pdf_file, rag_mode):
    """Loads and splits the PDF, then builds either a FAISS vector store or a BM25 index."""
    if pdf_file is None:
        return "No file uploaded.", None, ""

    start_time = time.time()
    current_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    
    log_output = f"[{current_time}] 📁 Loading PDF: {os.path.basename(pdf_file.name)}...\n"
    
    try:
        # 1. Document Loading and Splitting (Shared Step)
        loader = PyPDFLoader(pdf_file.name)
        docs = loader.load()
        
        text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
        final_chunks = text_splitter.split_documents(docs)
        log_output += f" -> Successfully generated {len(final_chunks)} semantic chunks.\n"
        
        # 2. Strategy Execution based on Mode Selection
        if rag_mode == "Vector (FAISS)":
            log_output += f" -> [Vector Mode] Generating embeddings & building FAISS database...\n"
            vector_store = FAISS.from_documents(final_chunks, embeddings)
            retriever = vector_store.as_retriever(search_kwargs={"k": 3})
        else:
            log_output += f" -> [Vectorless Mode] Parsing string structures & building BM25 Text index...\n"
            retriever = BM25Retriever.from_documents(final_chunks)
            retriever.k = 3
        
        elapsed = time.time() - start_time
        log_output += f"✅ Success! Engine compiled via {rag_mode} in {elapsed:.2f} seconds."
        
        return log_output, retriever, f"PDF successfully indexed via {rag_mode}! Ready for queries."
        
    except Exception as e:
        return f"❌ Error processing PDF: {str(e)}", None, "Processing failed."

def answer_query(query, retriever, chat_history):
    """Retrieves context using the active retriever engine and runs inference against local LLM."""
    current_time = datetime.datetime.now().strftime("%H:%M:%S")
    
    if retriever is None:
        chat_history.append((query, "⚠️ Please upload and process a document first!"))
        return "", chat_history

    if not query.strip():
        return "", chat_history

    # Establish the unified retrieval chain using whatever retriever object was cached in state
    rag_chain = create_retrieval_chain(retriever, question_answer_chain)
    
    # Run Inference
    response = rag_chain.invoke({"input": query})
    answer = response['answer']
    
    # Format answer with timestamp marker
    formatted_answer = f"[{current_time}]\n{answer}"
    
    chat_history.append((query, formatted_answer))
    return "", chat_history

# -------------------------------------------------------------------------
# 3. Gradio UI Architecture
# -------------------------------------------------------------------------
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    # App State to pass the active retriever instance (either LangChain Vector VectorStoreRetriever OR BM25Retriever)
    active_retriever_state = gr.State(None)
    
    gr.Markdown(
        """
        # 🏦 Multi-Mode Secure Document RAG Engine
        Analyze private documents completely offline. Switch instantly between **Semantic Vector Spaces** or **Vectorless Keyword Layouts**.
        """
    )
    
    with gr.Row():
        # Left Panel: Control Strategy & Document Ingestion
        with gr.Column(scale=1):
            gr.Markdown("### 1. Configuration & Ingestion")
            
            # Mode toggle
            rag_mode_selection = gr.Radio(
                choices=["Vector (FAISS)", "Vectorless (BM25)"],
                value="Vector (FAISS)",
                label="Retrieval Engine Strategy"
            )
            
            file_input = gr.File(label="Upload PDF / Book", file_types=[".pdf"])
            process_btn = gr.Button("Build Retrieval Knowledge Base", variant="primary")
            
            status_text = gr.Textbox(label="Engine Status", interactive=False, value="Awaiting document...")
            log_monitor = gr.TextArea(label="Pipeline Terminal Logs", interactive=False, lines=6)
            
        # Right Panel: Interactive Chat Interface
        with gr.Column(scale=2):
            gr.Markdown("### 2. Conversational Analysis")
            chatbot = gr.Chatbot(label="Document Retrieval Chat", height=470)
            
            with gr.Row():
                query_input = gr.Textbox(
                    label="Ask a document retrieval question...",
                    placeholder="e.g., What is the primary difference between working capital and capital budgeting?",
                    scale=4
                )
                submit_btn = gr.Button("Submit Query", variant="secondary", scale=1)

    # --- UI Event Wire-up ---
    
    # File processing sequence (Passes file and current mode selection)
    process_btn.click(
        fn=process_pdf,
        inputs=[file_input, rag_mode_selection],
        outputs=[log_monitor, active_retriever_state, status_text]
    )
    
    # Query execution sequences (Triggers on Click or hitting Enter)
    submit_btn.click(
        fn=answer_query,
        inputs=[query_input, active_retriever_state, chatbot],
        outputs=[query_input, chatbot]
    )
    query_input.submit(
        fn=answer_query,
        inputs=[query_input, active_retriever_state, chatbot],
        outputs=[query_input, chatbot]
    )

# Launch local server
if __name__ == "__main__":
    demo.launch(server_name="127.0.0.1", server_port=7860, share=False)