Update src/streamlit_app.py

src/streamlit_app.py

@@ -2,271 +2,246 @@
 import os
 import re
 import streamlit as st
-import torch
+import shutil
+import hashlib
+import pickle
 
-# ==========================================================
-# ✅ PAGE CONFIG
-# ==========================================================
+# Protect against pages that call heavy imports before environment set
 st.set_page_config(page_title="Enterprise Knowledge Assistant", layout="wide")
-print("CUDA available:", torch.cuda.is_available())
 
-# ==========================================================
-# ⚙️ CACHE DIR
-# ==========================================================
-CACHE_DIR = "/tmp/hf_cache"
-os.makedirs(CACHE_DIR, exist_ok=True)
-os.environ.update({
-    "HF_HOME": CACHE_DIR,
-    "TRANSFORMERS_CACHE": CACHE_DIR,
-    "HF_DATASETS_CACHE": CACHE_DIR,
-    "HF_MODULES_CACHE": CACHE_DIR,
-})
-
-# ==========================================================
-# 📦 IMPORTS
+# --- Local helpers / small cache cleanup (keeps /tmp sane) ---
+def clean_cache(max_size_gb: float = 2.0):
+    folders = [
+        "/root/.cache/huggingface",
+        "/root/.cache/transformers",
+        "/root/.cache/torch",
+    ]
+    for folder in folders:
+        if os.path.exists(folder):
+            try:
+                size_gb = sum(
+                    os.path.getsize(os.path.join(dp, f))
+                    for dp, _, files in os.walk(folder)
+                    for f in files
+                ) / (1024**3)
+                if size_gb > max_size_gb:
+                    shutil.rmtree(folder, ignore_errors=True)
+            except Exception:
+                pass
+
+clean_cache()
+
+# ==========================================================
+# Imports (after HF cache set to avoid extra downloads)
 # ==========================================================
 from ingestion import extract_text_from_pdf, chunk_text
 from vectorstore import build_faiss_index
 from qa import retrieve_chunks, generate_answer, cache_embeddings, embed_chunks, genai_generate
 
 # ==========================================================
-# 🧠 SMART SUGGESTION GENERATOR
-# ==========================================================
-def generate_dynamic_suggestions_from_toc(toc, chunks, doc_name="Document"):
-    if not toc or not chunks:
-        return []
-    titles = []
-    for sec, raw_title in toc:
-        title = re.sub(r"^\s*[\dA-Za-z.\-]+\s*", "", raw_title)
-        title = re.sub(r"\.{2,}\s*\d+$", "", title).strip()
-        if 4 < len(title) < 120:
-            titles.append(title)
-    context_sample = " ".join(chunks[:3])[:4000]
-    prompt = f"""
-    You are generating concise, context-aware questions based on the document "{doc_name}".
-    Use this Table of Contents and sample content for inspiration.
-
-    TABLE OF CONTENTS:
-    {chr(10).join(['- ' + t for t in titles[:8]])}
-
-    TEXT SAMPLE:
-    {context_sample}
-
-    Generate 5–7 short, relevant, strictly document-based questions.
-    Each question should be under 18 words.
-    """
-    try:
-        ai_response = genai_generate(prompt)
-        questions = re.findall(r"[-•]?\s*(.+?)\?", ai_response)
-        clean_qs = [q.strip("•-— ").strip() + "?" for q in questions if 8 < len(q) < 120]
-        seen, final = set(), []
-        for q in clean_qs:
-            if q.lower() not in seen:
-                seen.add(q.lower())
-                final.append(q)
-        return final[:7]
-    except Exception:
-        return ["What is this document about?", "How do I begin using this guide?"]
-
-# ==========================================================
-# 🎨 STYLING — MINIMAL, ENTERPRISE UI
+# SESSION - initialize keys early to avoid widget duplication
 # ==========================================================
