Update src/streamlit_app.py

src/streamlit_app.py

@@ -3,28 +3,24 @@ import re
 import shutil
 import streamlit as st
 import torch
-from typing import List, Tuple
 
-# -------------------------
-# Environment / diagnostics
-# -------------------------
+# ==========================================================
+# ✅ Environment Diagnostics
+# ==========================================================
 print("CUDA available:", torch.cuda.is_available())
 if torch.cuda.is_available():
-    try:
-        print("GPU:", torch.cuda.get_device_name(0))
-    except Exception:
-        pass
+    print("GPU:", torch.cuda.get_device_name(0))
 else:
     print("Running on CPU")
 
-# -------------------------
-# Page config
-# -------------------------
+# ==========================================================
+# ✅ Page Configuration
+# ==========================================================
 st.set_page_config(page_title="Enterprise Knowledge Assistant", layout="wide")
 
-# -------------------------
-# Cache cleanup (prevent HF cache overflow on spaces)
-# -------------------------
+# ==========================================================
+# 🧹 Cache Management
+# ==========================================================
 def clean_cache(max_size_gb: float = 2.0):
     folders = [
         "/root/.cache/huggingface",
@@ -33,27 +29,23 @@ def clean_cache(max_size_gb: float = 2.0):
     ]
     total_deleted = 0.0
     for folder in folders:
-        try:
-            if os.path.exists(folder):
-                size_bytes = sum(
-                    os.path.getsize(os.path.join(dp, f))
-                    for dp, _, files in os.walk(folder)
-                    for f in files
-                )
-                size_gb = size_bytes / (1024 ** 3)
-                if size_gb > max_size_gb or "torch" in folder:
-                    shutil.rmtree(folder, ignore_errors=True)
-                    total_deleted += size_gb
-        except Exception:
-            pass
+        if os.path.exists(folder):
+            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 or "torch" in folder:
+                shutil.rmtree(folder, ignore_errors=True)
+                total_deleted += size_gb
     os.makedirs("/tmp/hf_cache", exist_ok=True)
     print(f"🧹 Cache cleanup done. Removed ~{total_deleted:.2f} GB.")
 
 clean_cache()
 
