Update app.py

app.py

@@ -1,795 +1,365 @@
 """
-DocuMind AI - Minimal Enterprise RAG Chat UI
-Single-file Streamlit app matching a clean document-intelligence interface.
-
-Features:
-- Minimal light UI inspired by the provided reference image
-- Left sidebar with New Chat + Upload
-- Native Streamlit sidebar open/close toggle (top-left)
-- Chat history with user/assistant bubbles
-- Optional source citations from backend
-- Works with an existing FastAPI backend
+DocuMind - Streamlit Frontend for HuggingFace Spaces
+Minimal & Clean UI implementation based on the provided mockup.
 """
 
-from __future__ import annotations
-
-import html
+import streamlit as st
 import json
 import time
-from datetime import datetime
-
 import requests
-import streamlit as st
-
 
 # ============================================================================
-# PAGE CONFIG
+# STREAMLIT CONFIG (Must be first)
 # ============================================================================
+
 st.set_page_config(
     page_title="DocuMind AI",
     page_icon="🧠",
     layout="wide",
-    initial_sidebar_state="expanded",
+    initial_sidebar_state="expanded"
 )
 
-
 # ============================================================================
-# THEME + MINIMAL CSS
+# CUSTOM CSS (Theming to match the minimal clean mockup)
 # ============================================================================
