Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,247 +1,153 @@
-# streamlit_app.py
+# ==========================================================
+# streamlit_app.py — Restored Stable UI (Pre-rerun Fix)
+# ==========================================================
 import os
 import re
 import streamlit as st
-import shutil
-import hashlib
-import pickle
-
-# Protect against pages that call heavy imports before environment set
-st.set_page_config(page_title="Enterprise Knowledge Assistant", layout="wide")
-
-# --- 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
 
-# ==========================================================
-# SESSION - initialize keys early to avoid widget duplication
-# ==========================================================
-initial_state = {
-    "user_query_input": "",
-    "show_more": False,
-    "selected_suggestion": None,
-    "query_suggestions_fixed": None,
-    "last_doc": None,
-    "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
+st.set_page_config(page_title="Enterprise Knowledge Assistant", layout="wide")
 
 # ==========================================================
-# Sidebar (stable, no unexpected reruns)
+# 🧭 SIDEBAR — Original Clean Version
 # ==========================================================
 with st.sidebar:
     st.markdown("### 🧭 Response Mode")
-    st.session_state["reasoning_mode"] = st.radio(
+    mode = st.radio(
         "",
         ("Strict (Document-only)", "Extended (Document + general)"),
         index=0,
-        help="Strict = answers only from the document. Extended = may include general context."
-    ) == "Extended (Document + general)"
+        help="Strict = answers only from the document. Extended = may include general context.",
+    )
 
     st.markdown("---")
-    # Avoid forcing reruns when toggled: Streamlit manages state
-    st.checkbox("Show advanced settings (for developers)", key="show_advanced")
-
-    if st.session_state.show_advanced:
+    show_advanced = st.checkbox("Show advanced settings (for developers)", value=False)
+    if show_advanced:
         st.markdown("### Developer Settings")
-        # 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"
-        )
+        chunk_size = st.slider("Chunk Size (characters)", 200, 1500, 1000, step=50)
+        overlap = st.slider("Chunk Overlap (characters)", 50, 200, 120, step=10)
+        top_k = st.slider("Top K Results", 1, 10, 5)
+    else:
+        chunk_size, overlap, top_k = 1000, 120, 5
+
     st.markdown("---")
     st.caption("✨ Built by Shubham Sharma")
 
 # ==========================================================
-# Page header
+# 🧠 SESSION STATE — Clean and Minimal
 # ==========================================================
-st.title("📄 Enterprise Knowledge Assistant")
-st.caption("Ask questions about SAP documentation and enterprise PDFs — powered by reasoning and retrieval.")
+for key, val in {
+    "user_query_input": "",
+    "show_more": False,
+    "selected_suggestion": None,
+    "query_suggestions_fixed": None,
+    "last_doc": None,
+}.items():
+    if key not in st.session_state:
+        st.session_state[key] = val
+
+def set_user_query(q, idx):
+    st.session_state["user_query_input"] = q
+    st.session_state["selected_suggestion"] = idx
+    st.experimental_rerun()
 
 # ==========================================================
-# Document selection + upload
+# 📘 APP HEADER
 # ==========================================================
