Update app.py

app.py

@@ -1,399 +1,188 @@
 """
-DocuMind AI - Streamlit Frontend
-Modern single-file UI inspired by the provided design
+DocuMind - Streamlit Frontend for HuggingFace Spaces
+Updated visually to match the specific dark-sidebar & green-chat UI mockup.
 """
 
 import streamlit as st
-import requests
 import json
 import time
+import requests
+
+# ============================================================================
+# STREAMLIT CONFIG (Must be first)
+# ============================================================================
 
-# =========================================================
-# PAGE CONFIG
-# =========================================================
 st.set_page_config(
-    page_title="DocuMind AI",
+    page_title="DocuMind AI - Digital Curator",
     page_icon="🧠",
     layout="wide",
-    initial_sidebar_state="expanded",
+    initial_sidebar_state="expanded"
 )
 
-API_URL = "http://127.0.0.1:8000"
-API_TIMEOUT = 30
-
-# =========================================================
-# CSS
-# =========================================================
-st.markdown(
-    """
-    <style>
-        :root {
-            --sidebar-bg: linear-gradient(180deg, #0f2340 0%, #0b1b33 100%);
-            --main-bg: #f7f9fc;
-            --card-bg: #f1f4f8;
-            --text-dark: #1a1f2b;
-            --muted: #6b7280;
-            --line: rgba(255,255,255,0.10);
-            --accent: #1f7a6d;
-            --accent-2: #2da8c9;
-            --user-bubble: linear-gradient(135deg, #0f6a58 0%, #0a4f47 100%);
-            --assistant-bubble: #eef1f5;
-        }
-
-        html, body, [class*="css"] {
-            font-family: "Inter", "Segoe UI", sans-serif;
-        }
-
-        .stApp {
-            background: var(--main-bg);
-        }
-
-        /* Reduce default padding */
-        .block-container {
-            padding-top: 0.9rem;
-            padding-bottom: 1rem;
-            padding-left: 1.2rem;
-            padding-right: 1.2rem;
-        }
-
-        /* Sidebar */
-        section[data-testid="stSidebar"] {
-            background: var(--sidebar-bg);
-            border-right: 1px solid rgba(255,255,255,0.06);
-        }
-
-        section[data-testid="stSidebar"] * {
-            color: white !important;
-        }
-
-        section[data-testid="stSidebar"] .stMarkdown p,
-        section[data-testid="stSidebar"] .stMarkdown li,
-        section[data-testid="stSidebar"] .stCaption {
-            color: rgba(255,255,255,0.84) !important;
-        }
-
-        /* Sidebar buttons */
-        section[data-testid="stSidebar"] .stButton > button {
-            width: 100%;
-            border-radius: 16px;
-            border: 1px solid rgba(255,255,255,0.16);
-            background: linear-gradient(135deg, rgba(50,139,203,0.95) 0%, rgba(60,180,197,0.85) 100%);
-            color: white;
-            font-weight: 600;
-            padding: 0.85rem 1rem;
-            box-shadow: 0 10px 30px rgba(0,0,0,0.18);
-            transition: transform 0.15s ease, box-shadow 0.15s ease;
-        }
-
-        section[data-testid="stSidebar"] .stButton > button:hover {
-            transform: translateY(-1px);
-            box-shadow: 0 14px 34px rgba(0,0,0,0.22);
-        }
-
-        /* Upload button style */
-        section[data-testid="stSidebar"] .upload-card .stButton > button {
-            background: transparent;
-            border: 1px solid rgba(255,255,255,0.18);
-            box-shadow: none;
-        }
-
-        /* Custom sidebar cards */
-        .sidebar-brand {
-            display: flex;
-            align-items: center;
-            gap: 0.8rem;
-            margin: 0.25rem 0 1rem 0;
-        }
-
-        .brand-icon {
-            width: 3rem;
-            height: 3rem;
-            border-radius: 16px;
-            background: linear-gradient(135deg, #1e6a7a 0%, #2c8fd4 100%);
-            display: flex;
-            align-items: center;
-            justify-content: center;
-            font-size: 1.35rem;
-            font-weight: 800;
-            color: white;
-            box-shadow: 0 10px 24px rgba(0,0,0,0.24);
-        }
-
-        .brand-title {
-            font-size: 1.45rem;
-            font-weight: 800;
-            line-height: 1.1;
-        }
-
-        .brand-subtitle {
-            font-size: 0.95rem;
-            opacity: 0.78;
-            margin-top: 0.12rem;
-        }
-
-        .section-label {
-            margin-top: 1.2rem;
-            margin-bottom: 0.65rem;
-            font-size: 0.95rem;
-            letter-spacing: 0.12em;
-            text-transform: uppercase;
-            opacity: 0.78;
-            font-weight: 700;
-        }
-
-        .upload-card {
-            margin-top: 0.35rem;
-            padding: 1rem;
-            border-radius: 18px;
-            border: 1px solid rgba(255,255,255,0.16);
-            background: rgba(255,255,255,0.05);
-            box-shadow: inset 0 1px 0 rgba(255,255,255,0.05);
-        }
-
-        .upload-card .stFileUploader {
-            margin-bottom: 0.5rem;
-        }
-
-        .limit-pill {
-            display: inline-block;
-            padding: 0.38rem 0.7rem;
-            border-radius: 999px;
-            background: rgba(255,255,255,0.12);
-            font-size: 0.78rem;
-            font-weight: 700;
-            letter-spacing: 0.04em;
-            margin-top: 0.6rem;
-        }
-
-        .divider-line {
-            height: 1px;
-            background: var(--line);
-            margin: 1.25rem 0;
-        }
-
-        .recents-item {
-            display: flex;
-            align-items: flex-start;
-            gap: 0.65rem;
-            margin: 0.9rem 0;
-            color: rgba(255,255,255,0.92);
-            font-size: 1rem;
-            line-height: 1.25;
-        }
-
-        .recent-icon {
-            width: 1.35rem;
-            flex: 0 0 1.35rem;
-            opacity: 0.9;
-        }
-
-        .bottom-user {
-            margin-top: 2rem;
-            padding-top: 1rem;
-            border-top: 1px solid rgba(255,255,255,0.10);
-            display: flex;
-            align-items: center;
-            gap: 0.75rem;
-        }
-
-        .avatar {
-            width: 2.6rem;
-            height: 2.6rem;
-            border-radius: 999px;
-            background: linear-gradient(135deg, #f0b56b 0%, #d97d4a 100%);
-            display: flex;
-            align-items: center;
-            justify-content: center;
-            font-size: 1.1rem;
-            color: white;
-            font-weight: 700;
-        }
-
-        .user-name {
-            font-weight: 700;
-            line-height: 1.1;
-        }
-
-        .user-role {
-            font-size: 0.9rem;
-            opacity: 0.76;
-        }
-
-        /* Main top badge */
-        .session-badge-wrap {
-            display: flex;
-            justify-content: flex-end;
-            margin-bottom: 0.8rem;
-        }
-
-        .session-badge {
-            padding: 0.55rem 0.9rem;
-            border-radius: 999px;
-            background: #e7ebf0;
-            color: #445063;
-            font-weight: 700;
-            letter-spacing: 0.04em;
-            font-size: 0.85rem;
-        }
-
-        /* Hero */
-        .hero {
-            text-align: center;
-            margin-top: 1.1rem;
-            margin-bottom: 1.8rem;
-        }
-
-        .hero h1 {
-            font-size: clamp(2.1rem, 4vw, 3.3rem);
-            line-height: 1.08;
-            margin-bottom: 0.7rem;
-            font-weight: 800;
-            color: #0f1722;
-        }
-
-        .hero p {
-            font-size: 1.1rem;
-            color: #556070;
-            line-height: 1.5;
-            max-width: 46rem;
-            margin: 0 auto;
-        }
-
-        /* Chat area spacing */
-        .chat-wrap {
-            max-width: 56rem;
-            margin: 0 auto;
-            padding-top: 0.5rem;
-            padding-bottom: 6rem;
-        }
-
-        /* Message rows */
-        .msg-row {
-            display: flex;
-            margin: 1rem 0;
-            align-items: flex-start;
-            gap: 0.7rem;
-        }
-
-        .msg-row.user {
-            justify-content: flex-end;
-        }
-
-        .msg-row.assistant {
-            justify-content: flex-start;
-        }
-
-        .bubble {
-            max-width: min(42rem, 86%);
-            padding: 1rem 1.15rem;
-            border-radius: 1.25rem;
-            font-size: 1.02rem;
-            line-height: 1.45;
-            box-shadow: 0 10px 25px rgba(15, 23, 34, 0.08);
-            white-space: pre-wrap;
-            word-wrap: break-word;
-        }
-
-        .bubble.user {
-            background: var(--user-bubble);
-            color: white;
-            border-bottom-right-radius: 0.45rem;
-        }
-
-        .bubble.assistant {
-            background: var(--assistant-bubble);
-            color: #202633;
-            border-bottom-left-radius: 0.45rem;
-        }
-
-        .assistant-avatar {
-            width: 2.2rem;
-            height: 2.2rem;
-            border-radius: 999px;
-            background: #22324f;
-            color: white;
-            display: flex;
-            align-items: center;
-            justify-content: center;
-            font-weight: 700;
-            font-size: 0.92rem;
-            flex: 0 0 2.2rem;
-            margin-top: 0.2rem;
-        }
-
-        .timestamp {
-            font-size: 0.82rem;
-            color: #7a8494;
-            margin-top: 0.35rem;
-            margin-left: 0.3rem;
-        }
-
-        /* Chat input area */
-        .stChatInput {
-            position: sticky;
-            bottom: 0.8rem;
-            z-index: 20;
-        }
-
-        /* Make the input look more like the screenshot */
-        [data-testid="stChatInput"] {
-            background: transparent;
-        }
-
-        [data-testid="stChatInput"] > div {
-            border-radius: 18px !important;
-            border: 1px solid rgba(15, 23, 34, 0.08) !important;
-            background: #eef2f6 !important;
-            box-shadow: 0 10px 30px rgba(15, 23, 34, 0.06);
-        }
-
-        [data-testid="stChatInput"] textarea {
-            background: transparent !important;
-            color: #1f2937 !important;
-        }
+# ============================================================================
+# CUSTOM CSS INJECTION (Theming to match the new mockup)
+# ============================================================================
+custom_css = """
+<style>
+    /* Global Font and Backgrounds */
+    @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 App Background */
+    .stApp {
+        background-color: #F8F9FA;
+    }
+    
+    /* Sidebar Background & Styling (Dark Theme) */
+    [data-testid="stSidebar"] {
+        background-color: #172430 !important;
+        border-right: none;
+    }
+    [data-testid="stSidebar"] * {
+        color: #FFFFFF !important;
+    }
+    
+    /* Hide default header and footer */
+    header {visibility: hidden;}
+    footer {visibility: hidden;}
+    
+    /* Sidebar Headers & Text */
+    .sidebar-title {
+        font-size: 1.25rem;
+        font-weight: 700;
+        color: #FFFFFF;
+        margin-bottom: 2px;
+    }
+    .sidebar-subtitle {
+        font-size: 0.8rem;
+        color: #8A9CA8;
+        margin-bottom: 25px;
+    }
+    .recent-header {
+        font-size: 0.75rem;
+        font-weight: 600;
+        color: #8A9CA8;
+        margin-top: 20px;
+        margin-bottom: 15px;
+        letter-spacing: 1px;
+    }
+    .recent-item {
+        color: #D1D8DD;
+        font-size: 0.95rem;
+        padding: 8px 0;
+        display: flex;
+        align-items: center;
+        gap: 10px;
+        cursor: pointer;
+    }
+    
+    /* Upload Box override for Dark Sidebar */
+    [data-testid="stFileUploader"] {
+        border: 1px solid #2B3B4A;
+        border-radius: 12px;
+        padding: 15px 10px;
+        background-color: transparent;
+        text-align: center;
+    }
+    [data-testid="stFileUploader"] > div > div {
+        background-color: transparent;
+    }
+    
+    /* Chat bubbles */
+    [data-testid="stChatMessage"] {
+        border-radius: 12px;
+        padding: 15px 20px;
+        margin-bottom: 10px;
+        max-width: 85%;
+    }
+    /* Assistant Bubble */
+    [data-testid="stChatMessage"]:has([data-testid="assistantAvatar"]) {
+        background-color: #EEF0F2;
+        color: #111827;
+        margin-right: auto;
+    }
+    /* User Bubble */
+    [data-testid="stChatMessage"]:has([data-testid="userAvatar"]) {
+        background-color: #0C4A42;
+        color: white !important;
+        margin-left: auto;
+    }
+    [data-testid="stChatMessage"]:has([data-testid="userAvatar"]) p {
+        color: white !important;
+    }
+
+    /* Welcome Header */
+    .welcome-header {
+        text-align: center;
+        color: #111827;
+        font-size: 2.5rem;
+        font-weight: 700;
+        margin-top: 60px;
+        margin-bottom: 15px;
+    }
+    .welcome-sub {
+        text-align: center;
+        color: #4B5563;
+        font-size: 1.1rem;
+        margin-bottom: 50px;
+        line-height: 1.5;
+    }
+
+    /* Top Nav Badge */
+    .top-badge-container {
+        display: flex;
+        justify-content: flex-end;
+        padding: 10px 20px 0 0;
+    }
+    .status-badge {
+        background-color: #E2E8F0;
+        color: #475569;
+        padding: 6px 12px;
+        border-radius: 6px;
+        font-size: 0.75rem;
+        font-weight: 700;
+        letter-spacing: 0.5px;
+    }
+</style>
+"""
+st.markdown(custom_css, unsafe_allow_html=True)
 
-        /* Hide default Streamlit stuff that adds clutter */
-        #MainMenu {visibility: hidden;}
-        footer {visibility: hidden;}
-        header {visibility: hidden;}
 
-        /* Expanders */
-        details {
-            border-radius: 14px;
-        }
+# ============================================================================
+# API CONFIGURATION & UTILITIES
+# ============================================================================
 
-        /* Recents spacing */
-        .recent-doc {
-            font-size: 1rem;
-            line-height: 1.25;
-            margin-bottom: 0.9rem;
-            opacity: 0.95;
-        }
+API_URL = "http://127.0.0.1:8000"
+API_TIMEOUT = 30  # seconds
 
-        .muted-small {
-            font-size: 0.82rem;
-            opacity: 0.78;
-        }
-    </style>
-    """,
-    unsafe_allow_html=True,
-)
+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
 
