Update app.py

app.py

@@ -1,6 +1,6 @@
 """
-DocuMind - Streamlit Frontend for HuggingFace Spaces
-Simplified version that works with existing FastAPI backend
+DocuMind - Enterprise RAG Chatbot
+Redesigned UI matching modern document intelligence interface
 """
 
 import streamlit as st
@@ -8,226 +8,685 @@ import requests
 import json
 import os
 import time
+from datetime import datetime
 
 # ============================================================================
-# STREAMLIT CONFIG (Must be first)
+# PAGE CONFIG (Must be first)
 # ============================================================================
-
 st.set_page_config(
-    page_title="DocuMind - Enterprise RAG",
+    page_title="DocuMind AI",
     page_icon="🧠",
     layout="wide",
     initial_sidebar_state="expanded"
 )
 
 # ============================================================================
-# API CONFIGURATION
+# CUSTOM CSS — Full UI Override
 # ============================================================================
+st.markdown("""
+<style>
+@import url('https://fonts.googleapis.com/css2?family=Syne:wght@400;500;600;700;800&family=DM+Sans:ital,opsz,wght@0,9..40,300;0,9..40,400;0,9..40,500;0,9..40,600;1,9..40,300&display=swap');
+
+/* ── ROOT VARIABLES ── */
+:root {
+    --bg-base:       #0d0f14;
+    --bg-card:       #141720;
+    --bg-sidebar:    #111318;
+    --bg-input:      #1a1d27;
+    --bg-user-msg:   #1e2235;
+    --bg-bot-msg:    #181b26;
+    --accent:        #4ade80;
+    --accent-dim:    #22c55e22;
+    --accent-muted:  #4ade8044;
+    --border:        #ffffff0f;
+    --border-hover:  #ffffff1a;
+    --text-primary:  #f0f2f8;
+    --text-secondary:#8892aa;
+    --text-muted:    #50586e;
+    --badge-active:  #4ade80;
+    --badge-bg:      #4ade8022;
+    --radius-sm:     8px;
+    --radius-md:     14px;
+    --radius-lg:     20px;
+    --shadow:        0 4px 24px #00000055;
+}
+
+/* ── GLOBAL RESET ── */
+*, *::before, *::after { box-sizing: border-box; }
+
+html, body, [data-testid="stAppViewContainer"] {
+    font-family: 'DM Sans', sans-serif;
+    background: var(--bg-base) !important;
+    color: var(--text-primary) !important;
+}
+
+/* ── HIDE STREAMLIT CHROME ── */
+#MainMenu, footer, header,
+[data-testid="stToolbar"],
+[data-testid="stDecoration"],
+.stDeployButton,
+[data-testid="collapsedControl"] { display: none !important; }
+
+/* ── SIDEBAR ── */
+[data-testid="stSidebar"] {
+    background: var(--bg-sidebar) !important;
+    border-right: 1px solid var(--border) !important;
+    min-width: 240px !important;
+    max-width: 240px !important;
+}
+[data-testid="stSidebar"] > div:first-child {
+    padding: 0 !important;
+}
+[data-testid="stSidebarContent"] {
+    padding: 0 !important;
+    background: transparent !important;
+}
+
+/* ── MAIN AREA ── */
+[data-testid="stMain"] {
+    background: var(--bg-base) !important;
+}
+.main .block-container {
+    padding: 0 !important;
+    max-width: 100% !important;
+}
+
+/* ── HIDE DEFAULT SIDEBAR WIDGETS LABELS ── */
+[data-testid="stSidebar"] .stMarkdown p {
+    color: var(--text-secondary);
+    font-size: 0.78rem;
+    font-family: 'DM Sans', sans-serif;
+}
+
+/* ── FILE UPLOADER ── */
+[data-testid="stFileUploader"] {
+    background: var(--bg-input) !important;
+    border: 1px dashed var(--border-hover) !important;
+    border-radius: var(--radius-md) !important;
+    padding: 8px !important;
+}
+[data-testid="stFileUploader"] label {
+    color: var(--text-secondary) !important;
+    font-size: 0.82rem !important;
+}
+[data-testid="stFileUploadDropzone"] {
+    background: transparent !important;
+    border: none !important;
+}
+[data-testid="stFileUploadDropzone"] p,
+[data-testid="stFileUploadDropzone"] span {
+    color: var(--text-muted) !important;
+    font-size: 0.78rem !important;
+}
+
+/* ── BUTTONS ── */
+.stButton > button {
+    background: var(--bg-input) !important;
+    color: var(--text-primary) !important;
+    border: 1px solid var(--border-hover) !important;
+    border-radius: var(--radius-sm) !important;
+    font-family: 'DM Sans', sans-serif !important;
+    font-size: 0.82rem !important;
+    font-weight: 500 !important;
+    padding: 6px 14px !important;
+    transition: all 0.2s ease !important;
+    width: 100% !important;
+}
+.stButton > button:hover {
+    background: var(--accent-dim) !important;
+    border-color: var(--accent-muted) !important;
+    color: var(--accent) !important;
+}
+
+/* ── CHAT INPUT ── */
+[data-testid="stChatInput"] {
+    background: var(--bg-input) !important;
+    border: 1px solid var(--border-hover) !important;
+    border-radius: var(--radius-lg) !important;
+    padding: 4px 8px !important;
+}
+[data-testid="stChatInput"] textarea {
+    background: transparent !important;
+    color: var(--text-primary) !important;
+    font-family: 'DM Sans', sans-serif !important;
+    font-size: 0.9rem !important;
+    border: none !important;
+    outline: none !important;
+}
+[data-testid="stChatInput"] textarea::placeholder {
+    color: var(--text-muted) !important;
+}
+[data-testid="stChatInput"] button {
+    background: var(--accent) !important;
+    border-radius: 50% !important;
+    color: #000 !important;
+}
+
+/* ── CHAT MESSAGES ── */
+[data-testid="stChatMessage"] {
+    background: transparent !important;
+    border: none !important;
+    padding: 4px 0 !important;
+}
+
+/* User messages */
+[data-testid="stChatMessage"][data-testid*="user"],
+.stChatMessage:has([data-testid="chatAvatarIcon-user"]) {
+    flex-direction: row-reverse !important;
+}
+
+/* ── SPINNER ── */
+.stSpinner > div {
+    border-top-color: var(--accent) !important;
+}
+
+/* ── ALERTS / WARNINGS ── */
+.stAlert {
+    background: var(--bg-card) !important;
+    border: 1px solid var(--border) !important;
+    border-radius: var(--radius-md) !important;
+    color: var(--text-secondary) !important;
+    font-size: 0.82rem !important;
+}
+.stSuccess {
+    border-color: var(--accent-muted) !important;
+    background: var(--accent-dim) !important;
+}
+
+/* ── DIVIDER ── */
+hr {
+    border-color: var(--border) !important;
+    margin: 8px 0 !important;
+}
+
+/* ── EXPANDER ── */
+.streamlit-expander {
+    background: var(--bg-card) !important;
+    border: 1px solid var(--border) !important;
+    border-radius: var(--radius-md) !important;
+}
+.streamlit-expander summary {
+    color: var(--text-secondary) !important;
+    font-size: 0.82rem !important;
+}
+
+/* ── SCROLLBAR ── */
+::-webkit-scrollbar { width: 4px; height: 4px; }
+::-webkit-scrollbar-track { background: transparent; }
+::-webkit-scrollbar-thumb { background: var(--border-hover); border-radius: 4px; }
+::-webkit-scrollbar-thumb:hover { background: var(--text-muted); }
+</style>
+""", unsafe_allow_html=True)
 
