Update app.py

app.py

@@ -1,242 +1,396 @@
 """
-DocuMind - Modern UI + FastAPI Backend (Single File)
+DocuMind - Streamlit Frontend for HuggingFace Spaces
+Updated visually to match the modern "Digital Curator" UI mockup.
 """
 
 import streamlit as st
-import requests
 import json
+import os
 import time
+import requests
+
+# ============================================================================
+# STREAMLIT CONFIG (Must be first)
+# ============================================================================
 
-# =========================================================
-# CONFIG
-# =========================================================
 st.set_page_config(
-    page_title="DocuMind AI",
+    page_title="DocuMind AI - Digital Curator",
     page_icon="🧠",
     layout="wide",
     initial_sidebar_state="expanded"
 )
 
-API_URL = "http://127.0.0.1:8000"
-API_TIMEOUT = 30
-
-# =========================================================
-# CUSTOM CSS (UI MAGIC)
-# =========================================================
-st.markdown("""
+# ============================================================================
+# 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)
 
-/* Background */
-body {
-    background-color: #f7f9fb;
-}
-
-/* Sidebar */
-section[data-testid="stSidebar"] {
-    background-color: #ffffff;
-    border-right: 1px solid #e6e8eb;
-}
-
-/* Buttons */
-.stButton>button {
-    border-radius: 12px;
-    background-color: #0f4c5c;
-    color: white;
-    border: none;
-    padding: 10px;
-    width: 100%;
-}
-
-/* Header */
-.main-title {
-    text-align: center;
-    font-size: 34px;
-    font-weight: 600;
-    margin-top: 20px;
-}
-
-.sub-text {
-    text-align: center;
-    color: gray;
-    margin-bottom: 30px;
-}
-
-/* Chat bubbles */
-.chat-user {
-    background-color: #0f4c5c;
-    color: white;
-    padding: 12px 16px;
-    border-radius: 18px;
-    margin: 10px 0;
-    max-width: 70%;
-    margin-left: auto;
-}
-
-.chat-assistant {
-    background-color: #f1f3f5;
-    color: black;
-    padding: 12px 16px;
-    border-radius: 18px;
-    margin: 10px 0;
-    max-width: 70%;
-}
-
-/* Input box */
-[data-testid="stChatInput"] {
-    position: fixed;
-    bottom: 20px;
-    left: 300px;
-    right: 20px;
-}
 
-</style>
-""", unsafe_allow_html=True)
+# ============================================================================
+# API CONFIGURATION & UTILITIES
+# ============================================================================
+
+API_URL = "http://127.0.0.1:8000"
+API_TIMEOUT = 30  # seconds
+
+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 FUNCTIONS
-# =========================================================
 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":
-            res = requests.get(url, **kwargs)
+            response = requests.get(url, **kwargs)
+        elif method == "POST":
+            response = requests.post(url, **kwargs)
         else:
-            res = requests.post(url, **kwargs)
-
-        return res, None
+            return None, "Invalid method"
+        
+        return response, 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, str(e)
+        return None, f"❌ Error: {str(e)}"
 
-# =========================================================
-# SESSION STATE
-# =========================================================
+# Initialize Session States
 if "messages" not in st.session_state:
     st.session_state.messages = []
-
-# =========================================================
-# SIDEBAR
-# =========================================================
+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 not st.session_state.backend_checked:
+    st.session_state.backend_healthy = check_backend_health(retries=1) # Fast check
+    st.session_state.backend_checked = True
+
+# ============================================================================
+# UI: SIDEBAR
+# ============================================================================
 with st.sidebar:
-    st.markdown("## 🧠 DocuMind AI")
-    st.caption("Digital Curator")
-
-    if st.button("➕ New Chat"):
+    # 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>
+            <div>
+                <div class="sidebar-title">DocuMind AI</div>
+                <div class="sidebar-subtitle">Digital Curator</div>
+            </div>
+        </div>
+    """, unsafe_allow_html=True)
+    
+    # New Chat Button
+    if st.button("＋ New Chat", use_container_width=True, type="secondary"):
         st.session_state.messages = []
-
-    st.markdown("### 📤 Upload")
-
-    uploaded_file = st.file_uploader("", type=["pdf", "txt", "docx"])
-
-    if uploaded_file and st.button("Upload Document"):
-        files = {"file": (uploaded_file.name, uploaded_file.getvalue())}
-        res, err = safe_api_call("POST", "/ingest", files=files)
-
-        if err:
-            st.error(err)
-        elif res and res.status_code == 200:
-            st.success("Uploaded successfully!")
-        else:
-            st.error("Upload failed")
-
-    st.markdown("---")
-    st.markdown("### 📄 Recents")
-
-    res, err = safe_api_call("GET", "/sources")
-    if res and res.status_code == 200:
-        docs = res.json().get("documents", [])
-        for d in docs:
-            st.markdown(f"• {d}")
-
-# =========================================================
-# MAIN HEADER
-# =========================================================
-st.markdown(
-    '<div class="main-title">How can I assist your research today?</div>',
-    unsafe_allow_html=True
-)
-
-st.markdown(
-    '<div class="sub-text">I can analyze documents, summarize findings, or answer questions.</div>',
-    unsafe_allow_html=True
-)
-
-# =========================================================
-# CHAT DISPLAY
-# =========================================================
-for msg in st.session_state.messages:
-    if msg["role"] == "user":
-        st.markdown(
-            f'<div class="chat-user">{msg["content"]}</div>',
-            unsafe_allow_html=True
-        )
+        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
+    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)
     else:
-        st.markdown(
-            f'<div class="chat-assistant">{msg["content"]}</div>',
-            unsafe_allow_html=True
-        )
+        # 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>
+        </div>
+    """, unsafe_allow_html=True)
+
+# ============================================================================
+# UI: MAIN INTERFACE
+# ============================================================================
+
+# 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>
+    </div>
+""", unsafe_allow_html=True)
 
-        if msg.get("sources"):
-            with st.expander("📚 Sources"):
-                for src in msg["sources"]:
-                    st.caption(src.get("source", "Unknown"))
+# 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("")
+
+user_input = st.chat_input("Ask a question...")
+
+# 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)
 
-# =========================================================
-# CHAT INPUT
-# =========================================================
-user_input = st.chat_input("Ask a question about your documents...")
 
 if user_input:
-    # Show user message instantly
-    st.session_state.messages.append({
-        "role": "user",
-        "content": user_input
-    })
-
-    with st.spinner("Thinking..."):
-        res, err = safe_api_call(
-            "POST",
-            "/query",
-            json={"question": user_input},
-            stream=True
-        )
-
-        full_response = ""
-        sources = []
-
-        if err:
-            full_response = err
-
-        elif res:
-            try:
-                for line in res.iter_lines():
-                    if line:
-                        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 Exception as e:
-                full_response = str(e)
-
-    # Save assistant response
+    # 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
-    })
-
-    st.rerun()
-
-# =========================================================
-# FOOTER
-# =========================================================
-st.markdown("---")
-st.markdown(
-    "<div style='text-align:center; font-size:12px; color:gray;'>DocuMind AI • Enterprise RAG</div>",
-    unsafe_allow_html=True
-)
+    })