Update app.py

app.py

@@ -2,12 +2,27 @@ 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
 # ──────────────────────────────────────────────────────────────────
@@ -21,6 +36,7 @@ st.set_page_config(
 # 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()