-# Use localhost for internal communication
+# ============================================================================
+# API CONFIG
+# ============================================================================
 API_URL = "http://127.0.0.1:8000"
-API_TIMEOUT = 30  # seconds
+API_TIMEOUT = 60
 
 # ============================================================================
-# UTILITY FUNCTIONS
+# SESSION STATE INIT
 # ============================================================================
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+if "backend_ok" not in st.session_state:
+    st.session_state.backend_ok = False
+if "backend_checked" not in st.session_state:
+    st.session_state.backend_checked = False
+if "chat_sessions" not in st.session_state:
+    st.session_state.chat_sessions = ["Q3 Financial Report", "Product Roadmap 2...", "Research Summary"]
+if "mode_summary" not in st.session_state:
+    st.session_state.mode_summary = False
+if "mode_citations" not in st.session_state:
+    st.session_state.mode_citations = True
+if "ingested_docs" not in st.session_state:
+    st.session_state.ingested_docs = []
 
-def check_backend_health(retries=5):
-    """Check if backend API is running"""
-    for attempt in range(retries):
-        try:
-            response = requests.get(f"{API_URL}/docs", timeout=5)
-            if response.status_code == 200:
-                return True
-        except:
-            if attempt < retries - 1:
-                time.sleep(1)
-    return False
-
-def safe_api_call(method, endpoint, **kwargs):
-    """Safely call API with error handling"""
+# ============================================================================
+# HELPERS
+# ============================================================================
+def check_backend():
     try:
