Update app.py

app.py

@@ -1,186 +1,399 @@
 """
-DocuMind - Streamlit Frontend for HuggingFace Spaces
-Updated visually to match the modern "Digital Curator" UI mockup.
+DocuMind AI - Streamlit Frontend
+Modern single-file UI inspired by the provided design
 """
 
 import streamlit as st
+import requests
 import json
-import os
 import time
-import requests
-
-# ============================================================================
-# STREAMLIT CONFIG (Must be first)
-# ============================================================================
 
+# =========================================================
+# PAGE CONFIG
+# =========================================================
 st.set_page_config(
-    page_title="DocuMind AI - Digital Curator",
+    page_title="DocuMind AI",
     page_icon="🧠",
     layout="wide",
-    initial_sidebar_state="expanded"
+    initial_sidebar_state="expanded",
 )
 
-# ============================================================================
-# CUSTOM CSS INJECTION (Theming to match the 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: #FDFDFD;
-    }
-    
-    /* Sidebar Background & Styling */
-    [data-testid="stSidebar"] {
-        background-color: #F4F6F8 !important;
-        border-right: 1px solid #E5E9ED;
-    }
-    
-    /* Hide default header and footer */
-    header {visibility: hidden;}
-    footer {visibility: hidden;}
-    
-    /* Sidebar custom elements */
-    .sidebar-title {
-        font-size: 1.25rem;
-        font-weight: 700;
-        color: #111827;
-        margin-bottom: 2px;
-    }
-    .sidebar-subtitle {
-        font-size: 0.8rem;
-        color: #9CA3AF;
-        margin-bottom: 20px;
-    }
-    .recent-item {
-        color: #4B5563;
-        font-size: 0.9rem;
-        padding: 8px 0;
-        display: flex;
-        align-items: center;
-        gap: 10px;
-        cursor: pointer;
-    }
-    
-    /* Upload Box override */
-    [data-testid="stFileUploader"] > div {
-        background-color: #0A3D4D;
-        color: white;
-        border-radius: 12px;
-        padding: 10px;
-        border: none;
-    }
-    [data-testid="stFileUploader"] small {
-        color: #8BA8B0 !important;
-    }
-    
-    /* Chat bubbles */
-    [data-testid="stChatMessage"] {
-        border-radius: 12px;
-        padding: 15px 20px;
-        margin-bottom: 10px;
-    }
-    /* Assistant Bubble */
-    [data-testid="stChatMessage"]:has([data-testid="assistantAvatar"]) {
-        background-color: #EFEFEF;
-        color: #111827;
-    }
-    /* User Bubble */
-    [data-testid="stChatMessage"]:has([data-testid="userAvatar"]) {
-        background-color: #0A3D4D;
-        color: white;
-    }
-    
-    /* User text color fix for dark background */
-    [data-testid="stChatMessage"]:has([data-testid="userAvatar"]) p {
-        color: white !important;
-    }
-
-    /* Welcome Header */
-    .welcome-header {
-        text-align: center;
-        color: #0A3D4D;
-        font-size: 2.5rem;
-        font-weight: 700;
-        margin-top: 50px;
-        margin-bottom: 10px;
-    }
-    .welcome-sub {
-        text-align: center;
-        color: #4B5563;
-        font-size: 1.1rem;
-        margin-bottom: 50px;
-    }
-
-    /* Top Nav */
-    .top-nav {
-        display: flex;
-        justify-content: space-between;
-        align-items: center;
-        padding: 10px 0 30px 0;
-        border-bottom: 1px solid #E5E9ED;
-        margin-bottom: 30px;
-    }
-    .nav-title {
-        color: #0A3D4D;
-        font-weight: 600;
-        font-size: 1.1rem;
-        display: flex;
-        align-items: center;
-        gap: 10px;
-    }
-    .status-badge {
-        background-color: #D1E5E8;
-        color: #0A3D4D;
-        padding: 2px 8px;
-        border-radius: 12px;
-        font-size: 0.7rem;
-        font-weight: 700;
-    }
-</style>
-"""
-st.markdown(custom_css, unsafe_allow_html=True)
+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;
+        }
 
-# ============================================================================
-# API CONFIGURATION & UTILITIES
-# ============================================================================
+        html, body, [class*="css"] {
+            font-family: "Inter", "Segoe UI", sans-serif;
+        }
 
-API_URL = "http://127.0.0.1:8000"
-API_TIMEOUT = 30  # seconds
+        .stApp {
+            background: var(--main-bg);
+        }
 
-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
+        /* 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;
+        }
+
+        /* Hide default Streamlit stuff that adds clutter */
+        #MainMenu {visibility: hidden;}
+        footer {visibility: hidden;}
+        header {visibility: hidden;}
+
+        /* Expanders */
+        details {
+            border-radius: 14px;
+        }
+
+        /* Recents spacing */
+        .recent-doc {
+            font-size: 1rem;
+            line-height: 1.25;
+            margin-bottom: 0.9rem;
+            opacity: 0.95;
+        }
+
+        .muted-small {
+            font-size: 0.82rem;
+            opacity: 0.78;
+        }
+    </style>
+    """,
+    unsafe_allow_html=True,
+)
+
+# =========================================================
+# 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"
@@ -189,208 +402,233 @@ def safe_api_call(method, endpoint, **kwargs):
     except Exception as e:
         return None, f"❌ Error: {str(e)}"
 
