Update src/streamlit_app.py

src/streamlit_app.py

@@ -5,7 +5,7 @@ import streamlit as st
 import torch
 
 # ==========================================================
-# ✅ Environment Diagnostics
+# ✅ Environment Setup
 # ==========================================================
 print("CUDA available:", torch.cuda.is_available())
 if torch.cuda.is_available():
@@ -13,20 +13,13 @@ if torch.cuda.is_available():
 else:
     print("Running on CPU")
 
-# ==========================================================
-# ✅ Page Configuration
-# ==========================================================
 st.set_page_config(page_title="Enterprise Knowledge Assistant", layout="wide")
 
 # ==========================================================
-# 🧹 Cache Management
+# 🧹 Cache Cleanup
 # ==========================================================
 def clean_cache(max_size_gb: float = 2.0):
-    folders = [
-        "/root/.cache/huggingface",
-        "/root/.cache/transformers",
-        "/root/.cache/torch",
-    ]
+    folders = ["/root/.cache/huggingface", "/root/.cache/transformers", "/root/.cache/torch"]
     total_deleted = 0.0
     for folder in folders:
         if os.path.exists(folder):
@@ -44,7 +37,7 @@ def clean_cache(max_size_gb: float = 2.0):
 clean_cache()
 
 # ==========================================================
-# ⚙️ Hugging Face Cache Configuration
+# ⚙️ Hugging Face Cache
 # ==========================================================
 CACHE_DIR = "/tmp/hf_cache"
 os.makedirs(CACHE_DIR, exist_ok=True)
@@ -66,7 +59,6 @@ from qa import retrieve_chunks, generate_answer, cache_embeddings, embed_chunks,
 # 🧠 Smart Suggestion Generator
 # ==========================================================
 def generate_dynamic_suggestions_from_toc(toc, chunks, doc_name="Document"):
-    """Generate natural-sounding, concise professional questions from TOC + context."""
     if not toc or not chunks:
         return []
 
@@ -79,15 +71,17 @@ def generate_dynamic_suggestions_from_toc(toc, chunks, doc_name="Document"):
 
     context_sample = " ".join(chunks[:3])[:4000]
     prompt = f"""
-    You are an enterprise documentation assistant.
-    Generate 5–7 professional, relevant questions a user might ask about this document: "{doc_name}".
-    Use the table of contents and content sample below.
+    You are generating intelligent, short, and context-aware questions from a document titled "{doc_name}".
+    Use the following Table of Contents and sample content for reference.
 
     TABLE OF CONTENTS:
-    {chr(10).join(['- ' + t for t in titles[:10]])}
+    {chr(10).join(['- ' + t for t in titles[:8]])}
 
     CONTENT SAMPLE:
     {context_sample}
+
+    Generate 5–7 concise, professional, and relevant questions a user might ask about this document.
+    Each question should be under 20 words and directly based on the context.
     """
 
     try:
@@ -101,18 +95,7 @@ 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 start using this process?",
-            "What configurations are needed?",
-        ]
-
-# ==========================================================
-# 📁 Paths
-# ==========================================================
-BASE_DIR = os.path.dirname(__file__)
-LOGO_PATH = os.path.join(BASE_DIR, "logo.png")
-SAMPLE_PATH = os.path.join(BASE_DIR, "sample.pdf")
+        return ["What is this document about?", "How do I start using this process?"]
 
 # ==========================================================
 # 🖥️ Header
@@ -124,9 +107,6 @@ st.caption("Ask questions about SAP documentation and enterprise PDFs — powere
 # 🧭 Sidebar
 # ==========================================================
 with st.sidebar:
-    if os.path.exists(LOGO_PATH):
-        st.image(LOGO_PATH, width=150)
-
     if "reasoning_mode" not in st.session_state:
         st.session_state.reasoning_mode = False
     st.session_state.reasoning_mode = st.toggle(
@@ -148,28 +128,46 @@ with st.sidebar:
     st.caption("✨ Built by Shubham Sharma")
 
 # ==========================================================
-# 🧾 Document Handling
+# 🧾 Document Handling + Ask Section
 # ==========================================================
 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,
-    "pending_query": None,
-}.items():
-    if key not in st.session_state:
-        st.session_state[key] = default
-
-# ----------------------------------------------------------
-# 📄 Document Handling + Ask a Question
-# ----------------------------------------------------------
+# Style for chips
+st.markdown("""
+<style>
+.suggest-chip {
+    background-color: #1f2937;
+    border: 1px solid #374151;
+    border-radius: 20px;
+    color: #f9fafb;
+    padding: 8px 14px;
+    cursor: pointer;
+    font-size: 14px;
+    transition: all 0.2s ease-in-out;
+}
+.suggest-chip:hover {
+    background-color: #2563eb;
+    border-color: #3b82f6;
+    color: #ffffff;
+    box-shadow: 0 0 10px rgba(59,130,246,0.5);
+}
+</style>
+""", unsafe_allow_html=True)
+
+# Initialize session state
+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
+
+def set_user_query(q: str):
+    st.session_state["user_query_input"] = q
+
 if doc_choice == "-- Select --":
     st.info("⬅️ Please choose a document from the sidebar 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")
@@ -193,10 +191,7 @@ else:
                 st.text_area("TOC Preview", toc_text, height=180)
                 query_suggestions = generate_dynamic_suggestions_from_toc(toc, chunks, os.path.basename(temp_path))
             else:
-                query_suggestions = [
-                    "What is this document about?",
-                    "How do I start using this process?",
-                ]
+                query_suggestions = ["What is this document about?", "How do I start using this process?"]
 
         with st.spinner("⚙️ Preparing search index..."):
             embeddings = cache_embeddings(os.path.basename(temp_path), chunks, embed_chunks)
@@ -211,34 +206,21 @@ else:
         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)))
+            visible_suggestions = query_suggestions if st.session_state.show_more else query_suggestions[:3]
+            cols = st.columns(min(3, len(visible_suggestions)))
 
-            for i, q in enumerate(visible):
-                col = cols[i % 3]
-                if col.button(f"🔍 {q}", key=f"suggest_{i}"):
-                    st.session_state.pending_query = q
-                    st.experimental_rerun()
+            for i, q in enumerate(visible_suggestions):
+                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, key="toggle_show_more"):
+            if st.button(toggle_text):
                 st.session_state.show_more = not st.session_state.show_more
                 st.experimental_rerun()
 
-        # --- Handle query sync ---
-        if st.session_state.pending_query:
-            st.session_state.user_query_input = st.session_state.pending_query
-            st.session_state.pending_query = None
-
-        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")
 
         if user_query.strip():
             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)
@@ -250,9 +232,6 @@ else:
                 for i, r in enumerate(retrieved, start=1):
                     st.markdown(f"**Chunk {i}:** {r}")
 
-# ----------------------------------------------------------
-# 📑 Document Preview
-# ----------------------------------------------------------
 if chunks:
     st.markdown("---")
     st.subheader("📑 Document Preview")