+        r = requests.get(f"{API_URL}/docs", timeout=5)
+        return r.status_code == 200
+    except:
+        return False
+
+def api_call(method, endpoint, **kwargs):
+    try:
+        kwargs.setdefault("timeout", API_TIMEOUT)
         url = f"{API_URL}{endpoint}"
-        kwargs.setdefault('timeout', API_TIMEOUT)
-        
-        if method == "GET":
-            response = requests.get(url, **kwargs)
-        elif method == "POST":
-            response = requests.post(url, **kwargs)
-        else:
-            return None, "Invalid method"
-        
-        return response, None
+        fn = requests.get if method == "GET" else requests.post
+        return fn(url, **kwargs), None
     except requests.exceptions.Timeout:
         return None, "⏱️ Request timed out"
     except requests.exceptions.ConnectionError:
         return None, "⚠️ Backend not responding"
     except Exception as e:
-        return None, f"❌ Error: {str(e)}"
+        return None, f"Error: {e}"
+
+def now_time():
+    return datetime.now().strftime("%-I:%M %p")
+
+# Backend health check (once per session)
+if not st.session_state.backend_checked:
+    with st.spinner("Starting up..."):
+        time.sleep(2)
+        st.session_state.backend_ok = check_backend()
+    st.session_state.backend_checked = True
+
+# Refresh doc list
+try:
+    resp, _ = api_call("GET", "/sources")
+    if resp and resp.status_code == 200:
+        st.session_state.ingested_docs = resp.json().get("documents", [])
+except:
+    pass
 
 # ============================================================================
-# STREAMLIT FRONTEND UI
+# SIDEBAR
 # ============================================================================
+with st.sidebar:
+    # ── Logo ──
+    st.markdown("""
+    <div style="padding: 22px 20px 16px; border-bottom: 1px solid #ffffff0f;">
+        <div style="display:flex; align-items:center; gap:10px; margin-bottom:2px;">
+            <div style="width:32px;height:32px;background:linear-gradient(135deg,#4ade80,#22c55e);
+                        border-radius:8px;display:flex;align-items:center;justify-content:center;
+                        font-size:16px;flex-shrink:0;">🧠</div>
+            <div>
+                <div style="font-family:'Syne',sans-serif;font-weight:700;font-size:1rem;
+                            color:#f0f2f8;line-height:1.1;">DocuMind AI</div>
+                <div style="font-size:0.68rem;color:#50586e;letter-spacing:0.04em;">Digital Curator</div>
+            </div>
+        </div>
+    </div>
+    """, unsafe_allow_html=True)
 
-st.title("🧠 DocuMind")
-st.markdown("**Enterprise Document Intelligence Chatbot**")
-st.markdown("---")
+    # ── New Chat Button ──
+    st.markdown('<div style="padding:14px 16px 8px;">', unsafe_allow_html=True)
+    if st.button("＋  New Chat", key="new_chat"):
+        st.session_state.messages = []
+        st.rerun()
+    st.markdown('</div>', unsafe_allow_html=True)
 
-# Health check on load
-if "backend_checked" not in st.session_state:
-    st.session_state.backend_checked = False
-    st.session_state.backend_healthy = False
+    # ── Upload Section ──
+    st.markdown("""
+    <div style="padding:0 16px 8px;">
+        <div style="background:#1a1d27;border:1px solid #ffffff0f;border-radius:12px;padding:12px;">
+            <div style="display:flex;justify-content:space-between;align-items:center;margin-bottom:8px;">
+                <span style="font-size:0.78rem;color:#8892aa;font-weight:500;">📎 Upload</span>
+                <span style="font-size:0.68rem;color:#50586e;background:#ffffff08;
+                             padding:2px 6px;border-radius:4px;">200MB LIMIT</span>
+            </div>
+        </div>
+    </div>
+    """, unsafe_allow_html=True)
 