-# =========================================================
-# API HELPERS
-# =========================================================
 def safe_api_call(method, endpoint, **kwargs):
+    """Safely call API with error handling"""
     try:
         url = f"{API_URL}{endpoint}"
-        kwargs.setdefault("timeout", API_TIMEOUT)
-
+        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
     except requests.exceptions.Timeout:
         return None, "⏱️ Request timed out"
@@ -402,233 +191,169 @@ def safe_api_call(method, endpoint, **kwargs):
     except Exception as e:
         return None, f"❌ Error: {str(e)}"
 
-def check_backend_health(retries=3):
-    for attempt in range(retries):
-        try:
-            response = requests.get(f"{API_URL}/docs", timeout=5)
-            if response.status_code == 200:
-                return True
-        except Exception:
-            if attempt < retries - 1:
-                time.sleep(1)
-    return False
-
-# =========================================================
-# SESSION STATE
-# =========================================================
+# Initialize Session States
 if "messages" not in st.session_state:
     st.session_state.messages = []
-
 if "backend_checked" not in st.session_state:
     st.session_state.backend_checked = False
     st.session_state.backend_healthy = False
 
-if "selected_doc" not in st.session_state:
-    st.session_state.selected_doc = None
-
-# =========================================================
-# BACKEND CHECK
-# =========================================================
+# Health check on load
 if not st.session_state.backend_checked:
-    with st.spinner("Starting backend..."):
-        time.sleep(1.0)
-        st.session_state.backend_healthy = check_backend_health()
+    st.session_state.backend_healthy = check_backend_health(retries=1)
     st.session_state.backend_checked = True
 
-# =========================================================
-# SIDEBAR
-# =========================================================
+# ============================================================================
+# UI: SIDEBAR
+# ============================================================================
 with st.sidebar:
-    st.markdown(
-        """
-        <div class="sidebar-brand">
-            <div class="brand-icon">Ai</div>
+    # Logo & Title
+    st.markdown("""
+        <div style="display: flex; align-items: center; gap: 12px; margin-bottom: 10px;">
+            <div style="font-size: 2rem;">◧</div>
             <div>
-                <div class="brand-title">DocuMind AI</div>
-                <div class="brand-subtitle">Digital Curator</div>
+                <div class="sidebar-title">DocuMind AI</div>
+                <div class="sidebar-subtitle">Digital Curator</div>
             </div>
         </div>
-        """,
-        unsafe_allow_html=True,
-    )
-
-    if st.button("＋ New Chat", use_container_width=True):
+    """, unsafe_allow_html=True)
+    
+    # New Chat Button
+    if st.button("＋ New Chat", use_container_width=True, type="primary"):
         st.session_state.messages = []
