app.py

import streamlit as st
import uuid
import os
import base64
import weakref 
from langchain_community.retrievers import BM25Retriever
from langchain_core.messages import HumanMessage, AIMessage

from data_processing import process_and_ingest, SessionDocStore, cleanup_session_index
from rag_engine import run_advanced_rag

# ──────────────────────────────────────────────────────────────────
# AUTOMATIC CLEANUP HANDLER (The Janitor)
# ──────────────────────────────────────────────────────────────────

class SessionJanitor:
    """
    This object lives in the session state.
    When the session ends (Refresh/Close), this object is destroyed.
    The 'weakref.finalize' automatically calls the cleanup function.
    """
    def __init__(self, session_id):
        self.session_id = session_id
        # Register the cleanup function to run when this object dies
        weakref.finalize(self, cleanup_session_index, session_id)
# ──────────────────────────────────────────────────────────────────
# 1. PAGE CONFIG & STATE INITIALIZATION
# ──────────────────────────────────────────────────────────────────

st.set_page_config(
    page_title="Deep RAG Analyzer",
    page_icon="🤖",
    layout="wide"
)

# Initialize Session State
if "session_id" not in st.session_state:
    st.session_state.session_id = str(uuid.uuid4())
    st.session_state.janitor = SessionJanitor(st.session_state.session_id)

if "doc_store" not in st.session_state:
    st.session_state.doc_store = SessionDocStore()

if "messages" not in st.session_state:
    st.session_state.messages = [] # Format: {"role": "user/assistant", "content": "text", "images": []}

if "bm25" not in st.session_state:
    st.session_state.bm25 = None

if "processed_file" not in st.session_state:
    st.session_state.processed_file = None
    
if "uploader_key" not in st.session_state:
    st.session_state.uploader_key = str(uuid.uuid4())

# ──────────────────────────────────────────────────────────────────
# 2. SIDEBAR (Upload & Reset)
# ──────────────────────────────────────────────────────────────────

with st.sidebar:
    st.title("📁 Document Upload")
    uploaded_file = st.file_uploader("Upload PDF", type=["pdf"],key=st.session_state.uploader_key)
    
    if uploaded_file and uploaded_file.name != st.session_state.processed_file:
        with st.spinner("Partitioning & Embedding Generation (This may take around few seconds to few minutes depending upon file size & content)..."):
            # Save to temp file for processing
            temp_path = f"temp_{uploaded_file.name}"
            with open(temp_path, "wb") as f:
                f.write(uploaded_file.getbuffer())
                
            try:
                # Run Pipeline
                documents = process_and_ingest(
                    temp_path, 
                    st.session_state.session_id, 
                    st.session_state.doc_store
                )
                
                # Setup Retriever
                bm25 = BM25Retriever.from_documents(documents)
                bm25.k = 3
                st.session_state.bm25 = bm25
                st.session_state.processed_file = uploaded_file.name
                
                st.success(f"Processed {len(documents)} chunks!")
            except Exception as e:
                st.error(f"Error: {e}")
            finally:
                # Cleanup temp file
                if os.path.exists(temp_path):
                    os.remove(temp_path)

    st.markdown("---")
    if st.button("🗑️ Clear Chat & Reset"):
        try:
            # 1. Attempt Cleanup with visual feedback
            with st.spinner(f"Deleting vector data for {st.session_state.session_id}..."):
                cleanup_session_index(st.session_state.session_id)
            
            # 2. Reset Local State (Only if cleanup succeeded)
            st.session_state.session_id = str(uuid.uuid4())
            st.session_state.doc_store = SessionDocStore()
            st.session_state.messages = []
            st.session_state.bm25 = None
            st.session_state.processed_file = None
            st.session_state.uploader_key = str(uuid.uuid4())
            
            # 3. Re-attach Janitor for new session
            st.session_state.janitor = SessionJanitor(st.session_state.session_id)
            
            st.success("Cache cleared successfully!")
            st.rerun()
            
        except Exception as e:
            st.error(f"Cleanup failed! Check Pinecone Console.\nError: {e}")

    st.markdown("### ℹ️ How to Use")
    st.info(
        """
        1. **Upload a Single PDF only and only a single file type. Multiple PDF's are not accepted**.
        2. Wait for the **"Processed"** success message.
        3. Ask questions in the chat.
        4. The system uses **Hybrid Search** (Keyword + Vector) and **Re-ranking** for accuracy.
        5. Always Click on **Clear Chat** to start a fresh session with a new pdf.
        """
    )

# ──────────────────────────────────────────────────────────────────
# 3. CHAT INTERFACE
# ──────────────────────────────────────────────────────────────────

st.title("🤖 Advanced Multimodal RAG (with chat history)")
st.caption("Using GPT-4.1, IntFloat-e5-basev2 as Embedding model and Pinecone as Vector Database (Do not upload extremly large PDFs!)")

# Display History
for msg in st.session_state.messages:
    with st.chat_message(msg["role"]):
        st.markdown(msg["content"])
        #if "images" in msg and msg["images"]:
            # Display images in a row
            #cols = st.columns(len(msg["images"]))
            #for idx, img_b64 in enumerate(msg["images"]):
                #with cols[idx]:
                    #if "," in img_b64: img_b64 = img_b64.split(",")[1]
                    #st.image(base64.b64decode(img_b64), use_container_width=True)

# Chat Input
if prompt := st.chat_input("Ask about your document..."):
    
    if not st.session_state.bm25:
        st.error("Please upload and process a PDF first !")
        st.stop()

    # 1. Display User Message
    st.chat_message("user").markdown(prompt)
    st.session_state.messages.append({"role": "user", "content": prompt})
    
    # 2. Prepare History for RAG
    lc_history = []
    for m in st.session_state.messages:
        if m["role"] == "user":
            lc_history.append(HumanMessage(content=m["content"]))
        else:
            lc_history.append(AIMessage(content=m["content"]))

    # 3. Generate Response
    with st.chat_message("assistant"):
        with st.spinner("Thinking (retrieving the context)..."):
            try:
                answer = run_advanced_rag(
                    prompt,
                    st.session_state.session_id,
                    st.session_state.bm25,
                    st.session_state.doc_store,
                    lc_history
                )
                
                st.markdown(answer)
                
                # Display Images if found
                #if images:
                    #st.write("---")
                    #st.caption("📸 Retrieved Visual Context:")
                    #cols = st.columns(min(3, len(images)))
                    #for idx, img_b64 in enumerate(images[:3]):
                        #with cols[idx]:
                            #if "," in img_b64: img_b64 = img_b64.split(",")[1]
                            #st.image(base64.b64decode(img_b64), use_container_width=True)
                
                # Save to history
                st.session_state.messages.append({
                    "role": "assistant", 
                    "content": answer
                    #"images": images[:3] if images else []
                })
                
            except Exception as e:
                st.error(f"Error generating response: {e}")