-if not st.session_state.backend_checked:
-    with st.spinner("🔄 Starting backend..."):
-        time.sleep(2)  # Give backend time to start
-        st.session_state.backend_healthy = check_backend_health()
-    st.session_state.backend_checked = True
+    with st.container():
+        st.markdown('<div style="padding:0 16px 6px;">', unsafe_allow_html=True)
+        uploaded_file = st.file_uploader(
+            "Drop file here",
+            type=["txt", "pdf", "docx"],
+            label_visibility="collapsed"
+        )
+        if uploaded_file and st.button("Ingest Document", key="ingest_btn"):
+            with st.spinner("Processing..."):
+                files = {"file": (uploaded_file.name, uploaded_file.getvalue())}
+                resp, err = api_call("POST", "/ingest", files=files)
+                if err:
+                    st.error(err)
+                elif resp and resp.status_code == 200:
+                    st.success(f"✅ Ingested!")
+                    st.session_state.ingested_docs.append(uploaded_file.name)
+                else:
+                    detail = ""
+                    try:
+                        detail = resp.json().get("detail", resp.text)
+                    except:
+                        detail = resp.text if resp else "Unknown error"
+                    st.error(f"❌ {detail}")
+        st.markdown('</div>', unsafe_allow_html=True)
 
-# Show status
-if not st.session_state.backend_healthy:
-    st.warning(
-        "⚠️ Backend is starting up. This may take 30-60 seconds on first load. "
-        "Please refresh the page if you see this message for more than 1 minute."
-    )
+    # ── Recents ──
+    st.markdown("""
+    <div style="padding:16px 20px 8px;">
+        <div style="font-size:0.68rem;color:#50586e;letter-spacing:0.08em;
+                    font-weight:600;text-transform:uppercase;margin-bottom:10px;">Recents</div>
+    </div>
+    """, unsafe_allow_html=True)
 
-# --- Sidebar for Document Upload ---
-with st.sidebar:
-    st.header("🏢 Document Knowledge Base")
-    st.markdown("Upload PDFs, DOCX, or TXT documents to add them to the system.")
-    
-    uploaded_file = st.file_uploader("Upload a new document", type=["txt", "pdf", "docx"])
-    
-    if uploaded_file and st.button("Ingest Document", key="ingest_btn"):
-        with st.spinner("Ingesting document (creating chunks & embeddings)..."):
-            files = {"file": (uploaded_file.name, uploaded_file.getvalue())}
-            response, error = safe_api_call("POST", "/ingest", files=files)
-            
-            if error:
-                st.error(error)
-            elif response and response.status_code == 200:
-                st.success(f"✅ {uploaded_file.name} ingested successfully!")
-            else:
-                st.error(f"❌ Failed to ingest: {response.text if response else 'Unknown error'}")
-    
-    st.divider()
-    
-    st.subheader("📄 Indexed Documents")
-    response, error = safe_api_call("GET", "/sources")
-    
-    if error:
-        st.warning(f"Could not fetch documents: {error}")
-    elif response and response.status_code == 200:
-        documents = response.json().get("documents", [])
-        if documents:
-            for doc in documents:
-                st.markdown(f"- `{doc}`")
-        else:
-            st.info("No documents indexed yet.")
-    else:
-        st.info("Backend not ready yet...")
+    # Show indexed docs or placeholder sessions
+    display_docs = st.session_state.ingested_docs if st.session_state.ingested_docs else st.session_state.chat_sessions
+    for doc in display_docs[:6]:
+        label = doc if len(doc) <= 22 else doc[:19] + "..."
+        st.markdown(f"""
+        <div style="display:flex;align-items:center;gap:8px;padding:7px 20px;
+                    cursor:pointer;transition:background 0.15s;border-radius:0;">
+            <span style="color:#50586e;font-size:0.75rem;">🕓</span>
+            <span style="font-size:0.8rem;color:#8892aa;">{label}</span>
+        </div>
+        """, unsafe_allow_html=True)
 