-        st.session_state.selected_doc = None
         st.rerun()
-
-    st.markdown('<div class="section-label">Upload</div>', unsafe_allow_html=True)
-
-    st.markdown('<div class="upload-card">', unsafe_allow_html=True)
-    uploaded_file = st.file_uploader("Upload Document", type=["pdf", "docx", "txt"], label_visibility="collapsed")
-
-    if uploaded_file:
-        st.markdown('<div class="muted-small">Selected file: <b>{}</b></div>'.format(uploaded_file.name), unsafe_allow_html=True)
-
-        if st.button("Upload", use_container_width=True):
-            with st.spinner("Uploading and indexing..."):
-                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: {uploaded_file.name}")
-                else:
-                    st.error(f"Upload failed: {response.text if response else 'Unknown error'}")
-
-    st.markdown('<div class="limit-pill">200MB LIMIT</div>', unsafe_allow_html=True)
-    st.markdown("</div>", unsafe_allow_html=True)
-
-    st.markdown('<div class="divider-line"></div>', unsafe_allow_html=True)
-
-    st.markdown('<div class="section-label">Recents</div>', unsafe_allow_html=True)
-
+    
+    st.write("") # Spacer
+    
+    # Upload Area
+    uploaded_file = st.file_uploader("Upload Document", type=["txt", "pdf", "docx"], label_visibility="collapsed")
+    # Custom badge for the limit underneath the uploader
+    st.markdown("""
+        <div style="text-align: center; margin-top: -15px; margin-bottom: 20px;">
+            <span style="background-color: #2B3B4A; color: #8A9CA8; font-size: 0.7rem; padding: 3px 8px; border-radius: 12px; font-weight: 600;">200MB LIMIT</span>
+        </div>
+    """, unsafe_allow_html=True)
+    
+    if uploaded_file and st.button("Process Document", use_container_width=True):
+        with st.spinner("Analyzing document..."):
+            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("✅ Added to knowledge base!")
+            else:
+                st.error("❌ Upload failed.")
+    
+    st.divider()
+    
+    # Recents List
+    st.markdown("<div class='recent-header'>RECENTS</div>", unsafe_allow_html=True)
+    
+    # Try fetching real recents, otherwise show static UI
     response, error = safe_api_call("GET", "/sources")