-doc_choice = st.radio(
-    "Select a document:",
-    ("-- Select --", "Sample PDF", "Upload Custom PDF"),
-    index=0,
-    key="doc_choice"
+st.markdown(
+    """
+    <h1 style="text-align:center;">📄 Enterprise Knowledge Assistant</h1>
+    <p style="text-align:center;color:gray;">
+        Query SAP documentation and enterprise PDFs — powered by reasoning and retrieval.
+    </p>
+    """,
+    unsafe_allow_html=True,
 )
 
+# ==========================================================
+# 📂 DOCUMENT UPLOAD / SELECTION
+# ==========================================================
+st.markdown("### Select a document:")
+doc_choice = st.radio("", ("-- Select --", "Sample PDF", "Upload Custom PDF"))
+
 temp_path = None
-if doc_choice == "-- Select --":
-    st.info("⬅️ Please choose a document from the sidebar to begin.")
-elif doc_choice == "Sample PDF":
+if 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")
+    st.success("📘 Sample PDF loaded successfully. Ask questions below.")
+elif doc_choice == "Upload Custom PDF":
+    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.")
-
-# ==========================================================
-# Utility: set query (used by suggestion buttons)
-# ==========================================================
-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()
+        st.success(f"✅ '{uploaded_file.name}' uploaded successfully — ready to query below.")
+else:
+    st.info("⬅️ Please select or upload a document to begin.")
 
 # ==========================================================
-# Document processing and suggestion generation
+# 🧠 PROCESS DOCUMENT (only when selected)
 # ==========================================================
-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"):
+    text, toc = extract_text_from_pdf(temp_path)
+    chunks = chunk_text(text, chunk_size=chunk_size, overlap=overlap)
+    st.success("✅ Document ready — you can now ask questions below.")
+
+    embeddings = cache_embeddings(os.path.basename(temp_path), chunks, embed_chunks)
+    index = build_faiss_index(embeddings)
+
+    # ==========================================================
+    # 💬 SUGGESTED QUESTIONS (Smart Defaults)
+    # ==========================================================
+    if not st.session_state["query_suggestions_fixed"]:
+        st.session_state["query_suggestions_fixed"] = [
+            "What is the purpose of this document?",
+            "How can integration be set up in SAP Cloud?",
+            "What are the prerequisites mentioned?",
+            "What steps are involved in configuration?",
+            "How to troubleshoot integration issues?",
+            "What is the key functionality covered?"
+        ]
+
+    st.markdown("### Ask the Assistant")
+
+    visible_qs = st.session_state["query_suggestions_fixed"][:3] if not st.session_state["show_more"] else st.session_state["query_suggestions_fixed"]
+    cols = st.columns(3)
+    for i, q in enumerate(visible_qs):
+        if cols[i % 3].button(f"🔹 {q}", key=f"suggest_{i}"):
+            set_user_query(q, i)
+
+    if st.button("Show more ▼" if not st.session_state["show_more"] else "Show less ▲"):
         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",
-)
+    # ==========================================================
+    # 🧩 QUERY INPUT
+    # ==========================================================
+    user_query = st.text_input("Type your question or click one above:", value=st.session_state["user_query_input"], key="user_query_input")
 
-# ==========================================================
-# Answer generation (runs only when user_query non-empty)
-# ==========================================================
-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"])
+    # ==========================================================
+    # 🤖 GENERATE ANSWER
+    # ==========================================================
+    if user_query:
+        reasoning_mode = mode == "Extended (Document + general)"
+        st.markdown(f"**Mode:** {'🧠 Extended' if reasoning_mode else '📄 Strict Document'}")
+        with st.spinner("Thinking..."):
+            retrieved = retrieve_chunks(user_query, index, chunks, top_k=top_k, embeddings=embeddings)
+            answer = generate_answer(user_query, retrieved, reasoning_mode=reasoning_mode)
 
-    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)
+        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)
 
-    with st.expander("📄 Supporting Context", expanded=False):
-        for i, r in enumerate(retrieved, start=1):
-            st.markdown(f"**Chunk {i}:** {r}")
+        with st.expander("📘 Supporting Context", expanded=False):
+            for i, chunk in enumerate(retrieved, 1):
+                st.markdown(f"**Chunk {i}:** {chunk.strip()}")
 
 # ==========================================================
-# Document preview (collapsed)
+# 🗂️ DOCUMENT PREVIEW
 # ==========================================================
-if chunks:
+if temp_path:
     st.markdown("---")
-    with st.expander("📑 Document Preview", expanded=False):
-        st.text_area("Extracted text (first 1000 chars)", text[:1000], height=180)
+    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.")
-
-# End of file