Update src/streamlit_app.py

src/streamlit_app.py

@@ -5,46 +5,13 @@ import streamlit as st
 import torch
 
 # ==========================================================
-# ✅ Environment Diagnostics
-# ==========================================================
-print("CUDA available:", torch.cuda.is_available())
-if torch.cuda.is_available():
-    print("GPU:", torch.cuda.get_device_name(0))
-else:
-    print("Running on CPU")
-
-# ==========================================================
-# ✅ Page Configuration
+# ✅ Environment Setup
 # ==========================================================
 st.set_page_config(page_title="Enterprise Knowledge Assistant", layout="wide")
+print("CUDA available:", torch.cuda.is_available())
 
 # ==========================================================
-# 🧹 Cache Management
-# ==========================================================
-def clean_cache(max_size_gb: float = 2.0):
-    folders = [
-        "/root/.cache/huggingface",
-        "/root/.cache/transformers",
-        "/root/.cache/torch",
-    ]
-    total_deleted = 0.0
-    for folder in folders:
-        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()
-
-# ==========================================================
-# ⚙️ Hugging Face Cache Configuration
+# ⚙️ Hugging Face Cache Setup
 # ==========================================================
 CACHE_DIR = "/tmp/hf_cache"
 os.makedirs(CACHE_DIR, exist_ok=True)
@@ -56,37 +23,40 @@ os.environ.update({
 })
 
 # ==========================================================
-# 📦 Imports AFTER Environment Setup
+# 📦 Imports
 # ==========================================================
 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 (Based on TOC + Content)
+# 🧠 Smart Suggestion Generator
 # ==========================================================
 def generate_dynamic_suggestions_from_toc(toc, chunks, doc_name="Document"):
-    if not toc and not chunks:
+    """Generate clean, context-aware questions dynamically from TOC and text."""
+    if not toc or not chunks:
         return []
 
     titles = []
-    for sec, raw_title in toc or []:
+    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])[:3000]
+    context_sample = " ".join(chunks[:3])[:4000]
     prompt = f"""
-    You are helping a user explore a document titled "{doc_name}".
+    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[:10]])}
+    {chr(10).join(['- ' + t for t in titles[:8]])}
 
-    CONTENT SAMPLE:
+    TEXT SAMPLE:
     {context_sample}
 
-    Generate 5–7 short, specific professional questions.
-    Each should be under 20 words, well-phrased, and relevant to the document.
+    Generate 5–7 questions that are short, relevant, and strictly document-based.
+    Each question should be under 18 words.
     """
 
     try:
@@ -100,17 +70,66 @@ def generate_dynamic_suggestions_from_toc(toc, chunks, doc_name="Document"):
                 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?",
-        ]
-
-# ==========================================================
-# 🖥️ Header
-# ==========================================================
-st.title("📄 Enterprise Knowledge Assistant")
-st.caption("Query SAP documentation and enterprise PDFs — powered by reasoning and retrieval.")
+        return ["What is this document about?", "How do I start using this process?"]
+
+# ==========================================================
+# 🎨 Styling — Customer-Ready Minimal Theme
+# ==========================================================
+st.markdown("""
+<style>
+div.block-container {
+    padding-top: 1.5rem;
+    max-width: 1050px;
+}
+h1, h2, h3, h4 {
+    font-weight: 600;
+    color: #f3f4f6;
+}
+hr {
+    border: none;
+    border-top: 1px solid #2c2c2c;
+    margin: 1rem 0;
+}
+.suggest-chip {
+    background-color: #1f2937;
+    border: 1px solid #374151;
+    border-radius: 16px;
+    color: #e5e7eb;
+    padding: 6px 12px;
+    cursor: pointer;
+    font-size: 13px;
+    transition: all 0.2s ease-in-out;
+}
+.suggest-chip:hover {
+    background-color: #2563eb;
+    border-color: #3b82f6;
+    color: white;
+    box-shadow: 0 0 8px rgba(59,130,246,0.4);
+}
+.answer-box {
+    background: linear-gradient(135deg, #0f172a, #1e293b);
+    border-left: 4px solid #3b82f6;
+    border-radius: 8px;
+    padding: 14px 16px;
+    color: #f1f5f9;
+    margin-top: 1rem;
+    box-shadow: 0 0 10px rgba(59,130,246,0.1);
+}
+.stTextInput > div > div > input {
+    background-color: #0f172a;
+    color: #f1f5f9;
+    border-radius: 6px;
+    border: 1px solid #334155;
+    padding: 6px 10px;
+}
+.stTextArea > div > div > textarea {
+    background-color: #0f172a;
+    color: #f1f5f9;
+    border-radius: 6px;
+    border: 1px solid #334155;
+}
+</style>
+""", unsafe_allow_html=True)
 
 # ==========================================================
 # 🧭 Sidebar
@@ -118,49 +137,40 @@ st.caption("Query SAP documentation and enterprise PDFs — powered by reasoning
 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 = detailed reasoning · OFF = concise factual answers"
-    )
 