-    if error:
-        st.markdown('<div class="muted-small">Could not fetch documents.</div>', unsafe_allow_html=True)
-    elif response and response.status_code == 200:
-        docs = response.json().get("documents", [])
-        if docs:
-            for doc in docs[:6]:
-                doc_label = doc if len(doc) <= 26 else doc[:24] + "..."
-                st.markdown(
-                    f"""
-                    <div class="recents-item">
-                        <div class="recent-icon">🕘</div>
-                        <div class="recent-doc">{doc_label}</div>
-                    </div>
-                    """,
-                    unsafe_allow_html=True,
-                )
-        else:
-            st.markdown('<div class="muted-small">No documents indexed yet.</div>', unsafe_allow_html=True)
+    if not error and response and response.status_code == 200:
+        documents = response.json().get("documents", [])
+        for doc in documents[:5]: # Show top 5
+            st.markdown(f"<div class='recent-item'>🕒 {doc[:20]}{'...' if len(doc)>20 else ''}</div>", unsafe_allow_html=True)
     else:
-        st.markdown('<div class="muted-small">Backend not ready yet.</div>', unsafe_allow_html=True)
-
-    st.markdown('<div class="bottom-user">', unsafe_allow_html=True)
-    st.markdown(
-        """
-        <div class="avatar">👤</div>
-        <div>
-            <div class="user-name">Alex Sterling</div>
-            <div class="user-role">Pro Curator</div>
-        </div>
-        """,
-        unsafe_allow_html=True,
-    )
-    st.markdown("</div>", unsafe_allow_html=True)
-
-# =========================================================
-# MAIN HEADER
-# =========================================================
-st.markdown(
-    """
-    <div class="session-badge-wrap">
-        <div class="session-badge">CURRENT SESSION ACTIVE</div>
+        # Static Mockup UI items based on the image
+        st.markdown("<div class='recent-item'>🕒 Q3 Financial Report</div>", unsafe_allow_html=True)
+        st.markdown("<div class='recent-item'>🕒 Product Roadmap 2024</div>", unsafe_allow_html=True)
+        st.markdown("<div class='recent-item'>🕒 Research Summary</div>", unsafe_allow_html=True)
+
+# ============================================================================
+# UI: MAIN INTERFACE
+# ============================================================================
+
+# Top Navigation Badge
+st.markdown("""
+    <div class="top-badge-container">
+        <span class="status-badge">CURRENT SESSION ACTIVE</span>
     </div>
-    """,
-    unsafe_allow_html=True,
-)
+""", unsafe_allow_html=True)
 