-# -------------------------
-# HF cache env
-# -------------------------
+# ==========================================================
+# ⚙️ Hugging Face Cache Configuration
+# ==========================================================
 CACHE_DIR = "/tmp/hf_cache"
 os.makedirs(CACHE_DIR, exist_ok=True)
 os.environ.update({
@@ -63,106 +55,73 @@ os.environ.update({
     "HF_MODULES_CACHE": CACHE_DIR
 })
 
-# -------------------------
-# App-specific imports
-# -------------------------
+# ==========================================================
+# 📦 Imports AFTER Environment Setup
+# ==========================================================
 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 (fixed f-string)
-# -------------------------
-def generate_dynamic_suggestions_from_toc(toc: List[Tuple[str, str]], chunks: List[str], doc_name="Document") -> List[str]:
-    if not chunks:
+# ==========================================================
+# 🧠 Smart Suggestion Generator (Based on TOC + Content)
+# ==========================================================
+def generate_dynamic_suggestions_from_toc(toc, chunks, doc_name="Document"):
+    if not toc and not chunks:
         return []
 
     titles = []
     for sec, raw_title in toc or []:
         title = re.sub(r"^\s*[\dA-Za-z.\-]+\s*", "", raw_title)
         title = re.sub(r"\.{2,}\s*\d+$", "", title).strip()
-        if 4 < len(title) < 140:
+        if 4 < len(title) < 120:
             titles.append(title)
 
-    if not titles:
-        sample = " ".join(chunks[:4])
-        sents = re.split(r'(?<=[.?!])\s+', sample)
-        suggestions = []
-        for s in sents:
-            if re.search(r'\b(set up|configure|install|enable|procedure|process|how to|step)\b', s, re.I):
-                s = re.sub(r'[.?!]+$', '', s.strip())
-                q = s[0].upper() + s[1:]
-                if len(q) < 140:
-                    suggestions.append(q if q.endswith('?') else q + '?')
-        return suggestions[:7]
+    context_sample = " ".join(chunks[:3])[:3000]
+    prompt = f"""
+    You are helping a user explore a document titled "{doc_name}".
+    TABLE OF CONTENTS:
+    {chr(10).join(['- ' + t for t in titles[:10]])}
 
-    try:
-        prompt = f"""Generate 5-7 user-friendly, short questions based on the document's table of contents:
-Document: {doc_name}
-TOC:
-{chr(10).join(['- ' + t for t in titles[:15]])}"""
-        ai_resp = genai_generate(prompt)
-        qs = re.findall(r'([^\n?.!]+\?)', ai_resp)
-        clean_qs = [q.strip() for q in qs if 8 < len(q) < 140]
-        if clean_qs:
-            return list(dict.fromkeys(clean_qs))[:7]
-    except Exception:
-        pass
+    CONTENT SAMPLE:
+    {context_sample}
 
-    suggestions = []
-    for t in titles[:15]:
-        low = t.lower()
-        if re.search(r'\b(set up|install|configure|enable|define|create|prepare)\b', low):
-            cleaned = re.sub(r'[^a-zA-Z0-9 \-]', '', low)
-            suggestions.append(f"How do I {cleaned}?")
-        elif re.search(r'\b(purpose|overview|objective|scope|what is)\b', low):
-            suggestions.append(f"What is {t.strip().rstrip('.')}?")
-        elif re.search(r'\b(step|procedure|process|task)\b', low):
-            suggestions.append(f"What are the steps for {t.strip().rstrip('.')}?")
-        else:
-            suggestions.append(f"What is described in '{t}'?")
-
-    seen, final = set(), []
-    for s in suggestions:
-        s = re.sub(r'\s+', ' ', s).strip()
-        if s.lower() not in seen:
-            seen.add(s.lower())
-            final.append(s)
-    return final[:7]
+    Generate 5–7 short, specific professional questions.
+    Each should be under 20 words, well-phrased, and relevant to the document.
+    """
 
-# -------------------------
-# CSS for modern UI
-# -------------------------
-st.markdown("""
-<style>
-.section-header {font-weight:700;font-size:1.2rem;margin-top:22px;margin-bottom:8px;color:#f3f4f6;}
-.info-card {background:linear-gradient(180deg,rgba(16,24,39,0.9),rgba(6,10,14,0.9));padding:12px 16px;border-radius:10px;color:#e6eef6;box-shadow:0 6px 20px rgba(2,6,23,0.5);}
-.suggest-chip {display:inline-block;margin:6px 8px;padding:10px 16px;border-radius:999px;background:rgba(31,41,55,0.65);border:1px solid rgba(148,163,184,0.08);color:#e6eef6;font-size:14px;cursor:pointer;transition:transform .12s ease,box-shadow .12s ease;}
-.suggest-chip:hover {transform:translateY(-4px);box-shadow:0 8px 20px rgba(37,99,235,0.12);background:rgba(37,99,235,0.12);}
-.suggest-chip.active {background:linear-gradient(90deg,rgba(59,130,246,0.12),rgba(99,102,241,0.06));border:1px solid rgba(99,102,241,0.22);color:#eaf2ff;box-shadow:0 8px 30px rgba(37,99,235,0.12);}
-.stTextInput>div>div>input {background:rgba(17,24,39,0.75);border-radius:8px;padding:12px 14px;color:#e6eef6;border:1px solid rgba(255,255,255,0.03);}
-.stTextInput>div>div>input:focus {box-shadow:0 6px 20px rgba(37,99,235,0.06);border:1px solid rgba(37,99,235,0.3);}
-.assistant-card {background:linear-gradient(180deg,rgba(6,10,14,0.7),rgba(10,15,20,0.6));border-left:4px solid rgba(59,130,246,0.9);padding:18px 20px;border-radius:10px;color:#f8fbff;margin-bottom:10px;box-shadow:0 10px 40px rgba(2,6,23,0.5);}
-</style>
-""", unsafe_allow_html=True)
+    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 configure this process?",
+            "What key steps are described?",
+        ]
 
-# -------------------------
-# App Header
-# -------------------------
+# ==========================================================
+# 🖥️ Header
+# ==========================================================
 st.title("📄 Enterprise Knowledge Assistant")
 st.caption("Query SAP documentation and enterprise PDFs — powered by reasoning and retrieval.")
 
-# -------------------------
-# Sidebar
-# -------------------------
+# ==========================================================
+# 🧭 Sidebar
+# ==========================================================
 with st.sidebar:
     if "reasoning_mode" not in st.session_state:
         st.session_state.reasoning_mode = False
-
     st.session_state.reasoning_mode = st.toggle(
         "🧠 Enable Reasoning Mode",
         value=st.session_state.reasoning_mode,
-        help="ON: expanded reasoning / OFF: strict doc lookup"
+        help="ON = detailed reasoning · OFF = concise factual answers"
     )
 
     st.markdown("---")
@@ -171,32 +130,34 @@ with st.sidebar:
 
     st.markdown("---")
     st.header("⚙️ Settings")
-    chunk_size = st.slider("Chunk Size (chars)", 200, 1500, 1000, step=50)
-    overlap = st.slider("Chunk Overlap (chars)", 50, 200, 120, step=10)
+    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)
     st.markdown("---")
     st.caption("✨ Built by Shubham Sharma")
 
