Update app.py

app.py

@@ -1,10 +1,15 @@
+"""
+DocuMind - Modern UI + FastAPI Backend (Single File)
+"""
+
 import streamlit as st
 import requests
 import json
-import os
 import time
-from datetime import datetime
 
+# =========================================================
+# CONFIG
+# =========================================================
 st.set_page_config(
     page_title="DocuMind AI",
     page_icon="🧠",
@@ -12,170 +17,226 @@ st.set_page_config(
     initial_sidebar_state="expanded"
 )
 
+API_URL = "http://127.0.0.1:8000"
+API_TIMEOUT = 30
+
+# =========================================================
+# CUSTOM CSS (UI MAGIC)
+# =========================================================
 st.markdown("""
 <style>
-@import url('https://fonts.googleapis.com/css2?family=Syne:wght@400;500;600;700;800&family=DM+Sans:ital,opsz,wght@0,9..40,300;0,9..40,400;0,9..40,500;0,9..40,600;1,9..40,300&display=swap');
-:root {
-    --bg-base:#0d0f14;--bg-card:#141720;--bg-sidebar:#111318;--bg-input:#1a1d27;
-    --bg-user-msg:#1e2235;--bg-bot-msg:#181b26;--accent:#4ade80;--accent-dim:#22c55e22;
-    --accent-muted:#4ade8044;--border:#ffffff0f;--border-hover:#ffffff1a;
-    --text-primary:#f0f2f8;--text-secondary:#8892aa;--text-muted:#50586e;
-    --radius-sm:8px;--radius-md:14px;--radius-lg:20px;
+
+/* Background */
+body {
+    background-color: #f7f9fb;
 }
-html,body,[data-testid="stAppViewContainer"]{font-family:'DM Sans',sans-serif;background:var(--bg-base)!important;color:var(--text-primary)!important}
-#MainMenu,footer,header,[data-testid="stToolbar"],[data-testid="stDecoration"],.stDeployButton{display:none!important}
-[data-testid="stSidebar"]{background:var(--bg-sidebar)!important;border-right:1px solid var(--border)!important;min-width:240px!important;max-width:240px!important}
-[data-testid="stSidebarContent"]{padding:0!important}
-.main .block-container{padding:0!important;max-width:100%!important}
-.stButton>button{background:var(--bg-input)!important;color:var(--text-primary)!important;border:1px solid var(--border-hover)!important;border-radius:var(--radius-sm)!important;padding:6px 14px!important;width:100%!important}
-.stButton>button:hover{background:var(--accent-dim)!important;color:var(--accent)!important}
-[data-testid="stChatInput"]{background:var(--bg-input)!important;border:1px solid var(--border-hover)!important;border-radius:var(--radius-lg)!important}
-[data-testid="stChatInput"] textarea{color:var(--text-primary)!important}
-</style>
-""", unsafe_allow_html=True)
 
-API_URL = "http://127.0.0.1:8000"
-API_TIMEOUT = 60
+/* Sidebar */
+section[data-testid="stSidebar"] {
+    background-color: #ffffff;
+    border-right: 1px solid #e6e8eb;
+}
 
-if "messages" not in st.session_state:
-    st.session_state.messages = []
-if "backend_ok" not in st.session_state:
-    st.session_state.backend_ok = False
-if "backend_checked" not in st.session_state:
-    st.session_state.backend_checked = False
-if "chat_sessions" not in st.session_state:
-    st.session_state.chat_sessions = ["Q3 Financial Report","Product Roadmap 2...","Research Summary"]
-if "mode_summary" not in st.session_state:
-    st.session_state.mode_summary = False
-if "mode_citations" not in st.session_state:
-    st.session_state.mode_citations = True
-if "ingested_docs" not in st.session_state:
-    st.session_state.ingested_docs = []
-
-def check_backend():
-    try:
-        r = requests.get(f"{API_URL}/docs", timeout=5)
-        return r.status_code == 200
-    except:
-        return False
+/* 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)
 
-def api_call(method, endpoint, **kwargs):
+# =========================================================
+# API FUNCTIONS
+# =========================================================
+def safe_api_call(method, endpoint, **kwargs):
     try:
-        kwargs.setdefault("timeout", API_TIMEOUT)
         url = f"{API_URL}{endpoint}"
-        if method == "GET":
-            return requests.get(url, **kwargs), None
-        return requests.post(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}"
+        kwargs.setdefault("timeout", API_TIMEOUT)
 
-def now_time():
-    return datetime.now().strftime("%I:%M %p").lstrip("0")
+        if method == "GET":
+            res = requests.get(url, **kwargs)
+        else:
+            res = requests.post(url, **kwargs)
 
-if not st.session_state.backend_checked:
-    with st.spinner("Starting up..."):
-        time.sleep(1)
-        st.session_state.backend_ok = check_backend()
-    st.session_state.backend_checked = True
+        return res, None
+    except Exception as e:
+        return None, str(e)
 
-try:
-    resp, _ = api_call("GET", "/sources")
-    if resp and resp.status_code == 200:
-        st.session_state.ingested_docs = resp.json().get("documents", [])
-except:
-    pass
+# =========================================================
+# SESSION STATE
+# =========================================================
+if "messages" not in st.session_state:
+    st.session_state.messages = []
 
+# =========================================================
+# SIDEBAR
+# =========================================================
 with st.sidebar:
-    st.markdown("""
-    <div style="padding:22px 20px 16px;border-bottom:1px solid #ffffff0f;">
-        <div style="display:flex;align-items:center;gap:10px;">
-            <div style="width:32px;height:32px;background:linear-gradient(135deg,#4ade80,#22c55e);
-                        border-radius:8px;display:flex;align-items:center;justify-content:center;">🧠</div>
-            <div>
-                <div style="font-weight:700;">DocuMind AI</div>
-                <div style="font-size:0.7rem;color:#50586e;">Digital Curator</div>
-            </div>
-        </div>
-    </div>
-    """, unsafe_allow_html=True)
-
-    if st.button("＋  New Chat"):
+    st.markdown("## 🧠 DocuMind AI")
+    st.caption("Digital Curator")
+
+    if st.button("➕ New Chat"):
         st.session_state.messages = []
-        st.rerun()
 
-    uploaded_file = st.file_uploader("Upload", type=["txt","pdf","docx"])
+    st.markdown("### 📤 Upload")
 
-    if uploaded_file and st.button("Ingest Document"):
+    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())}
-        resp, err = api_call("POST", "/ingest", files=files)
+        res, err = safe_api_call("POST", "/ingest", files=files)
+
         if err:
             st.error(err)
-        elif resp and resp.status_code == 200:
-            st.success("Ingested")
-            st.session_state.ingested_docs.append(uploaded_file.name)
+        elif res and res.status_code == 200:
+            st.success("Uploaded successfully!")
         else:
-            st.error("Failed")
-
-    for doc in st.session_state.ingested_docs[:5] or st.session_state.chat_sessions:
-        st.markdown(f"- {doc}")
+            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("## CURRENT SESSION 🟢 ACTIVE")
+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:
-    with st.chat_message(msg["role"]):
-        st.markdown(msg["content"])
-
-user_input = st.chat_input("Ask a question...")
+    if msg["role"] == "user":
+        st.markdown(
+            f'<div class="chat-user">{msg["content"]}</div>',
+            unsafe_allow_html=True
+        )
+    else:
+        st.markdown(
+            f'<div class="chat-assistant">{msg["content"]}</div>',
+            unsafe_allow_html=True
+        )
+
+        if msg.get("sources"):
+            with st.expander("📚 Sources"):
+                for src in msg["sources"]:
+                    st.caption(src.get("source", "Unknown"))
+
+# =========================================================
+# CHAT INPUT
+# =========================================================
+user_input = st.chat_input("Ask a question about your documents...")
 
 if user_input:
-    st.session_state.messages.append({"role":"user","content":user_input})
-    with st.chat_message("user"):
-        st.markdown(user_input)
+    # Show user message instantly
+    st.session_state.messages.append({
+        "role": "user",
+        "content": user_input
+    })
 
-    full_response = ""
-    sources = []
+    with st.spinner("Thinking..."):
+        res, err = safe_api_call(
+            "POST",
+            "/query",
+            json={"question": user_input},
+            stream=True
+        )
 
-    with st.chat_message("assistant"):
-        placeholder = st.empty()
-        resp, err = api_call("POST", "/query", json={"question": user_input}, stream=True)
+        full_response = ""
+        sources = []
 
         if err:
             full_response = err
-            placeholder.markdown(err)
-        elif resp:
+
+        elif res:
             try:
-                for line in resp.iter_lines(decode_unicode=True):
+                for line in res.iter_lines():
                     if line:
-                        data = json.loads(line)
-                        if data.get("type") == "sources":
-                            sources = data.get("data", [])
-                        elif data.get("type") == "token":
+                        data = json.loads(line.decode("utf-8"))
+
+                        if data.get("type") == "token":
                             full_response += data.get("content", "")
-                            placeholder.markdown(full_response + "▌")
-                placeholder.markdown(full_response)
+
+                        elif data.get("type") == "sources":
+                            sources = data.get("data", [])
+
             except Exception as e:
-                full_response = f"Error: {e}"
-                placeholder.markdown(full_response)
+                full_response = str(e)
 
+    # Save assistant response
     st.session_state.messages.append({
-        "role":"assistant",
-        "content":full_response,
-        "sources":sources
+        "role": "assistant",
+        "content": full_response,
+        "sources": sources
     })
 
-c1,c2,c3,_=st.columns([1,1,1,5])
-with c1:
-    if st.button("Summary Mode"):
-        st.session_state.mode_summary=not st.session_state.mode_summary
-        st.rerun()
-with c2:
-    if st.button("Source Citations"):
-        st.session_state.mode_citations=not st.session_state.mode_citations
-        st.rerun()
-with c3:
-    if st.button("Translate"):
-        st.toast("Coming soon")
+    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
+)