-st.markdown(
-    """
-    <div class="hero">
-        <h1>How can I assist your<br>research today?</h1>
-        <p>
-            I can analyze documents, summarize findings, or answer specific questions
-            about your archived data.
-        </p>
-    </div>
-    """,
-    unsafe_allow_html=True,
-)
+# Show backend warning cleanly if down
+if not st.session_state.backend_healthy:
+    st.warning("⚠️ Backend API is offline. Starting up or unreachable.")
 
-# =========================================================
-# CHAT AREA
-# =========================================================
-st.markdown('<div class="chat-wrap">', unsafe_allow_html=True)
+# Welcome Screen (If no messages)
+if len(st.session_state.messages) == 0:
+    st.markdown("<div class='welcome-header'>How can I assist your<br>research today?</div>", unsafe_allow_html=True)
+    st.markdown("<div class='welcome-sub'>I can analyze documents, summarize findings, or<br>answer specific questions about your archived data.</div>", unsafe_allow_html=True)
 
-for idx, msg in enumerate(st.session_state.messages):
-    if msg["role"] == "user":
-        st.markdown(
-            f"""
-            <div class="msg-row user">
-                <div class="bubble user">{msg["content"]}</div>
-            </div>
-            """,
-            unsafe_allow_html=True,
-        )
-    else:
-        st.markdown(
-            f"""
-            <div class="msg-row assistant">
-                <div class="assistant-avatar">Ai</div>
-                <div>
-                    <div class="bubble assistant">{msg["content"]}</div>
-                    <div class="timestamp">{msg.get("timestamp", "")}</div>
-                </div>
-            </div>
-            """,
-            unsafe_allow_html=True,
-        )
+# Container for chat messages
+chat_container = st.container()
+
+with chat_container:
+    for msg in st.session_state.messages:
+        is_user = msg["role"] == "user"
+        # Hide the default user avatar to match the clean dark bubble look, use logo for assistant
+        avatar_icon = "◧" if not is_user else None 
+        
+        with st.chat_message(msg["role"], avatar=avatar_icon):
+            st.markdown(msg["content"])
 