-st.markdown("""
-<style>
-div.block-container {padding-top: 1.2rem; max-width: 1080px;}
-h1, h2, h3 {color: #f3f4f6; font-weight: 600;}
-.suggest-chip {
-    background: #0f1724;
-    border: 1px solid #374151;
-    border-radius: 14px;
-    color: #e6eef8;
-    padding: 8px 12px;
-    cursor: pointer;
-    font-size: 13px;
-    margin: 6px 6px 6px 0;
-    display: inline-block;
-    transition: background 0.2s, transform 0.1s;
-}
-.suggest-chip:hover {background: #1e3a8a; transform: translateY(-2px);}
-.answer-box {
-    background: linear-gradient(180deg,#0b1220,#071027);
-    border-left: 4px solid #3b82f6;
-    border-radius: 8px;
-    padding: 16px 18px;
-    color: #e6eef8;
-    margin-top: 12px;
-    box-shadow: 0 4px 14px rgba(0,0,0,0.35);
-}
-.stTextInput > div > div > input {
-    background-color: #0f172a !important;
-    color: #f1f5f9 !important;
-    border-radius: 6px !important;
-    border: 1px solid #334155 !important;
-    padding: 8px 10px !important;
-}
-</style>
-""", unsafe_allow_html=True)
-
-# ==========================================================
-# 🧠 SESSION STATE (initialize before widgets)
-# ==========================================================
-for key, val in {
+initial_state = {
     "user_query_input": "",
     "show_more": False,
     "selected_suggestion": None,
     "query_suggestions_fixed": None,
     "last_doc": None,
-    "show_advanced": False,  # keep persistent checkbox state
-}.items():
-    if key not in st.session_state:
-        st.session_state[key] = val
-
+    "show_advanced": False,
+    "chunk_size": 1000,
+    "overlap": 120,
+    "top_k": 5,
+    "reasoning_mode": False,
+}
+for k, v in initial_state.items():
+    if k not in st.session_state:
+        st.session_state[k] = v
 
 # ==========================================================
-# 🧭 SIDEBAR — STABLE VERSION (no rerun reloads)
+# Sidebar (stable, no unexpected reruns)
 # ==========================================================
 with st.sidebar:
     st.markdown("### 🧭 Response Mode")
-    mode = st.radio(
+    st.session_state["reasoning_mode"] = st.radio(
         "",
         ("Strict (Document-only)", "Extended (Document + general)"),
         index=0,
-        help="Strict = answers only from the document. Extended = may include general context.",
-        key="response_mode"
-    )
+        help="Strict = answers only from the document. Extended = may include general context."
+    ) == "Extended (Document + general)"
 
     st.markdown("---")
+    # Avoid forcing reruns when toggled: Streamlit manages state
+    st.checkbox("Show advanced settings (for developers)", key="show_advanced")
 
-    # ✅ Checkbox linked to session state (no rerun logic, Streamlit handles it automatically)
-    st.checkbox(
-        "Show advanced settings (for developers)",
-        key="show_advanced",
-        help="Toggle developer settings without reloading document or clearing outputs."
-    )
-
-    # ✅ Sliders appear conditionally, state persists
     if st.session_state.show_advanced:
         st.markdown("### Developer Settings")
-        chunk_size = st.slider("Chunk Size (characters)", 200, 1500, 1000, step=50, key="chunk_slider")
-        overlap = st.slider("Chunk Overlap (characters)", 50, 200, 120, step=10, key="overlap_slider")
-        top_k = st.slider("Top K Results", 1, 10, 5, key="topk_slider")
-    else:
-        chunk_size, overlap, top_k = 1000, 120, 5
-
+        # Persist slider states to session_state keys to avoid duplication
+        st.session_state["chunk_size"] = st.slider(
+            "Chunk Size (characters)", 200, 1500, st.session_state["chunk_size"], step=50, key="chunk_slider"
+        )
+        st.session_state["overlap"] = st.slider(
+            "Chunk Overlap (characters)", 50, 200, st.session_state["overlap"], step=10, key="overlap_slider"
+        )
+        st.session_state["top_k"] = st.slider(
+            "Top K Results", 1, 10, st.session_state["top_k"], key="topk_slider"
+        )
     st.markdown("---")
     st.caption("✨ Built by Shubham Sharma")
 