-# Initialize Session States
+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
+# =========================================================
 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
 
-# Health check on load
+if "selected_doc" not in st.session_state:
+    st.session_state.selected_doc = None
+
+# =========================================================
+# BACKEND CHECK
+# =========================================================
 if not st.session_state.backend_checked:
-    st.session_state.backend_healthy = check_backend_health(retries=1) # Fast check
+    with st.spinner("Starting backend..."):
+        time.sleep(1.0)
+        st.session_state.backend_healthy = check_backend_health()
     st.session_state.backend_checked = True
 
-# ============================================================================
-# UI: SIDEBAR
-# ============================================================================
+# =========================================================
+# SIDEBAR
+# =========================================================
 with st.sidebar:
-    # Logo & Title
-    st.markdown("""
-        <div style="display: flex; align-items: center; gap: 10px; margin-bottom: 5px;">
-            <div style="background-color: #0A3D4D; color: white; border-radius: 8px; width: 32px; height: 32px; display: flex; justify-content: center; align-items: center;">🧠</div>
+    st.markdown(
+        """
+        <div class="sidebar-brand">
+            <div class="brand-icon">Ai</div>
             <div>
-                <div class="sidebar-title">DocuMind AI</div>
-                <div class="sidebar-subtitle">Digital Curator</div>
+                <div class="brand-title">DocuMind AI</div>
+                <div class="brand-subtitle">Digital Curator</div>
             </div>
         </div>
-    """, unsafe_allow_html=True)
-    
-    # New Chat Button
-    if st.button("＋ New Chat", use_container_width=True, type="secondary"):
+        """,
+        unsafe_allow_html=True,
+    )
+
+    if st.button("＋ New Chat", use_container_width=True):
         st.session_state.messages = []
+        st.session_state.selected_doc = None
         st.rerun()
-    
-    st.write("") # Spacer
-    
-    # Upload Area
-    uploaded_file = st.file_uploader("Upload Document (200MB LIMIT)", type=["txt", "pdf", "docx"], label_visibility="collapsed")
-    if uploaded_file and st.button("Process Document", use_container_width=True, type="primary"):
-        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.markdown("<br>", unsafe_allow_html=True)
-    
-    # Recents List
-    st.markdown("<div style='font-size: 0.75rem; font-weight: 600; color: #9CA3AF; margin-bottom: 10px; letter-spacing: 1px;'>RECENTS</div>", unsafe_allow_html=True)
-    
-    # Try fetching real recents, otherwise show static UI for design purposes
+
+    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)
+
     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]: # Show top 5
-            st.markdown(f"<div class='recent-item'>🕒 {doc[:20]}{'...' if len(doc)>20 else ''}</div>", unsafe_allow_html=True)
+    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)
     else:
-        # Fallback Mockup UI items if backend is empty/offline
-        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)
-
-    # Footer elements pushing to bottom
-    st.markdown("<div style='margin-top: 150px;'></div>", unsafe_allow_html=True)
-    st.markdown("<div class='recent-item'>⚙️ Settings</div>", unsafe_allow_html=True)
-    st.markdown("<div class='recent-item'>❓ Help</div>", unsafe_allow_html=True)
-    
-    st.divider()
-    
-    # User Profile
-    st.markdown("""
-        <div style="display: flex; align-items: center; gap: 10px;">
-            <div style="background-color: #0A3D4D; border-radius: 50%; width: 35px; height: 35px; display: flex; justify-content: center; align-items: center; font-size: 1.2rem;">👨‍💼</div>
-            <div>
-                <div style="font-weight: 600; font-size: 0.9rem; color: #111827;">Alex Sterling</div>
-                <div style="font-size: 0.75rem; color: #9CA3AF;">Pro Curator</div>
-            </div>
+        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)
+        """,
+        unsafe_allow_html=True,
+    )
+    st.markdown("</div>", unsafe_allow_html=True)
 