-        if msg.get("sources"):
-            with st.expander("Source citations"):
-                for i, src in enumerate(msg["sources"], start=1):
-                    st.markdown(f"**Source {i}**")
-                    st.caption(f"From: {src.get('source', 'Unknown')}")
-                    content = src.get("content", "")
-                    st.write(content[:800] + ("..." if len(content) > 800 else ""))
+# ============================================================================
+# UI: CHAT INPUT
+# ============================================================================
 
-st.markdown("</div>", unsafe_allow_html=True)
+st.write("") # Vertical spacing
 
-# =========================================================
-# CHAT INPUT
-# =========================================================
 user_input = st.chat_input("Ask a question...")
 
 if user_input:
-    timestamp = time.strftime("%I:%M %p").lstrip("0")
-    st.session_state.messages.append(
-        {"role": "user", "content": user_input, "timestamp": timestamp}
-    )
-
-    with st.spinner("Thinking..."):
-        response, error = safe_api_call(
-            "POST",
-            "/query",
-            json={"question": user_input},
-            stream=True,
-        )
-
-        full_response = ""
-        sources = []
-
-        if error:
-            full_response = error
-        elif response:
-            try:
-                for line in response.iter_lines():
-                    if not line:
-                        continue
-                    try:
-                        data = json.loads(line.decode("utf-8"))
-                        if data.get("type") == "token":
-                            full_response += data.get("content", "")
-                        elif data.get("type") == "sources":
-                            sources = data.get("data", [])
-                    except json.JSONDecodeError:
-                        continue
-            except Exception as e:
-                full_response = f"❌ Error processing response: {e}"
-
-    st.session_state.messages.append(
-        {
-            "role": "assistant",
-            "content": full_response,
-            "sources": sources,
-            "timestamp": timestamp,
-        }
-    )
-    st.rerun()
+    # Add and display user message immediately
+    st.session_state.messages.append({"role": "user", "content": user_input})
+    with chat_container:
+        with st.chat_message("user", avatar=None):
+            st.markdown(user_input)
+    
+    # Get and display assistant response
+    with chat_container:
+        with st.chat_message("assistant", avatar="◧"):
+            placeholder = st.empty()
+            full_response = ""
+            sources = []
+            
+            response, error = safe_api_call(
+                "POST", 
+                "/query", 
+                json={"question": user_input}, 
+                stream=True
+            )
+            
+            if error:
+                st.error(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") == "sources":
+                                    sources = data.get("data", [])
+                                elif data.get("type") == "token":
+                                    full_response += data.get("content", "")
+                                    placeholder.markdown(full_response + "▌")
+                            except json.JSONDecodeError:
+                                continue
+                    
+                    # Convert bullet points to checkmarks to match mockup design style
+                    styled_response = full_response.replace("- ", "✅ ")
+                    placeholder.markdown(styled_response)
+                                
+                except Exception as e:
+                    error_msg = f"❌ Error processing response: {str(e)}"
+                    placeholder.markdown(error_msg)
+                    full_response = error_msg
+            
+    # Save to state
+    st.session_state.messages.append({
+        "role": "assistant",
+        "content": full_response,
+        "sources": sources
+    })