-st.markdown(
-    """
-    <style>
-    @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700;800&display=swap');
-
-    :root {
-        --bg: #f6f8fb;
-        --panel: #ffffff;
-        --panel-soft: #f2f5f8;
-        --panel-dark: #0f4c5c;
-        --panel-dark-2: #0b3d49;
-        --text: #0f172a;
-        --muted: #667085;
-        --muted-2: #94a3b8;
-        --border: rgba(15, 23, 42, 0.08);
-        --shadow: 0 8px 30px rgba(15, 23, 42, 0.06);
-        --radius-xl: 28px;
-        --radius-lg: 22px;
-        --radius-md: 16px;
-        --radius-sm: 12px;
-        --accent: #0f4c5c;
-        --accent-soft: rgba(15, 76, 92, 0.10);
-    }
-
-    html, body, [data-testid="stAppViewContainer"] {
-        background: var(--bg) !important;
-        color: var(--text) !important;
-        font-family: 'Inter', sans-serif !important;
+custom_css = """
+<style>
+    /* Base Font */
+    @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&display=swap');
+    
+    html, body, [class*="css"] {
+        font-family: 'Inter', sans-serif;
+    }
+    
+    /* Main Backgrounds */
+    .stApp {
+        background-color: #FAFBFC;
     }
-
-    *, *::before, *::after { box-sizing: border-box; }
-
-    /* Hide Streamlit chrome */
-    #MainMenu, footer, header, [data-testid="stToolbar"], [data-testid="stDecoration"] {
-        display: none !important;
-    }
-
-    /* Sidebar */
     [data-testid="stSidebar"] {
-        background: linear-gradient(180deg, #f8fafc 0%, #f5f7fa 100%) !important;
-        border-right: 1px solid var(--border) !important;
-    }
-
-    [data-testid="stSidebarContent"] {
-        padding-top: 0 !important;
-    }
-
-    /* Main container */
-    .main .block-container {
-        padding: 0 !important;
-        max-width: 100% !important;
-    }
-
-    /* Sidebar collapse button */
-    [data-testid="collapsedControl"] {
-        color: #0f172a !important;
-        background: rgba(255,255,255,0.75) !important;
-        border: 1px solid var(--border) !important;
-        border-radius: 999px !important;
-        box-shadow: var(--shadow) !important;
-    }
-
-    /* Buttons */
-    .stButton > button {
-        border-radius: 18px !important;
-        border: 1px solid var(--border) !important;
-        background: var(--panel) !important;
-        color: var(--text) !important;
-        font-family: 'Inter', sans-serif !important;
-        font-weight: 600 !important;
-        transition: 0.18s ease !important;
-        box-shadow: none !important;
-        width: 100% !important;
-        padding: 0.8rem 1rem !important;
-    }
-    .stButton > button:hover {
-        transform: translateY(-1px) !important;
-        border-color: rgba(15, 76, 92, 0.18) !important;
-        background: rgba(15, 76, 92, 0.04) !important;
-    }
-
-    /* Sidebar file uploader */
-    [data-testid="stFileUploader"] {
-        background: var(--panel) !important;
-        border: 1px dashed rgba(15, 23, 42, 0.10) !important;
-        border-radius: 18px !important;
-        padding: 10px !important;
-        box-shadow: none !important;
-    }
-    [data-testid="stFileUploader"] label,
-    [data-testid="stFileUploadDropzone"] p,
-    [data-testid="stFileUploadDropzone"] span {
-        color: var(--muted) !important;
-        font-family: 'Inter', sans-serif !important;
-    }
-
-    /* Chat input */
-    [data-testid="stChatInput"] {
-        background: transparent !important;
-    }
-    [data-testid="stChatInput"] textarea {
-        background: var(--panel) !important;
-        border-radius: 24px !important;
-        border: 1px solid rgba(15, 23, 42, 0.08) !important;
-        box-shadow: var(--shadow) !important;
-        font-family: 'Inter', sans-serif !important;
-        padding: 1rem 1rem 1rem 1rem !important;
-        color: var(--text) !important;
-    }
-    [data-testid="stChatInput"] textarea::placeholder {
-        color: var(--muted-2) !important;
-    }
-    [data-testid="stChatInput"] button {
-        background: var(--accent) !important;
-        color: #fff !important;
-        border-radius: 999px !important;
-    }
-
-    /* Spacers */
-    hr {
-        border: none !important;
-        border-top: 1px solid rgba(15, 23, 42, 0.06) !important;
-        margin: 0 !important;
-    }
-
-    /* Scrollbar */
-    ::-webkit-scrollbar { width: 8px; height: 8px; }
-    ::-webkit-scrollbar-track { background: transparent; }
-    ::-webkit-scrollbar-thumb { background: rgba(100, 116, 139, 0.25); border-radius: 999px; }
-    ::-webkit-scrollbar-thumb:hover { background: rgba(100, 116, 139, 0.45); }
-
-    /* Hero */
-    .hero-wrap {
-        display: flex;
-        flex-direction: column;
-        align-items: center;
-        justify-content: center;
-        text-align: center;
-        padding: 5rem 1.25rem 2rem;
-        min-height: 48vh;
-    }
-    .hero-title {
-        font-size: clamp(2rem, 3.5vw, 3.5rem);
-        line-height: 1.08;
-        font-weight: 800;
-        letter-spacing: -0.04em;
-        color: #0f172a;
-        margin: 0 0 1rem 0;
+        background-color: #F4F6F8 !important;
+        border-right: 1px solid #E5E7EB;
     }
-    .hero-subtitle {
-        max-width: 760px;
-        font-size: 1.02rem;
-        line-height: 1.65;
-        color: var(--muted);
-        margin: 0;
+    
+    /* Hide top bar & footer */
+    header {visibility: hidden;}
+    footer {visibility: hidden;}
+    .block-container {
+        padding-top: 2rem !important;
+        max-width: 900px;
     }
 
-    /* Session header */
-    .session-row {
-        display: flex;
-        align-items: center;
-        gap: 0.75rem;
-        padding: 1rem 1.25rem 0.85rem 1.25rem;
-    }
-    .session-label {
-        font-size: 0.95rem;
-        font-weight: 800;
-        letter-spacing: 0.02em;
-        color: #0f4c5c;
-        text-transform: uppercase;
-    }
-    .session-badge {
-        display: inline-flex;
-        align-items: center;
-        padding: 0.28rem 0.65rem;
-        border-radius: 999px;
-        background: rgba(15, 76, 92, 0.10);
-        color: #0f4c5c;
-        font-size: 0.7rem;
-        font-weight: 800;
-        letter-spacing: 0.08em;
-    }
-
-    /* Sidebar content blocks */
-    .sidebar-brand {
-        padding: 1.1rem 1rem 1rem;
-        border-bottom: 1px solid rgba(15, 23, 42, 0.06);
-    }
-    .brand-row {
+    /* Sidebar Logo Header */
+    .sidebar-header {
         display: flex;
         align-items: center;
-        gap: 0.75rem;
+        gap: 12px;
+        margin-bottom: 24px;
+        margin-top: -30px;
     }
-    .brand-icon {
-        width: 44px;
-        height: 44px;
-        border-radius: 14px;
-        background: linear-gradient(135deg, #0f4c5c 0%, #116579 100%);
+    .logo-icon {
+        background-color: #064052;
         color: white;
+        width: 32px;
+        height: 32px;
+        border-radius: 8px;
         display: flex;
-        align-items: center;
         justify-content: center;
-        font-size: 1.1rem;
-        box-shadow: var(--shadow);
+        align-items: center;
+        font-weight: bold;
+        font-size: 16px;
     }
     .brand-title {
-        font-size: 1.05rem;
-        font-weight: 800;
-        color: #0f172a;
-        line-height: 1.1;
+        color: #111827;
+        font-weight: 600;
+        font-size: 1.1rem;
+        line-height: 1.2;
     }
     .brand-subtitle {
-        font-size: 0.78rem;
-        color: var(--muted-2);
-        margin-top: 0.15rem;
+        color: #9CA3AF;
+        font-size: 0.75rem;
     }
 
-    .sidebar-section {
-        padding: 1rem;
-    }
-    .sidebar-heading {
-        font-size: 0.72rem;
-        font-weight: 800;
-        letter-spacing: 0.12em;
-        text-transform: uppercase;
-        color: #64748b;
-        margin: 0 0 0.6rem 0;
-    }
-    .recents-item {
-        display: flex;
-        align-items: center;
-        gap: 0.6rem;
-        padding: 0.65rem 0.2rem;
-        color: #334155;
-        font-size: 0.92rem;
-        border-radius: 12px;
+    /* Override New Chat Button */
+    div[data-testid="stSidebar"] div[data-testid="stButton"] button {
+        width: 100%;
+        border-radius: 20px;
+        background-color: #EFEFEF;
+        color: #064052;
+        border: none;
+        font-weight: 500;
+        padding: 0.5rem 1rem;
     }
-    .recents-item span.icon {
-        color: #64748b;
-        font-size: 0.85rem;
+    div[data-testid="stSidebar"] div[data-testid="stButton"] button:hover {
+        background-color: #E5E7EB;
     }
 
-    /* Message bubbles */
-    .msg-user-wrap {
-        display: flex;
-        justify-content: flex-end;
-        margin: 0.8rem 0 0.25rem;
-    }
-    .msg-user {
-        max-width: min(720px, 66%);
-        background: linear-gradient(180deg, #0f4c5c 0%, #0b3d49 100%);
-        color: #fff;
-        border-radius: 22px 22px 6px 22px;
-        padding: 1rem 1.1rem;
-        box-shadow: var(--shadow);
-        font-size: 0.98rem;
-        line-height: 1.65;
-        border: 1px solid rgba(255,255,255,0.08);
-        white-space: pre-wrap;
-        word-break: break-word;
-    }
-    .msg-bot-row {
-        display: flex;
-        gap: 0.75rem;
-        align-items: flex-start;
-        margin: 1rem 0 0.3rem;
-    }
-    .bot-avatar {
-        width: 36px;
-        height: 36px;
-        border-radius: 12px;
-        background: linear-gradient(135deg, #0f4c5c 0%, #116579 100%);
-        color: #fff;
-        display: flex;
-        align-items: center;
-        justify-content: center;
-        font-size: 1rem;
-        flex-shrink: 0;
-        box-shadow: var(--shadow);
+    /* Override File Uploader to look like the dark Teal Box */
+    [data-testid="stFileUploader"] {
+        background-color: #064052;
+        border-radius: 16px;
+        padding: 1rem;
+        margin-top: 1rem;
+        border: none;
     }
-    .msg-bot {
-        max-width: min(820px, 72%);
-        background: #eef2f6;
-        color: #0f172a;
-        border: 1px solid rgba(15, 23, 42, 0.06);
-        border-radius: 8px 22px 22px 22px;
-        padding: 1rem 1.1rem;
-        font-size: 0.98rem;
-        line-height: 1.7;
-        white-space: normal;
-        word-break: break-word;
-        box-shadow: 0 2px 10px rgba(15, 23, 42, 0.03);
+    [data-testid="stFileUploader"] > section {
+        text-align: center;
     }
-    .msg-time {
-        font-size: 0.72rem;
-        color: #94a3b8;
-        margin-top: 0.35rem;
+    [data-testid="stFileUploader"] * {
+        color: #FFFFFF !important;
     }
-    .typing-cursor {
-        opacity: 0.85;
+    [data-testid="stFileUploader"] small {
+        display: none; /* Hide default size limit text */
     }
 
-    /* Source cards */
-    .source-card {
-        background: #fff;
-        border: 1px solid rgba(15, 23, 42, 0.07);
-        border-radius: 16px;
-        padding: 0.9rem 1rem;
-        margin: 0.45rem 0;
-        box-shadow: 0 2px 10px rgba(15, 23, 42, 0.03);
-    }
-    .source-head {
+    /* Recents Section */
+    .recents-title {
+        font-size: 0.7rem;
+        font-weight: 600;
+        color: #9CA3AF;
+        letter-spacing: 1px;
+        margin-top: 30px;
+        margin-bottom: 10px;
+    }
+    .recent-item {
+        color: #4B5563;
+        font-size: 0.85rem;
+        padding: 8px 0;
         display: flex;
-        justify-content: space-between;
         align-items: center;
-        margin-bottom: 0.35rem;
-        font-size: 0.8rem;
+        gap: 10px;
+        cursor: pointer;
+    }
+    
+    /* Top Bar Indicator */
+    .session-badge-container {
+        margin-bottom: 40px;
+    }
+    .session-title {
+        color: #064052;
+        font-weight: 600;
+        font-size: 0.9rem;
+    }
+    .badge-active {
+        background-color: #D1E5E8;
+        color: #064052;
+        padding: 2px 8px;
+        border-radius: 12px;
+        font-size: 0.65rem;
         font-weight: 700;
+        margin-left: 8px;
+        vertical-align: middle;
     }
-    .source-title { color: #0f4c5c; }
-    .source-score {
-        color: #64748b;
-        background: #f1f5f9;
-        padding: 0.1rem 0.45rem;
-        border-radius: 999px;
+
+    /* Empty State Headers */
+    .welcome-title {
+        text-align: center;
+        color: #064052;
+        font-size: 2.2rem;
         font-weight: 700;
+        margin-top: 2rem;
+        margin-bottom: 1rem;
     }
-    .source-file {
-        font-size: 0.78rem;
-        color: #64748b;
-        margin-bottom: 0.4rem;
-    }
-    .source-snippet {
-        font-size: 0.86rem;
-        line-height: 1.55;
-        color: #334155;
-        border-left: 2px solid rgba(15, 76, 92, 0.18);
-        padding-left: 0.7rem;
+    .welcome-subtitle {
+        text-align: center;
+        color: #4B5563;
+        font-size: 1rem;
+        margin-bottom: 3rem;
+        line-height: 1.5;
     }
 
-    /* Bottom spacing so chat input doesn't overlap content */
-    .bottom-space {
-        height: 1.75rem;
+    /* Chat Messages */
+    [data-testid="stChatMessage"] {
+        border-radius: 16px;
+        padding: 1.25rem;
+        margin-bottom: 1.5rem;
+        max-width: 85%;
+    }
+    
+    /* Assistant Bubble */
+    [data-testid="stChatMessage"]:has([data-testid="assistantAvatar"]) {
+        background-color: #EAEAEA;
+        color: #111827;
+        margin-right: auto;
+    }
+    
+    /* User Bubble */
+    [data-testid="stChatMessage"]:has([data-testid="userAvatar"]) {
+        background-color: #064052;
+        color: white;
+        margin-left: auto;
+    }
+    [data-testid="stChatMessage"]:has([data-testid="userAvatar"]) p {
+        color: white !important;
+    }
+    
+    /* Hide User Avatar to match clean look */
+    [data-testid="stChatMessage"]:has([data-testid="userAvatar"]) [data-testid="chatAvatarIcon-user"] {
+        display: none;
     }
-    </style>
-    """,
-    unsafe_allow_html=True,
-)
-
-
-# ============================================================================
-# CONFIG
-# ============================================================================
-API_URL = "http://127.0.0.1:8000"
-API_TIMEOUT = 60
-
 
-# ============================================================================
-# SESSION STATE
-# ============================================================================
-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 2024",
-        "Research Summary",
-    ]
-if "ingested_docs" not in st.session_state:
-    st.session_state.ingested_docs = []
-if "show_sources" not in st.session_state:
-    st.session_state.show_sources = True
+    /* Chat Input */
+    [data-testid="stChatInput"] {
+        background-color: #FFFFFF;
+        border: 1px solid #E5E7EB;
+        border-radius: 30px;
+        padding-left: 10px;
+        padding-right: 10px;
+        box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.05);
+    }
+</style>
+"""
+st.markdown(custom_css, unsafe_allow_html=True)
 
 
 # ============================================================================
-# HELPERS
+# API CONFIGURATION & UTILITIES
 # ============================================================================
-def check_backend() -> bool:
-    try:
-        r = requests.get(f"{API_URL}/docs", timeout=5)
-        return r.status_code == 200
-    except Exception:
-        return False
 
+API_URL = "http://127.0.0.1:8000"
+API_TIMEOUT = 30 
 
-def api_call(method: str, endpoint: str, **kwargs):
+def safe_api_call(method, endpoint, **kwargs):
     try:
-        kwargs.setdefault("timeout", API_TIMEOUT)
         url = f"{API_URL}{endpoint}"
-        fn = requests.get if method.upper() == "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: {e}"
-
-
-def now_time() -> str:
-    return datetime.now().strftime("%-I:%M %p")
-
-
-def esc(text: str) -> str:
-    return html.escape(text or "")
-
-
-def render_text_block(text: str) -> str:
-    """Convert plain text to safe HTML for display in a bubble."""
-    safe = esc(text).replace("\n", "<br>")
-    return safe
-
-
-def render_assistant_text(text: str) -> str:
-    """Render assistant content with lightweight bullet styling."""
-    lines = (text or "").splitlines()
-    out = []
-    for raw in lines:
-        line = raw.strip()
-        if not line:
-            continue
-
-        if line.startswith(("- ", "• ", "* ")):
-            item = esc(line[2:])
-            if ":" in item:
-                key, val = item.split(":", 1)
-                out.append(
-                    f"<div style='display:flex;gap:0.6rem;align-items:flex-start;margin:0.35rem 0;'>"
-                    f"<span style='color:#0f4c5c;font-weight:900;line-height:1.2;'>✓</span>"
-                    f"<span><strong>{key.strip()}:</strong> {val.strip()}</span>"
-                    f"</div>"
-                )
-            else:
-                out.append(
-                    f"<div style='display:flex;gap:0.6rem;align-items:flex-start;margin:0.35rem 0;'>"
-                    f"<span style='color:#0f4c5c;font-weight:900;line-height:1.2;'>✓</span>"
-                    f"<span>{item}</span>"
-                    f"</div>"
-                )
+        kwargs.setdefault('timeout', API_TIMEOUT)
+        
+        if method == "GET":
+            response = requests.get(url, **kwargs)
+        elif method == "POST":
+            response = requests.post(url, **kwargs)
         else:
-            out.append(f"<div style='margin:0.22rem 0;'>{esc(line)}</div>")
-    return "".join(out)
-
-
-# ============================================================================
-# BOOTSTRAP
-# ============================================================================
-if not st.session_state.backend_checked:
-    with st.spinner("Starting up..."):
-        time.sleep(1.2)
-        st.session_state.backend_ok = check_backend()
-    st.session_state.backend_checked = True
-
-try:
-    resp, _ = api_call("GET", "/sources")
-    if resp and resp.status_code == 200:
-        st.session_state.ingested_docs = resp.json().get("documents", [])
-except Exception:
-    pass
+            return None, "Invalid method"
+        return response, None
+    except Exception as e:
+        return None, f"❌ Error: {str(e)}"
 
+# Initialize Session States
+if "messages" not in st.session_state:
+    st.session_state.messages = []
 
 # ============================================================================
-# SIDEBAR
+# UI: SIDEBAR
 # ============================================================================
 with st.sidebar:
-    st.markdown(
-        """
-        <div class="sidebar-brand">
-            <div class="brand-row">
-                <div class="brand-icon">🧠</div>
-                <div>
-                    <div class="brand-title">DocuMind AI</div>
-                    <div class="brand-subtitle">Digital Curator</div>
-                </div>
+    # Top Branding
+    st.markdown("""
+        <div class="sidebar-header">
+            <div class="logo-icon">🧠</div>
+            <div>
+                <div class="brand-title">DocuMind AI</div>
+                <div class="brand-subtitle">Digital Curator</div>
             </div>
         </div>
-        """,
-        unsafe_allow_html=True,
-    )
-
-    st.markdown('<div class="sidebar-section">', unsafe_allow_html=True)
-    if st.button("＋  New Chat", key="new_chat"):
+    """, unsafe_allow_html=True)
+    
+    # New Chat
+    if st.button("＋ New Chat", use_container_width=True):
         st.session_state.messages = []
         st.rerun()
-    st.markdown("</div>", unsafe_allow_html=True)
+    
+    # Upload Area
+    uploaded_file = st.file_uploader("Upload", type=["txt", "pdf", "docx"], label_visibility="collapsed")
+    st.markdown("""
+        <div style="text-align: center; margin-top: -20px; margin-bottom: 20px;">
+            <span style="color: #8BA8B0; font-size: 0.65rem; font-weight: 600; letter-spacing: 0.5px;">200MB LIMIT</span>
+        </div>
+    """, unsafe_allow_html=True)
 
-    st.markdown('<div class="sidebar-section" style="padding-top:0;">', unsafe_allow_html=True)
-    st.markdown('<div class="sidebar-heading">Upload</div>', unsafe_allow_html=True)
-    uploaded_file = st.file_uploader(
-        "Upload document",
-        type=["txt", "pdf", "docx"],
-        label_visibility="collapsed",
-    )
-    if uploaded_file and st.button("Ingest Document", key="ingest_btn"):
+    if uploaded_file and st.button("Ingest Document", use_container_width=True):
         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("Ingested successfully")
-                st.session_state.ingested_docs.append(uploaded_file.name)
+            response, error = safe_api_call("POST", "/ingest", files=files)
+            if not error and response and response.status_code == 200:
+                st.success("✅ Added to knowledge base!")
             else:
-                detail = "Unknown error"
-                try:
-                    detail = resp.json().get("detail", resp.text)
-                except Exception:
-                    if resp is not None:
-                        detail = resp.text
-                st.error(detail)
-    st.markdown("</div>", unsafe_allow_html=True)
-
-    st.markdown('<div class="sidebar-section">', unsafe_allow_html=True)
-    st.markdown('<div class="sidebar-heading">Recents</div>', unsafe_allow_html=True)
-    recent_items = st.session_state.ingested_docs or st.session_state.chat_sessions
-    for item in recent_items[:6]:
-        label = item if len(item) <= 26 else item[:23] + "..."
-        st.markdown(
-            f"<div class='recents-item'><span class='icon'>🕘</span><span>{esc(label)}</span></div>",
-            unsafe_allow_html=True,
-        )
-    st.markdown("</div>", unsafe_allow_html=True)
+                st.error("❌ Upload failed.")
+    
+    # Recents 
+    st.markdown("<div class='recents-title'>RECENTS</div>", unsafe_allow_html=True)
+    
+    response, error = safe_api_call("GET", "/sources")
+    if not error and response and response.status_code == 200:
+        documents = response.json().get("documents", [])
+        for doc in documents[:5]:
+            st.markdown(f"<div class='recent-item'>🕒 {doc[:25]}</div>", unsafe_allow_html=True)
+    else:
+        # Static placeholders matching the mockup
+        st.markdown("<div class='recent-item'>🕒 Q3 Financial Report</div>", unsafe_allow_html=True)
+        st.markdown("<div class='recent-item'>🕒 Product Roadmap 2...</div>", unsafe_allow_html=True)
+        st.markdown("<div class='recent-item'>🕒 Research Summary</div>", unsafe_allow_html=True)
 
 
 # ============================================================================
-# MAIN UI
+# UI: MAIN CONTENT
 # ============================================================================
-st.markdown(
-    """
-    <div class="session-row">
-        <div class="session-label">CURRENT SESSION</div>
-        <div class="session-badge">ACTIVE</div>
-    </div>
-    <hr>
-    """,
-    unsafe_allow_html=True,
-)
 
-if not st.session_state.backend_ok:
-    st.markdown(
-        """
-        <div style="margin:1rem 1.25rem 0; padding:0.85rem 1rem; border-radius:16px;
-                    border:1px solid rgba(245, 158, 11, 0.18); background: rgba(245, 158, 11, 0.08);
-                    color:#b45309; font-size:0.9rem;">
-            ⚠️ Backend is starting up. If it still shows after a minute, refresh the page.
-        </div>
-        """,
-        unsafe_allow_html=True,
-    )
+# Top Session Indicator
+st.markdown("""
+    <div class="session-badge-container">
+        <span class="session-title">CURRENT SESSION</span>
+        <span class="badge-active">ACTIVE</span>
+    </div>
+""", unsafe_allow_html=True)
 
-# Empty state
-if not st.session_state.messages:
-    st.markdown(
-        """
-        <div class="hero-wrap">
-            <h1 class="hero-title">How can I assist your research today?</h1>
-            <p class="hero-subtitle">
-                I can analyze documents, summarize findings, or answer specific questions about your archived data.
-            </p>
-        </div>
-        """,
-        unsafe_allow_html=True,
-    )
+# Empty State
+if len(st.session_state.messages) == 0:
+    st.markdown("<div class='welcome-title'>How can I assist your research today?</div>", unsafe_allow_html=True)
+    st.markdown("<div class='welcome-subtitle'>I can analyze documents, summarize findings, or answer<br>specific questions about your archived data.</div>", unsafe_allow_html=True)
 
-# Chat history
+# Chat Container
 for msg in st.session_state.messages:
-    role = msg.get("role", "assistant")
-    content = msg.get("content", "")
-    ts = msg.get("time", "")
-    sources = msg.get("sources", [])
+    is_user = msg["role"] == "user"
+    avatar = None if is_user else "🧠"
+    
+    with st.chat_message(msg["role"], avatar=avatar):
+        st.markdown(msg["content"])
 
-    if role == "user":
-        st.markdown(
-            f"""
-            <div class="msg-user-wrap">
-                <div>
-                    <div class="msg-user">{render_text_block(content)}</div>
-                    <div class="msg-time" style="text-align:right; padding-right:0.4rem;">{esc(ts)}</div>
-                </div>
-            </div>
-            """,
-            unsafe_allow_html=True,
-        )
-    else:
-        st.markdown(
-            f"""
-            <div class="msg-bot-row">
-                <div class="bot-avatar">🧠</div>
-                <div>
-                    <div class="msg-bot">{render_assistant_text(content)}</div>
-                    <div class="msg-time">{esc(ts)}</div>
-                </div>
-            </div>
-            """,
-            unsafe_allow_html=True,
-        )
 
-        if sources and st.session_state.show_sources:
-            with st.expander(f"📚 {len(sources)} source(s) cited"):
-                for i, src in enumerate(sources, start=1):
-                    score = float(src.get("score", 0) or 0)
-                    source_name = esc(str(src.get("source", "Unknown")))
-                    snippet = esc(str(src.get("content", ""))[:400])
-                    if len(str(src.get("content", ""))) > 400:
-                        snippet += "..."
-
-                    st.markdown(
-                        f"""
-                        <div class="source-card">
-                            <div class="source-head">
-                                <span class="source-title">Source {i}</span>
-                                <span class="source-score">{score:.0%} match</span>
-                            </div>
-                            <div class="source-file">📄 {source_name}</div>
-                            <div class="source-snippet">{snippet}</div>
-                        </div>
-                        """,
-                        unsafe_allow_html=True,
-                    )
-
-
-# Chat input
-st.markdown('<div class="bottom-space"></div>', unsafe_allow_html=True)
+# ============================================================================
+# UI: CHAT INPUT
+# ============================================================================
 user_input = st.chat_input("Ask a question...")
 
 if user_input:
-    user_ts = now_time()
-    st.session_state.messages.append({"role": "user", "content": user_input, "time": user_ts})
-
-    # Render user bubble immediately
-    st.markdown(
-        f"""
-        <div class="msg-user-wrap">
-            <div>
-                <div class="msg-user">{render_text_block(user_input)}</div>
-                <div class="msg-time" style="text-align:right; padding-right:0.4rem;">{esc(user_ts)}</div>
-            </div>
-        </div>
-        """,
-        unsafe_allow_html=True,
-    )
-
-    full_response = ""
-    sources: list[dict] = []
-    assistant_placeholder = st.empty()
-
-    # Assistant shell
-    st.markdown(
-        """
-        <div class="msg-bot-row">
-            <div class="bot-avatar">🧠</div>
-            <div style="width:100%;">
-        """,
-        unsafe_allow_html=True,
-    )
-
-    resp, err = api_call("POST", "/query", json={"question": user_input}, stream=True)
-
-    if err:
-        full_response = err
-        assistant_placeholder.markdown(
-            f"""
-            <div class="msg-bot" style="border-color: rgba(239, 68, 68, 0.25); color:#b91c1c; background:#fff5f5;">
-                {esc(err)}
-            </div>
-            """,
-            unsafe_allow_html=True,
+    # 1. Show user message
+    st.session_state.messages.append({"role": "user", "content": user_input})
+    with st.chat_message("user", avatar=None):
+        st.markdown(user_input)
+    
+    # 2. Process and show assistant response
+    with st.chat_message("assistant", avatar="🧠"):
+        placeholder = st.empty()
+        full_response = ""
+        
+        response, error = safe_api_call(
+            "POST", 
+            "/query", 
+            json={"question": user_input}, 
+            stream=True
         )
-    elif resp:
-        try:
-            buf = ""
-            for line in resp.iter_lines():
-                if not line:
-                    continue
-                try:
-                    data = json.loads(line.decode("utf-8"))
-                except json.JSONDecodeError:
-                    continue
-
-                if data.get("type") == "sources":
-                    sources = data.get("data", []) or []
-                elif data.get("type") == "token":
-                    buf += data.get("content", "")
-                    full_response = buf
-                    assistant_placeholder.markdown(
-                        f"""
-                        <div class="msg-bot">{render_assistant_text(buf)}<span class="typing-cursor">▌</span></div>
-                        """,
-                        unsafe_allow_html=True,
-                    )
-
-            assistant_placeholder.markdown(
-                f"""
-                <div class="msg-bot">{render_assistant_text(full_response)}</div>
-                """,
-                unsafe_allow_html=True,
-            )
-        except Exception as e:
-            full_response = f"Error: {e}"
-            assistant_placeholder.markdown(
-                f"""
-                <div class="msg-bot" style="border-color: rgba(239, 68, 68, 0.25); color:#b91c1c; background:#fff5f5;">
-                    {esc(full_response)}
-                </div>
-                """,
-                unsafe_allow_html=True,
-            )
-
-    st.markdown("</div></div>", unsafe_allow_html=True)
-
-    if sources and st.session_state.show_sources:
-        with st.expander(f"📚 {len(sources)} source(s) cited"):
-            for i, src in enumerate(sources, start=1):
-                score = float(src.get("score", 0) or 0)
-                source_name = esc(str(src.get("source", "Unknown")))
-                snippet = esc(str(src.get("content", ""))[:400])
-                if len(str(src.get("content", ""))) > 400:
-                    snippet += "..."
-
-                st.markdown(
-                    f"""
-                    <div class="source-card">
-                        <div class="source-head">
-                            <span class="source-title">Source {i}</span>
-                            <span class="source-score">{score:.0%} match</span>
-                        </div>
-                        <div class="source-file">📄 {source_name}</div>
-                        <div class="source-snippet">{snippet}</div>
-                    </div>
-                    """,
-                    unsafe_allow_html=True,
-                )
-
-    bot_ts = now_time()
-    st.session_state.messages.append(
-        {
-            "role": "assistant",
-            "content": full_response,
-            "sources": sources,
-            "time": bot_ts,
-        }
-    )
-
-
-# ============================================================================
-# FOOTER NOTE
-# ============================================================================
-st.markdown(
-    """
-    <div style="height:1.5rem;"></div>
-    """,
-    unsafe_allow_html=True,
-)
+        
+        if error:
+            placeholder.markdown(error)
+            full_response = error
+        elif response:
+            try:
+                for line in response.iter_lines():
+                    if line:
+                        try:
+                            decoded_line = line.decode('utf-8')
+                            data = json.loads(decoded_line)
+                            if data.get("type") == "token":
+                                full_response += data.get("content", "")
+                                placeholder.markdown(full_response + "▌")
+                        except json.JSONDecodeError:
+                            continue
+                
+                # Render final response (Replace bullets with checkmarks to match mockup)
+                final_text = full_response.replace("- ", "✅ ")
+                placeholder.markdown(final_text)
+                            
+            except Exception as e:
+                placeholder.markdown(f"❌ Error: {str(e)}")
+                full_response = f"Error: {str(e)}"
+        
+    st.session_state.messages.append({"role": "assistant", "content": full_response})