Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,65 +1,214 @@
 # ==========================================================
-# 📂 DOCUMENT HANDLING — CLEAN, ACCURATE, AND BYTE-AWARE
+# streamlit_app.py — Stable Layout (English Only)
 # ==========================================================
+import os
+import re
 import hashlib
+import streamlit as st
+import torch
+from document_registry import DocumentRegistry
+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
 
+# ==========================================================
+# ✅ PAGE CONFIG
+# ==========================================================
+st.set_page_config(page_title="Enterprise Knowledge Assistant", layout="wide")
+print("CUDA available:", torch.cuda.is_available())
+
+# ==========================================================
+# ⚙️ CACHE SETUP
+# ==========================================================
+CACHE_DIR = "/tmp/hf_cache"
+os.makedirs(CACHE_DIR, exist_ok=True)
+os.environ.update({
+    "HF_HOME": CACHE_DIR,
+    "TRANSFORMERS_CACHE": CACHE_DIR,
+    "HF_DATASETS_CACHE": CACHE_DIR,
+    "HF_MODULES_CACHE": CACHE_DIR,
+})
+
+# ==========================================================
+# 🧠 SMART SUGGESTION GENERATOR
+# ==========================================================
+def generate_dynamic_suggestions_from_toc(toc, chunks, doc_name="Document"):
+    if not toc or not chunks:
+        return ["How do I start using this guide?", "What does this document cover?"]
+
+    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 a content assistant. Based on the TOC and text, generate 5–7 short natural questions.
+Each question <18 words, ends with '?', sounds human. Document: "{doc_name}"
+TOC:
+{chr(10).join(['- ' + t for t in titles[:8]])}
+Sample:
+{context_sample}
+"""
+
+    try:
+        ai_response = genai_generate(prompt)
+        lines = [ln.strip() for ln in ai_response.splitlines() if ln.strip()]
+        out = []
+        for ln in lines:
+            q = re.sub(r"^[\-\u2022\*\d\.\)\s]+", "", ln).strip()
+            if not q.endswith("?") and re.match(r"(?i)^(what|how|why|where|who|when|which|can|does|is|are)\b", q):
+                q += "?"
+            if 8 <= len(q) <= 140:
+                out.append(q)
+        uniq = []
+        seen = set()
+        for q in out:
+            if q.lower() not in seen:
+                seen.add(q.lower())
+                uniq.append(q)
+        return uniq[:7] or [f"What should I know about {t.rstrip('.')}?" for t in titles[:7]]
+    except Exception:
+        return ["How do I start using this guide?", "What does this document cover?"]
+
+# ==========================================================
+# 🎨 STYLING
+# ==========================================================
+st.markdown("""
+<style>
+div.block-container {padding-top:1.2rem;max-width:1080px;}
+h1,h2,h3{color:#f3f4f6;font-weight:600;}
+.suggest-chip{background:#0f1724;border:1px solid #374151;border-radius:14px;
+color:#e6eef8;padding:8px 12px;cursor:pointer;font-size:13px;margin:6px 6px 6px 0;
+display:inline-block;transition:background 0.2s,transform 0.1s;}
+.suggest-chip:hover{background:#1e3a8a;transform:translateY(-2px);}
+.answer-box{background:linear-gradient(180deg,#0b1220,#071027);
+border-left:4px solid #3b82f6;border-radius:8px;padding:16px 18px;color:#e6eef8;
+margin-top:12px;box-shadow:0 4px 14px rgba(0,0,0,0.35);}
+.stTextInput>div>div>input{background-color:#0f172a!important;color:#f1f5f9!important;
+border-radius:6px!important;border:1px solid #334155!important;padding:8px 10px!important;
+font-size:15px!important;}
+.small-link{font-size:13px;color:#60a5fa;cursor:pointer;}
+</style>
+""", unsafe_allow_html=True)
+
+# ==========================================================
+# 🧭 SIDEBAR
+# ==========================================================
+with st.sidebar:
+    st.markdown("### 🧭 Response Style")
+    mode = st.radio(
+        "",
+        ("Strict (Document-only)", "Extended (Document + General)"),
+        index=0,
+        help="Strict = answers only from the uploaded document."
+    )
+    st.markdown("---")
+
+    # 📚 Registry display
+    if "registry" in st.session_state:
+        registry = st.session_state["registry"]
+        registered = registry.list_docs() if hasattr(registry, "list_docs") else []
+        if registered:
+            with st.expander("📚 Registered Documents", expanded=False):
+                for i, d in enumerate(registered, 1):
+                    st.markdown(f"**{i}. {d.get('name','?')}** — {d.get('num_chunks','?')} chunks *(TOC: {d.get('toc_source','—')})*")
+        else:
+            st.caption("📭 No documents registered yet.")
+    else:
+        st.caption("📭 No registry initialized yet.")
+
+    st.markdown("---")
+    show_dev = st.checkbox("Show advanced settings (for developers)", value=False)
+    if show_dev:
+        st.markdown("### ⚙️ Developer Options")
+        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)
+    else:
+        chunk_size, overlap, top_k = 1000, 120, 5
+    st.markdown("---")
+    st.caption("✨ Built by Shubham Sharma")
+
+# ==========================================================
+# 🧠 SESSION STATE
+# ==========================================================
+for key, val in {
+    "user_query_input": "",
+    "show_more": False,
+    "selected_suggestion": None,
+    "query_suggestions_fixed": None,
+    "last_doc": None,
+    "doc_lang": "en",
+}.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()
+
+# ==========================================================
+# 📄 MAIN SECTION
+# ==========================================================
+st.title("📄 Enterprise Knowledge Assistant")
+st.caption("Query SAP documentation and enterprise PDFs — powered by reasoning and retrieval.")
+
+# ✅ FIXED: must be defined before document-handling logic
+doc_choice = st.radio(
+    "Select a document:",
+    ["-- Select --", "Sample PDF", "Upload Custom PDF"],
+    index=0
+)
+
+# ==========================================================
+# 📂 DOCUMENT HANDLING
+# ==========================================================
 def _hash_content(file_path):
-    """Generate a short SHA256 hash of the file's actual binary content."""
-    hasher = hashlib.sha256()
+    h = hashlib.sha256()
     with open(file_path, "rb") as f:
         while chunk := f.read(8192):
-            hasher.update(chunk)
-    return hasher.hexdigest()[:12]  # short unique hash for same-name files
-
+            h.update(chunk)
+    return h.hexdigest()[:12]
 
 def refresh_suggestions(doc_name, toc, chunks):
-    """Refresh dynamic suggestions and reset related states."""
     st.session_state["query_suggestions_fixed"] = generate_dynamic_suggestions_from_toc(
         toc, chunks, doc_name
     )
-    st.session_state["user_query_input"] = ""
-    st.session_state["selected_suggestion"] = None
-    st.session_state["show_more"] = False
-
+    st.session_state.update({"user_query_input": "", "selected_suggestion": None, "show_more": False})
 
-# --- Document selection ---
 if doc_choice == "-- Select --":
     st.info("⬅️ Select or upload a document to begin.")
 else:
     temp_path = None
-
-    # --- File selection ---
     if doc_choice == "Sample PDF":
         temp_path = os.path.join(os.path.dirname(__file__), "sample.pdf")
         st.markdown("✅ **Sample PDF selected.** Preparing document...")
     else:
-        uploaded_file = st.file_uploader(
-            "Upload a PDF document (max 200MB):",
-            type="pdf",
-            label_visibility="collapsed"
-        )
-        if uploaded_file:
-            temp_path = os.path.join("/tmp", uploaded_file.name)
+        uploaded = st.file_uploader("Upload a PDF document (max 200 MB):", type="pdf", label_visibility="collapsed")
+        if uploaded:
+            temp_path = os.path.join("/tmp", uploaded.name)
             with open(temp_path, "wb") as f:
-                f.write(uploaded_file.getbuffer())
+                f.write(uploaded.getbuffer())
         else:
             st.stop()
 
-    # --- Start processing if file exists ---
     if temp_path:
         doc_name = os.path.basename(temp_path)
-        file_hash = _hash_content(temp_path)
-        doc_identifier = f"{doc_name}_{file_hash}"
+        doc_hash = _hash_content(temp_path)
+        doc_identifier = f"{doc_name}_{doc_hash}"
 
-        # ✅ Step 0: Ensure registry exists
+        # --- registry check
         if "registry" not in st.session_state:
             st.session_state["registry"] = DocumentRegistry()
         registry = st.session_state["registry"]
 
-        # ✅ Step 1: Check if already registered
-        existing_doc = next((d for d in registry.list_docs() if d["name"] == doc_name), None)
-        if existing_doc:
-            doc_data = registry.get_doc(existing_doc["name"])
+        existing = next((d for d in registry.list_docs() if d["name"] == doc_name), None)
+        if existing:
+            doc_data = registry.get_doc(existing["name"])
             st.session_state.update({
                 "text": doc_data.get("text", ""),
                 "toc": doc_data.get("toc", []),
@@ -67,14 +216,12 @@ else:
                 "embeddings": doc_data.get("embeddings"),
                 "index": doc_data.get("index"),
                 "doc_ready": True,
-                "active_doc": existing_doc["name"],
+                "active_doc": existing["name"],
                 "status_text": f"✅ {doc_name} already processed — loaded from registry."
             })
-
-            refresh_suggestions(existing_doc["name"], st.session_state["toc"], st.session_state["chunks"])
+            refresh_suggestions(existing["name"], st.session_state["toc"], st.session_state["chunks"])
             st.experimental_rerun()
 
-        # ✅ Step 2: New document → process
         status = st.empty()
         status.info("📤 Upload complete — reading document...")
 
@@ -88,9 +235,7 @@ else:
 
         doc_id = registry.register(temp_path, chunks, embeddings, index)
         st.session_state["active_doc"] = doc_id
-
-        status.success("✅ Document processed successfully — all set to query your assistant!")
-
+        status.success("✅ Document processed successfully — ready to query!")
         refresh_suggestions(doc_name, toc, chunks)
 
         st.session_state.update({
@@ -101,50 +246,48 @@ else:
             "index": index,
             "doc_ready": True,
             "last_doc": doc_identifier,
-            "status_text": "✅ Document processed successfully — all set to query your assistant!"
+            "status_text": "✅ Document processed successfully — ready to query!"
         })
-
         st.experimental_rerun()
 
-    # --- Display Ready Message + Ask Section ---
-    if st.session_state.get("doc_ready"):
-        active_name = st.session_state.get("active_doc") or st.session_state.get("last_doc")
-        st.info(st.session_state.get("status_text", f"📄 {active_name or 'Document'} is ready for queries."))
-
-        st.markdown("### 💬 Ask the Assistant")
-        query_suggestions = st.session_state.get("query_suggestions_fixed", [])
-        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):
-                if cols[i % 3].button(f"💬 {q}", key=f"sugg_{i}"):
-                    st.session_state["user_query_input"] = q
-                    st.session_state["selected_suggestion"] = i
-                    st.experimental_rerun()
-
-            toggle_text = "Show less ▲" if st.session_state["show_more"] else "Show more ▼"
-            if st.button(toggle_text, help="Show or hide more suggestions"):
-                st.session_state["show_more"] = not st.session_state["show_more"]
-                st.experimental_rerun()
-
-        user_query = st.text_input(
-            "Type your question or click one above:",
-            key="user_query_input",
-            label_visibility="visible"
-        )
-
-        if user_query.strip():
-            reasoning_mode = mode == "Extended (Document + General)"
-            with st.spinner("💭 Generating your answer..."):
-                retrieved = retrieve_chunks(
-                    user_query,
-                    st.session_state["index"],
-                    st.session_state["chunks"],
-                    top_k=top_k,
-                    embeddings=st.session_state["embeddings"]
-                )
-                answer = generate_answer(user_query, retrieved, reasoning_mode=reasoning_mode)
-                st.session_state["retrieved"] = retrieved
-
-            st.markdown("### 🤖 Assistant’s Answer")
-            st.markdown(f"<div class='answer-box'>{answer}</div>", unsafe_allow_html=True)
+if st.session_state.get("doc_ready"):
+    st.info(st.session_state.get("status_text", "📄 Ready for queries."))
+    st.markdown("### 💬 Ask the Assistant")
+
+    suggs = st.session_state.get("query_suggestions_fixed", [])
+    if suggs:
+        visible = suggs if st.session_state["show_more"] else suggs[:3]
+        cols = st.columns(min(3, len(visible)))
+        for i, q in enumerate(visible):
+            if cols[i % 3].button(f"💬 {q}", key=f"sugg_{i}"):
+                set_user_query(q, 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 click one above:", key="user_query_input")
+    if user_query.strip():
+        reasoning_mode = mode == "Extended (Document + General)"
+        with st.spinner("💭 Generating your answer..."):
+            retrieved = retrieve_chunks(
+                user_query,
+                st.session_state["index"],
+                st.session_state["chunks"],
+                top_k=top_k,
+                embeddings=st.session_state["embeddings"]
+            )
+            answer = generate_answer(user_query, retrieved, reasoning_mode=reasoning_mode)
+            st.session_state["retrieved"] = retrieved
+        st.markdown("### 🤖 Assistant’s Answer")
+        st.markdown(f"<div class='answer-box'>{answer}</div>", unsafe_allow_html=True)
+
+# ==========================================================
+# 🎨 Sidebar scroll style
+# ==========================================================
+st.markdown("""
+<style>
+section[data-testid="stSidebar"] div.stExpander {max-height:480px;overflow-y:auto;}
+</style>
+""", unsafe_allow_html=True)