-# -------------------------
-# Session state
-# -------------------------
-for k, v in {
+# ==========================================================
+# 🧾 Document Handling
+# ==========================================================
+text, chunks, index, embeddings, toc = None, None, None, None, None
+
+# --- Initialize Session State ---
+for key, default in {
     "show_more": False,
     "user_query_input": "",
-    "selected_suggestion_idx": None,
+    "selected_suggestion": None,
 }.items():
-    if k not in st.session_state:
-        st.session_state[k] = v
+    if key not in st.session_state:
+        st.session_state[key] = default
 
-# -------------------------
-# Document processing
-# -------------------------
+# ==========================================================
+# 📄 Main Flow
+# ==========================================================
 BASE_DIR = os.path.dirname(__file__)
 SAMPLE_PATH = os.path.join(BASE_DIR, "sample.pdf")
-text, chunks, index, embeddings, toc = None, None, None, None, None
 
 if doc_choice == "-- Select --":
-    st.info("⬅️ Please choose a document from the sidebar to begin.")
+    st.info("⬅️ Please choose a document to begin.")
 else:
     if doc_choice == "Sample PDF":
         temp_path = SAMPLE_PATH
@@ -211,53 +172,87 @@ else:
         else:
             temp_path = None
 
+    # ------------------------------
+    # Document Processing
+    # ------------------------------
     if temp_path:
-        with st.spinner("🔍 Processing document..."):
+        with st.spinner("🔍 Processing your document..."):
             text, toc = extract_text_from_pdf(temp_path)
             chunks = chunk_text(text, chunk_size=chunk_size)
-            st.markdown("<div class='info-card'>✅ Document loaded successfully.</div>", unsafe_allow_html=True)
+            st.markdown("✅ Document loaded successfully.")
 
-            query_suggestions = generate_dynamic_suggestions_from_toc(toc, chunks, os.path.basename(temp_path))
+            query_suggestions = generate_dynamic_suggestions_from_toc(
+                toc, chunks, os.path.basename(temp_path)
+            )
 
         with st.spinner("⚙️ Preparing embeddings..."):
             embeddings = cache_embeddings(os.path.basename(temp_path), chunks, embed_chunks)
             index = build_faiss_index(embeddings)
-            st.markdown("<div class='info-card'>🚀 Document ready — ask questions below.</div>", unsafe_allow_html=True)
+            st.markdown("🚀 Document ready — you can now ask questions below.")
 
-        # -------------------------
-        # Ask a Question
-        # -------------------------
+        # ==========================================================
+        # 💬 Ask a Question
+        # ==========================================================
         st.markdown("## 🤖 Ask a Question")
-
         if query_suggestions:
             st.markdown("#### 💡 Suggested Questions")
+
             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):
                 col = cols[i % 3]
                 if col.button(f"🔍 {q}", key=f"q_{i}"):
+                    st.session_state.selected_suggestion = i
                     st.session_state.user_query_input = q
-                    st.session_state.selected_suggestion_idx = i
+
             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()
 
-        user_query = st.text_input("Type your question or pick one above:", value=st.session_state.user_query_input)
+        # ----------------------------------------------------------
+        # 🧠 Input Field
+        # ----------------------------------------------------------
+        user_query = st.text_input(
+            "Type your question or pick one above:",
+            value=st.session_state.user_query_input,
+            key="user_query_input"
+        )
 
+        # ----------------------------------------------------------
+        # 🧩 Generate Answer
+        # ----------------------------------------------------------
         if user_query.strip():
-            st.caption("Mode: 🧠 Reasoning" if st.session_state.reasoning_mode else "Mode: 📄 Strict Document")
-            with st.spinner("💭 Generating answer..."):
+            st.caption(
+                "Mode: 🧠 Reasoning"
+                if st.session_state.reasoning_mode
+                else "Mode: 📄 Strict Document"
+            )
+            with st.spinner("💭 Analyzing your document..."):
                 retrieved = retrieve_chunks(user_query, index, chunks, top_k=top_k, embeddings=embeddings)
-                answer = generate_answer(user_query, retrieved, reasoning_mode=st.session_state.reasoning_mode)
+                answer = generate_answer(
+                    user_query, retrieved, reasoning_mode=st.session_state.reasoning_mode
+                )
 
+            # Assistant Answer
             st.markdown("### ✅ Assistant’s Answer")
-            st.markdown(f"<div class='assistant-card'>{answer}</div>", unsafe_allow_html=True)
+            st.markdown(
+                f"<div style='background-color:#111827;padding:14px;border-radius:8px;color:#f1f5f9;'>"
+                f"💡 {answer}</div>",
+                unsafe_allow_html=True,
+            )
 
             with st.expander("📄 Supporting Context"):
                 for i, r in enumerate(retrieved, start=1):
-                    st.markdown(f"**Chunk {i}:** {r}")
+                    st.markdown(
+                        f"<div style='background-color:#1f2937;padding:10px;border-radius:6px;color:#e5e7eb;margin-bottom:8px;'>"
+                        f"<b>Chunk {i}:</b> {r}</div>",
+                        unsafe_allow_html=True,
+                    )
 
+        # ----------------------------------------------------------
+        # 📑 Document Preview
+        # ----------------------------------------------------------
         st.markdown("---")
         st.subheader("📑 Document Preview")
         st.text_area("Extracted text (first 1000 chars)", text[:1000], height=200)