-
 # ==========================================================
-# 🧩 QUERY HANDLER
+# Page header
 # ==========================================================
-def set_user_query(q, idx):
-    st.session_state["user_query_input"] = q
-    st.session_state["selected_suggestion"] = idx
-    # ⚠️ No explicit rerun — Streamlit does this automatically
-
+st.title("📄 Enterprise Knowledge Assistant")
+st.caption("Ask questions about SAP documentation and enterprise PDFs — powered by reasoning and retrieval.")
 
 # ==========================================================
-# 🧠 USER INPUT
+# Document selection + upload
 # ==========================================================
-user_query = st.text_input(
-    "Type your question or click one above:",
-    key="user_query_input"
+doc_choice = st.radio(
+    "Select a document:",
+    ("-- Select --", "Sample PDF", "Upload Custom PDF"),
+    index=0,
+    key="doc_choice"
 )
 
+temp_path = None
+if doc_choice == "-- Select --":
+    st.info("⬅️ Please choose a document from the sidebar to begin.")
+elif doc_choice == "Sample PDF":
+    temp_path = os.path.join(os.path.dirname(__file__), "sample.pdf")
+    st.success("📘 Using built-in Sample PDF.")
+else:
+    uploaded_file = st.file_uploader("📂 Upload your PDF", type="pdf")
+    if uploaded_file:
+        temp_path = os.path.join("/tmp", uploaded_file.name)
+        with open(temp_path, "wb") as f:
+            f.write(uploaded_file.getbuffer())
+        st.success(f"✅ '{uploaded_file.name}' uploaded successfully.")
 
 # ==========================================================
-# 📄 MAIN LAYOUT
+# Utility: set query (used by suggestion buttons)
 # ==========================================================
-st.title("📄 Enterprise Knowledge Assistant")
-st.caption("Query SAP documentation and enterprise PDFs — powered by reasoning and retrieval.")
-
-doc_choice = st.radio(
-    "Select a document:",
-    ["-- Select --", "Sample PDF", "Upload Custom PDF"],
-    index=0
+def set_user_query(q: str, idx: int = None):
+    # Only set values and rerun once to ensure text_input reflects it
+    st.session_state["user_query_input"] = q
+    st.session_state["selected_suggestion"] = idx
+    # Trigger a rerun so the single text_input shows the new value and downstream code runs
+    st.rerun()
+
+# ==========================================================
+# Document processing and suggestion generation
+# ==========================================================
+text, chunks, index, embeddings, toc = None, None, None, None, None
+if temp_path:
+    # If user switched files, clear cached suggestions so they regenerate once for this doc
+    if st.session_state["last_doc"] != os.path.basename(temp_path):
+        st.session_state["query_suggestions_fixed"] = None
+        st.session_state["selected_suggestion"] = None
+        st.session_state["last_doc"] = os.path.basename(temp_path)
+
+    with st.spinner("🔍 Processing your document..."):
+        # Extract -> chunk -> suggestions
+        text, toc = extract_text_from_pdf(temp_path)
+        chunks = chunk_text(text, chunk_size=st.session_state["chunk_size"], overlap=st.session_state["overlap"])
+        st.success("✅ Document loaded successfully.")
+        # Generate suggestions only once per document and store in session_state
+        if st.session_state["query_suggestions_fixed"] is None:
+            # Prefer genai (GPT) suggestions if available in qa.genai_generate
+            try:
+                # build prompt from toc or chunks (keeps consistent)
+                titles = []
+                if toc:
+                    for sec, raw_title in toc:
+                        t = re.sub(r"^\s*[\dA-Za-z.\-]+\s*", "", raw_title)
+                        t = re.sub(r"\.{2,}\s*\d+$", "", t).strip()
+                        if 4 < len(t) < 120:
+                            titles.append(t)
+                sample_text = " ".join(chunks[:4])[:4000]
+                prompt = f"Generate 6 short (<=18 words), document-focused questions based on TOC: {titles[:6]} and sample: {sample_text}"
+                ai_out = genai_generate(prompt)
+                # Parse lines/questions from model response robustly
+                qs = re.findall(r"(?m)^\s*[-•\d\)]*\s*(.+?)\?$", ai_out)
+                qs = [q.strip() + "?" for q in qs if 8 < len(q.strip()) < 120]
+                if not qs:
+                    # fallback: simple heuristics from TOC or chunk sentences
+                    qs = []
+                    if titles:
+                        for t in titles[:6]:
+                            qs.append(f"What is described in '{t}'?")
+                    else:
+                        sents = re.split(r'(?<=[.?!])\s+', sample_text)
+                        for s in sents[:6]:
+                            if len(s) > 20:
+                                qs.append((s.strip()[:80] + "...").strip() + "?")
+                    qs = qs[:6]
+                st.session_state["query_suggestions_fixed"] = qs
+            except Exception:
+                # deterministic fallback
+                st.session_state["query_suggestions_fixed"] = [
+                    "What is this document about?",
+                    "How do I get started with the process?",
+                    "What are the prerequisites?",
+                    "What steps are required?",
+                    "How to troubleshoot common issues?",
+                    "Where can I find the configuration?"
+                ]
+
+    # Build or load embeddings and index (embedding cache uses chunk-aware key)
+    with st.spinner("⚙️ Preparing search index..."):
+        embeddings = cache_embeddings(os.path.basename(temp_path), chunks, lambda c: embed_chunks(c))
+        index = build_faiss_index(embeddings)
+        st.success("🚀 Document ready — you can now ask questions below.")
+
+# ==========================================================
+# UI: Suggested Questions + Query input (only one text_input)
+# ==========================================================
+st.markdown("## 🤖 Ask the Assistant")
+
+if temp_path and st.session_state["query_suggestions_fixed"]:
+    st.markdown("#### 💡 Suggested Questions")
+    suggestions = st.session_state["query_suggestions_fixed"]
+    visible = suggestions if st.session_state["show_more"] else suggestions[:3]
+    cols = st.columns(min(3, len(visible)))
+    for i, q in enumerate(visible):
+        col = cols[i % 3]
+        # unique keys for each button
+        if col.button(f"🔍 {q}", key=f"suggest_{i}"):
+            set_user_query(q, idx=i)
+
+    toggle_text = "Show less ▲" if st.session_state["show_more"] else "Show more ▼"
+    if st.button(toggle_text, key="toggle_show_more"):
+        st.session_state["show_more"] = not st.session_state["show_more"]
+
+# --- single text input (only place we define it) ---
+user_query = st.text_input(
+    "Type your question or pick one above:",
+    value=st.session_state["user_query_input"],
+    key="user_query_input",
 )
 
 # ==========================================================
-# 📂 DOCUMENT HANDLING
+# Answer generation (runs only when user_query non-empty)
 # ==========================================================
-if doc_choice == "-- Select --":
-    st.info("⬅️ Select a document to begin.")
-else:
-    if doc_choice == "Sample PDF":
-        temp_path = os.path.join(os.path.dirname(__file__), "sample.pdf")
-        st.success("📘 Sample PDF loaded successfully. Ask questions below.")
-    else:
-        uploaded_file = st.file_uploader("", type="pdf", label_visibility="collapsed")
-        if uploaded_file:
-            temp_path = os.path.join("/tmp", uploaded_file.name)
-            with open(temp_path, "wb") as f:
-                f.write(uploaded_file.getbuffer())
-            st.success("✅ Document loaded successfully. You can now ask questions below.")
-        else:
-            temp_path = None
-
-    if temp_path:
-        with st.spinner("🔍 Processing document..."):
-            text, toc = extract_text_from_pdf(temp_path)
-            chunks = chunk_text(text, chunk_size=chunk_size)
-
-        with st.spinner("⚙️ Preparing index..."):
-            embeddings = cache_embeddings(os.path.basename(temp_path), chunks, embed_chunks, chunk_size=chunk_size, overlap=overlap)
-            index = build_faiss_index(embeddings)
-
-        st.success("✅ Ready. Ask below!")
-
-        doc_name = os.path.basename(temp_path)
-        if st.session_state["last_doc"] != doc_name:
-            query_suggestions = generate_dynamic_suggestions_from_toc(toc, chunks, doc_name)
-            st.session_state["query_suggestions_fixed"] = query_suggestions
-            st.session_state["last_doc"] = doc_name
-        else:
-            query_suggestions = st.session_state["query_suggestions_fixed"]
-
-        # ----------------------------------------------------------
-        # 💬 ASK ASSISTANT SECTION
-        # ----------------------------------------------------------
-        st.markdown("### Ask the Assistant")
-
-        if query_suggestions:
-            visible = query_suggestions if st.session_state["show_more"] else query_suggestions[:3]
-            cols = st.columns(min(3, len(visible)))
-            for i, q in enumerate(visible):
-                if cols[i % 3].button(f"🔍 {q}", key=f"sugg_{i}"):
-                    set_user_query(q, i)
+if temp_path and user_query and user_query.strip():
+    st.caption("Mode: 📄 Strict Document" if not st.session_state["reasoning_mode"] else "Mode: 🧠 Extended")
+    with st.spinner("💭 Generating your answer..."):
+        retrieved = retrieve_chunks(user_query, index, chunks, top_k=st.session_state["top_k"], embeddings=embeddings)
+        answer = generate_answer(user_query, retrieved, reasoning_mode=st.session_state["reasoning_mode"])
 
-            toggle_text = "Show less ▲" if st.session_state["show_more"] else "Show more ▼"
-            if st.button(toggle_text):
-                st.session_state["show_more"] = not st.session_state["show_more"]
-                st.experimental_rerun()
+    st.markdown("### ✅ Assistant’s Answer")
+    st.markdown(f"<div style='background:#0f172a;border-left:4px solid #3b82f6;padding:12px;border-radius:8px;color:#f1f5f9'>{answer}</div>", unsafe_allow_html=True)
 
-        user_query = st.text_input("Type your question or click one above:", key="user_query_input")
+    with st.expander("📄 Supporting Context", expanded=False):
+        for i, r in enumerate(retrieved, start=1):
+            st.markdown(f"**Chunk {i}:** {r}")
 
-        # ----------------------------------------------------------
-        # 💡 RESPONSE
-        # ----------------------------------------------------------
-        if user_query.strip():
-            with st.spinner("💭 Thinking..."):
-                retrieved = retrieve_chunks(user_query, index, chunks, top_k=top_k)
-                answer = generate_answer(user_query, retrieved)
-            st.markdown("### Assistant")
-            st.markdown(f"<div class='answer-box'>💡 {answer}</div>", unsafe_allow_html=True)
-
-            st.markdown("""
-            <script>window.scrollTo({top: document.body.scrollHeight, behavior: 'smooth'});</script>
-            """, unsafe_allow_html=True)
-
-            with st.expander("📄 Supporting Context"):
-                for i, r in enumerate(retrieved, start=1):
-                    st.markdown(f"**Chunk {i}:** {r}")
-
-            if toc:
-                with st.expander("📚 Table of Contents"):
-                    toc_text = "\n".join([f"{sec}. {title}" for sec, title in toc])
-                    st.text_area("", toc_text, height=140)
+# ==========================================================
+# Document preview (collapsed)
+# ==========================================================
+if chunks:
+    st.markdown("---")
+    with st.expander("📑 Document Preview", expanded=False):
+        st.text_area("Extracted text (first 1000 chars)", text[:1000], height=180)
+        st.caption(f"📦 {len(chunks)} chunks processed.")
 
-            with st.expander("📄 Document Preview"):
-                st.text_area("", text[:1000], height=140)
-                st.caption(f"{len(chunks)} chunks processed.")
+# End of file