-# --- Main Chat Interface ---
-st.subheader("💬 Chat with Your Documents")
+    # ── Bottom nav ──
+    st.markdown("""
+    <div style="position:fixed;bottom:0;left:0;width:240px;
+                border-top:1px solid #ffffff0f;background:#111318;padding:12px 16px;">
+        <div style="display:flex;flex-direction:column;gap:2px;margin-bottom:12px;">
+            <div style="display:flex;align-items:center;gap:8px;padding:7px 8px;
+                        border-radius:8px;cursor:pointer;">
+                <span style="font-size:0.85rem;">⚙️</span>
+                <span style="font-size:0.82rem;color:#8892aa;">Settings</span>
+            </div>
+            <div style="display:flex;align-items:center;gap:8px;padding:7px 8px;
+                        border-radius:8px;cursor:pointer;">
+                <span style="font-size:0.85rem;">❓</span>
+                <span style="font-size:0.82rem;color:#8892aa;">Help</span>
+            </div>
+        </div>
+        <div style="display:flex;align-items:center;gap:10px;">
+            <div style="width:28px;height:28px;background:linear-gradient(135deg,#667eea,#764ba2);
+                        border-radius:50%;display:flex;align-items:center;justify-content:center;
+                        font-size:12px;color:white;">A</div>
+            <div>
+                <div style="font-size:0.8rem;color:#f0f2f8;font-weight:500;">Alex Sterling</div>
+                <div style="font-size:0.68rem;color:#50586e;">Pro Curator</div>
+            </div>
+        </div>
+    </div>
+    """, unsafe_allow_html=True)
 
-# Initialize session state
-if "messages" not in st.session_state:
-    st.session_state.messages = []
+# ============================================================================
+# MAIN CHAT AREA
+# ============================================================================
+
+# ── Top bar ──
+col_top1, col_top2, col_top3 = st.columns([3, 6, 1])
+with col_top1:
+    st.markdown("""
+    <div style="padding:16px 0 12px;display:flex;align-items:center;gap:10px;">
+        <span style="font-size:0.72rem;color:#8892aa;font-weight:600;
+                     letter-spacing:0.08em;text-transform:uppercase;">Current Session</span>
+        <span style="background:#4ade8022;color:#4ade80;font-size:0.62rem;
+                     font-weight:600;padding:2px 8px;border-radius:4px;
+                     letter-spacing:0.06em;text-transform:uppercase;">ACTIVE</span>
+    </div>
+    """, unsafe_allow_html=True)
+with col_top3:
+    st.markdown("""
+    <div style="padding:16px 0 12px;display:flex;justify-content:flex-end;gap:12px;">
+        <span style="color:#8892aa;cursor:pointer;font-size:1.1rem;">🔔</span>
+        <span style="color:#8892aa;cursor:pointer;font-size:1.1rem;">👤</span>
+    </div>
+    """, unsafe_allow_html=True)
+
+st.markdown('<hr style="border-color:#ffffff0f;margin:0 0 0 0;"/>', unsafe_allow_html=True)
+
+# Backend warning
+if not st.session_state.backend_ok:
+    st.markdown("""
+    <div style="background:#f59e0b15;border:1px solid #f59e0b33;border-radius:10px;
+                padding:10px 16px;margin:12px 0;font-size:0.82rem;color:#f59e0b;">
+        ⚠️ Backend is starting up (30–60s on first load). Refresh if this persists.
+    </div>
+    """, unsafe_allow_html=True)
 
-# Display chat history
+# ── Empty state ──
+if not st.session_state.messages:
+    st.markdown("""
+    <div style="display:flex;flex-direction:column;align-items:center;
+                justify-content:center;padding:80px 20px 40px;text-align:center;">
+        <h1 style="font-family:'Syne',sans-serif;font-size:2rem;font-weight:700;
+                   color:#f0f2f8;margin:0 0 12px;line-height:1.2;">
+            How can I assist your research today?
+        </h1>
+        <p style="color:#8892aa;font-size:0.95rem;max-width:420px;line-height:1.6;margin:0;">
+            I can analyze documents, summarize findings, or answer<br>
+            specific questions about your archived data.
+        </p>
+    </div>
+    """, unsafe_allow_html=True)
+
+# ── Chat history ──
 for msg in st.session_state.messages:
-    with st.chat_message(msg["role"]):
-        st.markdown(msg["content"])
-        if "sources" in msg and msg.get("sources"):
-            with st.expander("📚 Show Sources"):
-                for idx, src in enumerate(msg["sources"]):
-                    score = src.get('score', 0)
-                    st.caption(f"**Source {idx+1}** [Relevance: {score:.2%}]")
-                    st.markdown(f"**From:** `{src.get('source', 'Unknown')}`")
-                    content = src.get('content', '')
-                    if len(content) > 500:
-                        st.markdown(f"> {content[:500]}...")
-                    else:
-                        st.markdown(f"> {content}")
-
-# Chat input
-user_input = st.chat_input("Ask a question about your documents...")
+    role = msg["role"]
+    content = msg["content"]
+    ts = msg.get("time", "")
+    sources = msg.get("sources", [])
+
+    if role == "user":
+        st.markdown(f"""
+        <div style="display:flex;justify-content:flex-end;margin:8px 0 2px;">
+            <div style="max-width:65%;">
+                <div style="background:#1e2235;border:1px solid #ffffff0f;
+                            border-radius:16px 16px 4px 16px;padding:14px 18px;
+                            color:#f0f2f8;font-size:0.9rem;line-height:1.6;">
+                    {content}
+                </div>
+                <div style="text-align:right;font-size:0.68rem;color:#50586e;
+                            margin-top:4px;padding-right:4px;">{ts}</div>
+            </div>
+        </div>
+        """, unsafe_allow_html=True)
+    else:
+        # Format bullet-like lines
+        lines = content.split("\n")
+        formatted = ""
+        for line in lines:
+            line = line.strip()
+            if not line:
+                continue
+            if line.startswith(("- ", "• ", "* ")):
+                item = line[2:]
+                # Bold first phrase before colon
+                if ":" in item:
+                    key, val = item.split(":", 1)
+                    formatted += f"""
+                    <div style="display:flex;align-items:flex-start;gap:10px;margin:6px 0;">
+                        <span style="color:#4ade80;margin-top:2px;flex-shrink:0;">✓</span>
+                        <span><strong style="color:#f0f2f8;">{key.strip()}:</strong>
+                              <span style="color:#c4cad9;">{val.strip()}</span></span>
+                    </div>"""
+                else:
+                    formatted += f"""
+                    <div style="display:flex;align-items:flex-start;gap:10px;margin:6px 0;">
+                        <span style="color:#4ade80;margin-top:2px;flex-shrink:0;">✓</span>
+                        <span style="color:#c4cad9;">{item}</span>
+                    </div>"""
+            else:
+                formatted += f'<p style="color:#c4cad9;margin:4px 0;line-height:1.7;">{line}</p>'
+
+        st.markdown(f"""
+        <div style="display:flex;justify-content:flex-start;margin:8px 0 2px;gap:10px;">
+            <div style="width:32px;height:32px;background:linear-gradient(135deg,#4ade80,#22c55e);
+                        border-radius:50%;display:flex;align-items:center;justify-content:center;
+                        font-size:14px;flex-shrink:0;margin-top:4px;">🧠</div>
+            <div style="max-width:68%;">
+                <div style="background:#181b26;border:1px solid #ffffff0f;
+                            border-radius:4px 16px 16px 16px;padding:14px 18px;
+                            font-size:0.9rem;line-height:1.6;">
+                    {formatted}
+                </div>
+                <div style="font-size:0.68rem;color:#50586e;margin-top:4px;padding-left:4px;">{ts}</div>
+            </div>
+        </div>
+        """, unsafe_allow_html=True)
+
+        if sources and st.session_state.mode_citations:
+            with st.expander(f"📚 {len(sources)} source(s) cited"):
+                for i, src in enumerate(sources):
+                    score = src.get("score", 0)
+                    st.markdown(f"""
+                    <div style="background:#1a1d27;border:1px solid #ffffff0a;border-radius:8px;
+                                padding:10px 14px;margin:6px 0;">
+                        <div style="display:flex;justify-content:space-between;margin-bottom:6px;">
+                            <span style="font-size:0.78rem;color:#4ade80;font-weight:600;">
+                                Source {i+1}</span>
+                            <span style="font-size:0.72rem;color:#8892aa;background:#ffffff08;
+                                         padding:1px 7px;border-radius:4px;">
+                                {score:.0%} match</span>
+                        </div>
+                        <div style="font-size:0.76rem;color:#50586e;margin-bottom:6px;">
+                            📄 {src.get('source','Unknown')}</div>
+                        <div style="font-size:0.82rem;color:#8892aa;line-height:1.6;
+                                    border-left:2px solid #4ade8044;padding-left:10px;">
+                            {src.get('content','')[:400]}{'...' if len(src.get('content',''))>400 else ''}
+                        </div>
+                    </div>
+                    """, unsafe_allow_html=True)
+
+# ── Chat Input ──
+st.markdown("<div style='height:16px;'></div>", unsafe_allow_html=True)
+user_input = st.chat_input("Ask a question...")
 
 if user_input:
-    # Add user message to history
-    st.session_state.messages.append({"role": "user", "content": user_input})
-    
-    # Display user message
-    with st.chat_message("user"):
-        st.markdown(user_input)
-    
-    # Get assistant response
-    with st.chat_message("assistant"):
+    ts = now_time()
+    st.session_state.messages.append({"role": "user", "content": user_input, "time": ts})
+
+    # Render user bubble immediately
+    st.markdown(f"""
+    <div style="display:flex;justify-content:flex-end;margin:8px 0 2px;">
+        <div style="max-width:65%;">
+            <div style="background:#1e2235;border:1px solid #ffffff0f;
+                        border-radius:16px 16px 4px 16px;padding:14px 18px;
+                        color:#f0f2f8;font-size:0.9rem;line-height:1.6;">{user_input}</div>
+            <div style="text-align:right;font-size:0.68rem;color:#50586e;
+                        margin-top:4px;padding-right:4px;">{ts}</div>
+        </div>
+    </div>
+    """, unsafe_allow_html=True)
+
+    full_response = ""
+    sources = []
+
+    # Assistant bubble wrapper
+    with st.container():
+        st.markdown("""
+        <div style="display:flex;justify-content:flex-start;margin:8px 0;gap:10px;">
+            <div style="width:32px;height:32px;background:linear-gradient(135deg,#4ade80,#22c55e);
+                        border-radius:50%;display:flex;align-items:center;justify-content:center;
+                        font-size:14px;flex-shrink:0;margin-top:4px;">🧠</div>
+            <div style="max-width:68%;min-width:120px;">
+        """, unsafe_allow_html=True)
+
         placeholder = st.empty()
-        full_response = ""
-        sources = []
-        
-        response, error = safe_api_call(
-            "POST",
-            "/query",
-            json={"question": user_input},
-            stream=True
-        )
-        
-        if error:
-            error_msg = error
-            st.error(error_msg)
-            full_response = error_msg
-        elif response:
+
+        resp, err = api_call("POST", "/query", json={"question": user_input}, stream=True)
+
+        if err:
+            full_response = err
+            placeholder.markdown(f"""
+            <div style="background:#181b26;border:1px solid #f87171aa;
+                        border-radius:4px 16px 16px 16px;padding:14px 18px;
+                        color:#f87171;font-size:0.9rem;">{err}</div>
+            """, unsafe_allow_html=True)
+        elif resp:
             try:
-                for line in response.iter_lines():
+                buf = ""
+                for line in resp.iter_lines():
                     if line:
                         try:
-                            decoded_line = line.decode('utf-8')
-                            data = json.loads(decoded_line)
-                            
+                            data = json.loads(line.decode("utf-8"))
                             if data.get("type") == "sources":
                                 sources = data.get("data", [])
                             elif data.get("type") == "token":
-                                full_response += data.get("content", "")
-                                placeholder.markdown(full_response + "▌")
+                                buf += data.get("content", "")
+                                full_response = buf
+                                placeholder.markdown(f"""
+                                <div style="background:#181b26;border:1px solid #ffffff0f;
+                                            border-radius:4px 16px 16px 16px;padding:14px 18px;
+                                            color:#c4cad9;font-size:0.9rem;line-height:1.7;">{buf}▌</div>
+                                """, unsafe_allow_html=True)
                         except json.JSONDecodeError:
                             continue
-                
-                placeholder.markdown(full_response)
-                
-                # Display sources if available
-                if sources:
-                    with st.expander("📚 Show Sources"):
-                        for idx, src in enumerate(sources):
-                            score = src.get('score', 0)
-                            st.caption(f"**Source {idx+1}** [Relevance: {score:.2%}]")
-                            st.markdown(f"**From:** `{src.get('source', 'Unknown')}`")
-                            content = src.get('content', '')
-                            if len(content) > 500:
-                                st.markdown(f"> {content[:500]}...")
-                            else:
-                                st.markdown(f"> {content}")
-            
+
+                # Final render
+                placeholder.markdown(f"""
+                <div style="background:#181b26;border:1px solid #ffffff0f;
+                            border-radius:4px 16px 16px 16px;padding:14px 18px;
+                            color:#c4cad9;font-size:0.9rem;line-height:1.7;">{full_response}</div>
+                """, unsafe_allow_html=True)
+
             except Exception as e:
-                error_msg = f"❌ Error processing response: {str(e)}"
-                placeholder.markdown(error_msg)
-                st.error(error_msg)
-                full_response = error_msg
-    
-    # Save assistant message
+                full_response = f"❌ Error: {e}"
+                placeholder.markdown(f"""
+                <div style="background:#181b26;border:1px solid #f87171aa;
+                            border-radius:4px 16px 16px 16px;padding:14px 18px;
+                            color:#f87171;font-size:0.9rem;">{full_response}</div>
+                """, unsafe_allow_html=True)
+
+        st.markdown("</div></div>", unsafe_allow_html=True)
+
+        if sources and st.session_state.mode_citations:
+            with st.expander(f"📚 {len(sources)} source(s) cited"):
+                for i, src in enumerate(sources):
+                    score = src.get("score", 0)
+                    st.markdown(f"""
+                    <div style="background:#1a1d27;border:1px solid #ffffff0a;border-radius:8px;
+                                padding:10px 14px;margin:6px 0;">
+                        <div style="display:flex;justify-content:space-between;margin-bottom:6px;">
+                            <span style="font-size:0.78rem;color:#4ade80;font-weight:600;">
+                                Source {i+1}</span>
+                            <span style="font-size:0.72rem;color:#8892aa;">
+                                {score:.0%} match</span>
+                        </div>
+                        <div style="font-size:0.76rem;color:#50586e;margin-bottom:6px;">
+                            📄 {src.get('source','Unknown')}</div>
+                        <div style="font-size:0.82rem;color:#8892aa;line-height:1.6;
+                                    border-left:2px solid #4ade8044;padding-left:10px;">
+                            {src.get('content','')[:400]}{'...' if len(src.get('content',''))>400 else ''}
+                        </div>
+                    </div>
+                    """, unsafe_allow_html=True)
+
+    # Save to session
+    bot_ts = now_time()
     st.session_state.messages.append({
         "role": "assistant",
         "content": full_response,
-        "sources": sources
+        "sources": sources,
+        "time": bot_ts
     })
 
-# Footer
-st.divider()
-st.markdown(
-    "<div style='text-align: center; color: var(--color-text-secondary); font-size: 0.85em;'>"
-    "DocuMind - Enterprise RAG Chatbot | "
-    "<a href='https://github.com/TejeshNaiduKona/DocuMind' target='_blank'>GitHub</a>"
-    "</div>",
-    unsafe_allow_html=True
-)
+# ── Bottom toolbar ──
+st.markdown("<div style='height:8px;'></div>", unsafe_allow_html=True)
+st.markdown('<hr style="border-color:#ffffff08;margin:0;"/>', unsafe_allow_html=True)
+
+c1, c2, c3, c4 = st.columns([1, 1, 1, 5])
+with c1:
+    active_s = "#4ade80" if st.session_state.mode_summary else "#50586e"
+    if st.button("⇅ Summary Mode", key="mode_sum"):
+        st.session_state.mode_summary = not st.session_state.mode_summary
+        st.rerun()
+with c2:
+    active_c = "#4ade80" if st.session_state.mode_citations else "#50586e"
+    if st.button("⊟ Source Citations", key="mode_cite"):
+        st.session_state.mode_citations = not st.session_state.mode_citations
+        st.rerun()
+with c3:
+    if st.button("⇌ Translate", key="mode_trans"):
+        st.toast("Translation mode coming soon!", icon="🌐")
+
+st.markdown("""
+<style>
+/* Bottom toolbar buttons */
+div[data-testid="stHorizontalBlock"] > div:nth-child(1) button,
+div[data-testid="stHorizontalBlock"] > div:nth-child(2) button,
+div[data-testid="stHorizontalBlock"] > div:nth-child(3) button {
+    background: transparent !important;
+    border: none !important;
+    color: #50586e !important;
+    font-size: 0.72rem !important;
+    font-weight: 500 !important;
+    letter-spacing: 0.04em !important;
+    padding: 4px 8px !important;
+    width: auto !important;
+}
+div[data-testid="stHorizontalBlock"] > div:nth-child(1) button:hover,
+div[data-testid="stHorizontalBlock"] > div:nth-child(2) button:hover,
+div[data-testid="stHorizontalBlock"] > div:nth-child(3) button:hover {
+    color: #4ade80 !important;
+    background: transparent !important;
+}
+</style>
+""", unsafe_allow_html=True)