-    st.markdown("---")
-    st.header("📚 Document Library")
-    doc_choice = st.radio("Choose a document:", ["-- Select --", "Sample PDF", "Upload Custom PDF"], index=1)
+    st.markdown("### ⚙️ Settings")
+    reasoning_mode = st.toggle("🧠 Enable Reasoning Mode", st.session_state.reasoning_mode)
+    st.session_state.reasoning_mode = reasoning_mode
 
-    st.markdown("---")
-    st.header("⚙️ Settings")
-    chunk_size = st.slider("Chunk Size (characters)", 200, 1500, 1000, step=50)
-    overlap = st.slider("Chunk Overlap (characters)", 50, 200, 120, step=10)
+    chunk_size = st.slider("Chunk Size", 200, 1500, 1000, step=50)
+    overlap = st.slider("Chunk Overlap", 50, 200, 120, step=10)
     top_k = st.slider("Top K Results", 1, 10, 5)
     st.markdown("---")
     st.caption("✨ Built by Shubham Sharma")
 
 # ==========================================================
-# 🧾 Document Handling
+# 📄 Main Flow
 # ==========================================================
+st.title("Enterprise Knowledge Assistant")
+st.caption("Query SAP documentation and enterprise PDFs — powered by reasoning and retrieval.")
+
 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": None,
-}.items():
-    if key not in st.session_state:
-        st.session_state[key] = default
+if "user_query_input" not in st.session_state:
+    st.session_state["user_query_input"] = ""
+if "show_more" not in st.session_state:
+    st.session_state["show_more"] = False
 
-# ==========================================================
-# 📄 Main Flow
-# ==========================================================
-BASE_DIR = os.path.dirname(__file__)
-SAMPLE_PATH = os.path.join(BASE_DIR, "sample.pdf")
+def set_user_query(q):
+    st.session_state["user_query_input"] = q
+
+doc_choice = st.radio("", ["-- Select --", "Sample PDF", "Upload Custom PDF"], index=1)
 
 if doc_choice == "-- Select --":
-    st.info("⬅️ Please choose a document to begin.")
+    st.info("⬅️ Select a document to begin.")
 else:
     if doc_choice == "Sample PDF":
-        temp_path = SAMPLE_PATH
+        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")
@@ -172,88 +182,60 @@ else:
         else:
             temp_path = None
 
-    # ------------------------------
-    # Document Processing
-    # ------------------------------
     if temp_path:
         with st.spinner("🔍 Processing your document..."):
             text, toc = extract_text_from_pdf(temp_path)
             chunks = chunk_text(text, chunk_size=chunk_size)
-            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..."):
+        with st.spinner("⚙️ Preparing search index..."):
             embeddings = cache_embeddings(os.path.basename(temp_path), chunks, embed_chunks)
             index = build_faiss_index(embeddings)
-            st.markdown("🚀 Document ready — you can now ask questions below.")
+            st.success("🚀 Document ready.")
 
-        # ==========================================================
+        # ----------------------------------------------------------
         # 💬 Ask a Question
-        # ==========================================================
-        st.markdown("## 🤖 Ask a Question")
-        if query_suggestions:
-            st.markdown("#### 💡 Suggested Questions")
+        # ----------------------------------------------------------
+        st.markdown("### Ask a Question")
 
+        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):
-                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
+                cols[i % 3].button(f"🔍 {q}", key=f"suggest_{i}", on_click=set_user_query, args=(q,))
 
             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()
 
-        # ----------------------------------------------------------
-        # 🧠 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"
-        )
+        user_query = st.text_input("Type your question or pick one above:", 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("💭 Analyzing your document..."):
+            with st.spinner("💭 Analyzing 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
-                )
-
-            # Assistant Answer
-            st.markdown("### ✅ Assistant’s Answer")
-            st.markdown(
-                f"<div style='background-color:#111827;padding:14px;border-radius:8px;color:#f1f5f9;'>"
-                f"💡 {answer}</div>",
-                unsafe_allow_html=True,
-            )
+                answer = generate_answer(user_query, retrieved, reasoning_mode=st.session_state.reasoning_mode)
+
+            st.markdown("### Assistant’s Answer")
+            st.markdown(f"<div class='answer-box'>💡 {answer}</div>", unsafe_allow_html=True)
 
             with st.expander("📄 Supporting Context"):
                 for i, r in enumerate(retrieved, start=1):
-                    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,
-                    )
+                    st.markdown(f"**Chunk {i}:** {r}")
+
+        # ----------------------------------------------------------
+        # 📚 Table of Contents
+        # ----------------------------------------------------------
+        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=150)
 
         # ----------------------------------------------------------
-        # 📑 Document Preview
+        # 📄 Document Preview
         # ----------------------------------------------------------
-        st.markdown("---")
-        st.subheader("📑 Document Preview")
-        st.text_area("Extracted text (first 1000 chars)", text[:1000], height=200)
-        st.caption(f"📦 {len(chunks)} chunks processed.")
+        if chunks:
+            with st.expander("📄 Document Preview"):
+                st.text_area("", text[:1000], height=150)
+                st.caption(f"{len(chunks)} chunks processed.")