-# ============================================================================
-# UI: MAIN INTERFACE
-# ============================================================================
+# =========================================================
+# MAIN HEADER
+# =========================================================
+st.markdown(
+    """
+    <div class="session-badge-wrap">
+        <div class="session-badge">CURRENT SESSION ACTIVE</div>
+    </div>
+    """,
+    unsafe_allow_html=True,
+)
 
-# Top Navigation Bar
-st.markdown("""
-    <div class="top-nav">
-        <div class="nav-title">CURRENT SESSION <span class="status-badge">ACTIVE</span></div>
-        <div style="color: #4B5563; font-size: 1.2rem; cursor: pointer;">🔔 &nbsp; 👤</div>
+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.")
-
-# Welcome Screen (If no messages)
-if len(st.session_state.messages) == 0:
-    st.markdown("<div class='welcome-header'>How can I assist your research today?</div>", unsafe_allow_html=True)
-    st.markdown("<div class='welcome-sub'>I can analyze documents, summarize findings, or answer<br>specific questions about your archived data.</div>", unsafe_allow_html=True)
-
-# Container for chat messages
-chat_container = st.container()
-
-with chat_container:
-    for msg in st.session_state.messages:
-        # Determine avatar based on role
-        is_user = msg["role"] == "user"
-        avatar_emoji = "👤" if is_user else "🧠"
-        
-        # We assign an avatar placeholder. We use CSS above to target user vs assistant bubbles.
-        with st.chat_message(msg["role"], avatar=avatar_emoji):
-            st.markdown(msg["content"])
-            
-            # Display Sources if they exist
-            if not is_user and "sources" in msg and msg.get("sources"):
-                with st.expander("📄 View Citations"):
-                    for idx, src in enumerate(msg["sources"]):
-                        score = src.get('score', 0)
-                        st.caption(f"**{src.get('source', 'Document')}** (Match: {score:.1%})")
-                        content = src.get('content', '')
-                        st.markdown(f"<span style='color: #4B5563; font-size: 0.9rem;'>{content[:300]}...</span>", unsafe_allow_html=True)
-
-# ============================================================================
-# UI: CHAT INPUT & BOTTOM TOOLBAR
-# ============================================================================
-
-# Add some vertical spacing to push the input to the bottom
-st.write("")
+    """,
+    unsafe_allow_html=True,
+)
 
-user_input = st.chat_input("Ask a question...")
+# =========================================================
+# CHAT AREA
+# =========================================================
+st.markdown('<div class="chat-wrap">', unsafe_allow_html=True)
 
-# Mockup Bottom Toolbar (Visual aesthetic matching the design)
-st.markdown("""
-    <div style="display: flex; justify-content: center; gap: 30px; margin-top: -15px; color: #9CA3AF; font-size: 0.75rem; font-weight: 600; letter-spacing: 0.5px;">
-        <span style="cursor: pointer; display: flex; align-items: center; gap: 5px;">✨ SUMMARY MODE</span>
-        <span style="cursor: pointer; display: flex; align-items: center; gap: 5px;">📄 SOURCE CITATIONS</span>
-        <span style="cursor: pointer; display: flex; align-items: center; gap: 5px;">🌐 TRANSLATE</span>
-    </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,
+        )
+
+        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 ""))
 
+st.markdown("</div>", unsafe_allow_html=True)
+
+# =========================================================
+# CHAT INPUT
+# =========================================================
+user_input = st.chat_input("Ask a question...")
 
 if user_input:
-    # 1. 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="👤"):
-            st.markdown(user_input)
-    
-    # 2. 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
-                    
-                    placeholder.markdown(full_response)
-                    
-                    if sources:
-                        with st.expander("📄 View Citations"):
-                            for idx, src in enumerate(sources):
-                                score = src.get('score', 0)
-                                st.caption(f"**{src.get('source', 'Document')}** (Match: {score:.1%})")
-                                content = src.get('content', '')
-                                st.markdown(f"<span style='color: #4B5563; font-size: 0.9rem;'>{content[:300]}...</span>", unsafe_allow_html=True)
-                                
-                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
-    })
+    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()