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}")