Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,8 +1,6 @@
 import os
 import re
 import shutil
-import pickle
-import faiss
 import streamlit as st
 import torch
 
@@ -24,7 +22,11 @@ st.set_page_config(page_title="Enterprise Knowledge Assistant", layout="wide")
 # 🧹 Cache Management
 # ==========================================================
 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):
@@ -58,74 +60,80 @@ os.environ.update({
 # ==========================================================
 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
+from qa import retrieve_chunks, generate_answer, cache_embeddings, embed_chunks, genai_generate
+
+# ==========================================================
+# 🧠 Smart Suggestion Generator (Strong Version)
+# ==========================================================
+def generate_dynamic_suggestions_from_toc(toc, chunks, doc_name="Document"):
+    """Generate 5–7 short, smart, natural suggestions based on document context."""
+    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 intelligent, concise questions for an enterprise assistant.
+    Document name: "{doc_name}"
+    TABLE OF CONTENTS:
+    {chr(10).join(['- ' + t for t in titles[:10]])}
+
+    CONTENT SAMPLE:
+    {context_sample}
+
+    Generate 5–7 smart, helpful, and professional questions that a user might ask about this document.
+    Keep them short (<20 words), clear, and well-formed.
+    """
+
+    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 start using this process?",
+            "What key configurations are included?",
+        ]
 
 # ==========================================================
-# 💾 Persistent Embedding Cache for Hugging Face
+# 📁 Paths
 # ==========================================================
-EMBED_STORE_DIR = "/data/embedding_cache"
-os.makedirs(EMBED_STORE_DIR, exist_ok=True)
-
-# ==========================================================
-# 🧠 Smart Dynamic Suggestion Generator (doc-based only)
-# ==========================================================
-def generate_dynamic_suggestions_from_doc(toc, chunks, doc_name="Document"):
-    """Generate clean, doc-based questions — no preset or AI hallucination."""
-    suggestions = []
-
-    # Case 1: Use TOC if available
-    if toc:
-        toc_titles = [re.sub(r"^\s*[\dA-Za-z.\-]+\s*", "", title).strip() for _, title in toc]
-        toc_titles = [re.sub(r"\.{2,}\s*\d+$", "", t).strip() for t in toc_titles if len(t.strip()) > 3]
-
-        for title in toc_titles[:8]:
-            title_clean = re.sub(r"[\d\-.]", "", title).strip()
-            if re.search(r"\b(how|what|why|when|where)\b", title_clean.lower()):
-                suggestions.append(title_clean.rstrip("?") + "?")
-            else:
-                suggestions.append(f"What is described in '{title_clean}'?")
-        return suggestions[:6]
-
-    # Case 2: Fallback to section headers or recurring patterns
-    text_sample = " ".join(chunks[:6])
-    headers = re.findall(r"(?<=\n)[A-Z][A-Za-z0-9\s:,-]{5,80}(?=\n)", text_sample)
-    for h in headers[:8]:
-        h = h.strip().rstrip(":")
-        suggestions.append(f"What is explained in '{h}'?")
-
-    # Case 3: If none found
-    if not suggestions:
-        suggestions = [
-            "What is the purpose of this document?",
-            "How can this process be started?",
-            "What are the key configurations explained?",
-        ]
-
-    # Deduplicate and format
-    seen, clean_final = set(), []
-    for q in suggestions:
-        q = q.strip().capitalize()
-        if q not in seen and 10 < len(q) < 120:
-            seen.add(q)
-            clean_final.append(q)
-    return clean_final[:6]
+BASE_DIR = os.path.dirname(__file__)
+LOGO_PATH = os.path.join(BASE_DIR, "logo.png")
+SAMPLE_PATH = os.path.join(BASE_DIR, "sample.pdf")
 
 # ==========================================================
 # 🖥️ Header
 # ==========================================================
 st.title("📄 Enterprise Knowledge Assistant")
-st.caption("Query SAP documentation and enterprise PDFs — powered by reasoning and retrieval.")
+st.caption("Ask questions about SAP documentation and enterprise PDFs — powered by reasoning and retrieval.")
 
 # ==========================================================
 # 🧭 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(
         "🧠 Enable Reasoning Mode",
         value=st.session_state.reasoning_mode,
-        help="ON = reasoning · OFF = strict factual lookup"
+        help="ON = expanded reasoning · OFF = factual answers"
     )
 
     st.markdown("---")
@@ -141,7 +149,7 @@ with st.sidebar:
     st.caption("✨ Built by Shubham Sharma")
 
 # ==========================================================
-# 🧾 Document Handling
+# 🧾 Document Handling + Question UI (Perfect Sync)
 # ==========================================================
 text, chunks, index, embeddings, toc = None, None, None, None, None
 
@@ -154,15 +162,12 @@ for key, default in {
     if key not in st.session_state:
         st.session_state[key] = default
 
-# ----------------------------------------------------------
-# 📄 Main Logic
-# ----------------------------------------------------------
+# --- Document Handling ---
 if doc_choice == "-- Select --":
     st.info("⬅️ Please choose a document from the sidebar to begin.")
 else:
-    # Load document
     if doc_choice == "Sample PDF":
-        temp_path = os.path.join(os.path.dirname(__file__), "sample.pdf")
+        temp_path = SAMPLE_PATH
         st.success("📘 Using built-in Sample PDF")
     else:
         uploaded_file = st.file_uploader("📂 Upload your PDF", type="pdf")
@@ -175,49 +180,31 @@ else:
             temp_path = None
 
     if temp_path:
-        # ----------------------------------------------------------
-        # 🔍 Extract + Chunk + Suggest
-        # ----------------------------------------------------------
         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_doc(toc, chunks, os.path.basename(temp_path))
+            query_suggestions = generate_dynamic_suggestions_from_toc(toc, chunks, os.path.basename(temp_path))
 
-        # ----------------------------------------------------------
-        # 💾 Persistent Embedding Cache (Hugging Face Safe)
-        # ----------------------------------------------------------
         with st.spinner("⚙️ Preparing search index..."):
-            doc_id = os.path.basename(temp_path).replace(".pdf", "")
-            emb_path = os.path.join(EMBED_STORE_DIR, f"{doc_id}_embeddings.pkl")
-            index_path = os.path.join(EMBED_STORE_DIR, f"{doc_id}_faiss.index")
-
-            if os.path.exists(emb_path) and os.path.exists(index_path):
-                with open(emb_path, "rb") as f:
-                    embeddings = pickle.load(f)
-                index = faiss.read_index(index_path)
-                st.success(f"📦 Loaded cached embeddings for '{doc_id}'.")
-            else:
-                embeddings = cache_embeddings(doc_id, chunks, embed_chunks)
-                index = build_faiss_index(embeddings)
-                with open(emb_path, "wb") as f:
-                    pickle.dump(embeddings, f)
-                faiss.write_index(index, index_path)
-                st.success(f"💾 Saved new embeddings for '{doc_id}'.")
+            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.")
 
         # ----------------------------------------------------------
-        # 💬 Ask a Question Section
+        # 💬 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"suggest_{i}"):
+                if col.button(f"🔍 {q}", key=f"q_{i}"):
                     st.session_state.selected_suggestion = i
                     st.session_state.user_query_input = q
 
@@ -226,17 +213,16 @@ else:
                 st.session_state.show_more = not st.session_state.show_more
                 st.experimental_rerun()
 
-        # --- Input box (synced with suggestion) ---
+        # --- Text Input synced with suggestion ---
         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 ---
+        # --- Answer Generation ---
         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)
@@ -248,9 +234,7 @@ else:
                 for i, r in enumerate(retrieved, start=1):
                     st.markdown(f"**Chunk {i}:** {r}")
 
-# ----------------------------------------------------------
-# 📑 Document Preview
-# ----------------------------------------------------------
+# --- Document Preview ---
 if chunks:
     st.markdown("---")
     st.subheader("📑 Document Preview")