Update app.py

app.py

@@ -1,224 +1,47 @@
 """
 DocuMind - Streamlit Frontend for HuggingFace Spaces
-Minimal & Clean UI implementation based on the provided mockup.
+Simplified version that works with existing FastAPI backend
 """
 
 import streamlit as st
 import json
+import os
 import time
-import requests
 
 # ============================================================================
 # STREAMLIT CONFIG (Must be first)
 # ============================================================================
-
 st.set_page_config(
-    page_title="DocuMind AI",
+    page_title="DocuMind - Enterprise RAG",
     page_icon="🧠",
     layout="wide",
     initial_sidebar_state="expanded"
 )
 
 # ============================================================================
-# CUSTOM CSS (Theming to match the minimal clean mockup)
+# API CONFIGURATION
 # ============================================================================
-custom_css = """
-<style>
-    /* Base Font */
-    @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&display=swap');
-    
-    html, body, [class*="css"] {
-        font-family: 'Inter', sans-serif;
-    }
-    
-    /* Main Backgrounds */
-    .stApp {
-        background-color: #FAFBFC;
-    }
-    [data-testid="stSidebar"] {
-        background-color: #F4F6F8 !important;
-        border-right: 1px solid #E5E7EB;
-    }
-    
-    /* Hide top bar & footer */
-    header {visibility: hidden;}
-    footer {visibility: hidden;}
-    .block-container {
-        padding-top: 2rem !important;
-        max-width: 900px;
-    }
-
-    /* Sidebar Logo Header */
-    .sidebar-header {
-        display: flex;
-        align-items: center;
-        gap: 12px;
-        margin-bottom: 24px;
-        margin-top: -30px;
-    }
-    .logo-icon {
-        background-color: #064052;
-        color: white;
-        width: 32px;
-        height: 32px;
-        border-radius: 8px;
-        display: flex;
-        justify-content: center;
-        align-items: center;
-        font-weight: bold;
-        font-size: 16px;
-    }
-    .brand-title {
-        color: #111827;
-        font-weight: 600;
-        font-size: 1.1rem;
-        line-height: 1.2;
-    }
-    .brand-subtitle {
-        color: #9CA3AF;
-        font-size: 0.75rem;
-    }
-
-    /* Override New Chat Button */
-    div[data-testid="stSidebar"] div[data-testid="stButton"] button {
-        width: 100%;
-        border-radius: 20px;
-        background-color: #EFEFEF;
-        color: #064052;
-        border: none;
-        font-weight: 500;
-        padding: 0.5rem 1rem;
-    }
-    div[data-testid="stSidebar"] div[data-testid="stButton"] button:hover {
-        background-color: #E5E7EB;
-    }
-
-    /* Override File Uploader to look like the dark Teal Box */
-    [data-testid="stFileUploader"] {
-        background-color: #064052;
-        border-radius: 16px;
-        padding: 1rem;
-        margin-top: 1rem;
-        border: none;
-    }
-    [data-testid="stFileUploader"] > section {
-        text-align: center;
-    }
-    [data-testid="stFileUploader"] * {
-        color: #FFFFFF !important;
-    }
-    [data-testid="stFileUploader"] small {
-        display: none; /* Hide default size limit text */
-    }
-
-    /* Recents Section */
-    .recents-title {
-        font-size: 0.7rem;
-        font-weight: 600;
-        color: #9CA3AF;
-        letter-spacing: 1px;
-        margin-top: 30px;
-        margin-bottom: 10px;
-    }
-    .recent-item {
-        color: #4B5563;
-        font-size: 0.85rem;
-        padding: 8px 0;
-        display: flex;
-        align-items: center;
-        gap: 10px;
-        cursor: pointer;
-    }
-    
-    /* Top Bar Indicator */
-    .session-badge-container {
-        margin-bottom: 40px;
-    }
-    .session-title {
-        color: #064052;
-        font-weight: 600;
-        font-size: 0.9rem;
-    }
-    .badge-active {
-        background-color: #D1E5E8;
-        color: #064052;
-        padding: 2px 8px;
-        border-radius: 12px;
-        font-size: 0.65rem;
-        font-weight: 700;
-        margin-left: 8px;
-        vertical-align: middle;
-    }
-
-    /* Empty State Headers */
-    .welcome-title {
-        text-align: center;
-        color: #064052;
-        font-size: 2.2rem;
-        font-weight: 700;
-        margin-top: 2rem;
-        margin-bottom: 1rem;
-    }
-    .welcome-subtitle {
-        text-align: center;
-        color: #4B5563;
-        font-size: 1rem;
-        margin-bottom: 3rem;
-        line-height: 1.5;
-    }
-
-    /* Chat Messages */
-    [data-testid="stChatMessage"] {
-        border-radius: 16px;
-        padding: 1.25rem;
-        margin-bottom: 1.5rem;
-        max-width: 85%;
-    }
-    
-    /* Assistant Bubble */
-    [data-testid="stChatMessage"]:has([data-testid="assistantAvatar"]) {
-        background-color: #EAEAEA;
-        color: #111827;
-        margin-right: auto;
-    }
-    
-    /* User Bubble */
-    [data-testid="stChatMessage"]:has([data-testid="userAvatar"]) {
-        background-color: #064052;
-        color: white;
-        margin-left: auto;
-    }
-    [data-testid="stChatMessage"]:has([data-testid="userAvatar"]) p {
-        color: white !important;
-    }
-    
-    /* Hide User Avatar to match clean look */
-    [data-testid="stChatMessage"]:has([data-testid="userAvatar"]) [data-testid="chatAvatarIcon-user"] {
-        display: none;
-    }
-
-    /* Chat Input */
-    [data-testid="stChatInput"] {
-        background-color: #FFFFFF;
-        border: 1px solid #E5E7EB;
-        border-radius: 30px;
-        padding-left: 10px;
-        padding-right: 10px;
-        box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.05);
-    }
-</style>
-"""
-st.markdown(custom_css, unsafe_allow_html=True)
-
+# Use localhost for internal communication
+API_URL = "http://127.0.0.1:8000"
+API_TIMEOUT = 30  # seconds
 
 # ============================================================================
-# API CONFIGURATION & UTILITIES
+# UTILITY FUNCTIONS
 # ============================================================================
-
-API_URL = "http://127.0.0.1:8000"
-API_TIMEOUT = 30 
+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
 
 def safe_api_call(method, endpoint, **kwargs):
+    """Safely call API with error handling"""
     try:
         url = f"{API_URL}{endpoint}"
         kwargs.setdefault('timeout', API_TIMEOUT)
@@ -229,118 +52,127 @@ def safe_api_call(method, endpoint, **kwargs):
             response = requests.post(url, **kwargs)
         else:
             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, f"❌ Error: {str(e)}"
 
-# Initialize Session States
-if "messages" not in st.session_state:
-    st.session_state.messages = []
-
 # ============================================================================
-# UI: SIDEBAR
+# STREAMLIT FRONTEND UI
 # ============================================================================
+st.title("🧠 DocuMind")
+st.markdown("**Enterprise Document Intelligence Chatbot**")
+st.markdown("---")
+
+# Health check on load
+if "backend_checked" not in st.session_state:
+    st.session_state.backend_checked = False
+    st.session_state.backend_healthy = False
+
+if not st.session_state.backend_checked:
+    with st.spinner("🔄 Starting backend..."):
+        time.sleep(2)  # Give backend time to start
+        st.session_state.backend_healthy = check_backend_health()
+    st.session_state.backend_checked = True
+
+# Show status
+if not st.session_state.backend_healthy:
+    st.warning(
+        "⚠️ Backend is starting up. This may take 30-60 seconds on first load. "
+        "Please refresh the page if you see this message for more than 1 minute."
+    )
+
+# --- Sidebar for Document Upload ---
 with st.sidebar:
-    # Top Branding
-    st.markdown("""
-        <div class="sidebar-header">
-            <div class="logo-icon">🧠</div>
-            <div>
-                <div class="brand-title">DocuMind AI</div>
-                <div class="brand-subtitle">Digital Curator</div>
-            </div>
-        </div>
-    """, unsafe_allow_html=True)
+    st.header("🏢 Document Knowledge Base")
+    st.markdown("Upload PDFs, DOCX, or TXT documents to add them to the system.")
     
-    # New Chat
-    if st.button("＋ New Chat", use_container_width=True):
-        st.session_state.messages = []
-        st.rerun()
+    uploaded_file = st.file_uploader("Upload a new document", type=["txt", "pdf", "docx"])
     
-    # Upload Area
-    uploaded_file = st.file_uploader("Upload", type=["txt", "pdf", "docx"], label_visibility="collapsed")
-    st.markdown("""
-        <div style="text-align: center; margin-top: -20px; margin-bottom: 20px;">
-            <span style="color: #8BA8B0; font-size: 0.65rem; font-weight: 600; letter-spacing: 0.5px;">200MB LIMIT</span>
-        </div>
-    """, unsafe_allow_html=True)
-
-    if uploaded_file and st.button("Ingest Document", use_container_width=True):
-        with st.spinner("Processing..."):
+    if uploaded_file and st.button("Ingest Document", key="ingest_btn"):
+        with st.spinner("Ingesting document (creating chunks & embeddings)..."):
             files = {"file": (uploaded_file.name, uploaded_file.getvalue())}
             response, error = safe_api_call("POST", "/ingest", files=files)
-            if not error and response and response.status_code == 200:
-                st.success("✅ Added to knowledge base!")
+            
+            if error:
+                st.error(error)
+            elif response and response.status_code == 200:
+                st.success(f"✅ {uploaded_file.name} ingested successfully!")
             else:
-                st.error("❌ Upload failed.")
+                st.error(f"❌ Failed to ingest: {response.text if response else 'Unknown error'}")
     
-    # Recents 
-    st.markdown("<div class='recents-title'>RECENTS</div>", unsafe_allow_html=True)
+    st.divider()
     
+    st.subheader("📄 Indexed Documents")
     response, error = safe_api_call("GET", "/sources")
-    if not error and response and response.status_code == 200:
+    
+    if error:
+        st.warning(f"Could not fetch documents: {error}")
+    elif response and response.status_code == 200:
         documents = response.json().get("documents", [])
-        for doc in documents[:5]:
-            st.markdown(f"<div class='recent-item'>🕒 {doc[:25]}</div>", unsafe_allow_html=True)
+        if documents:
+            for doc in documents:
+                st.markdown(f"- `{doc}`")
+        else:
+            st.info("No documents indexed yet.")
     else:
-        # Static placeholders matching the mockup
-        st.markdown("<div class='recent-item'>🕒 Q3 Financial Report</div>", unsafe_allow_html=True)
-        st.markdown("<div class='recent-item'>🕒 Product Roadmap 2...</div>", unsafe_allow_html=True)
-        st.markdown("<div class='recent-item'>🕒 Research Summary</div>", unsafe_allow_html=True)
+        st.info("Backend not ready yet...")
 
+# --- Main Chat Interface ---
+st.subheader("💬 Chat with Your Documents")
 
-# ============================================================================
-# UI: MAIN CONTENT
-# ============================================================================
-
-# Top Session Indicator
-st.markdown("""
-    <div class="session-badge-container">
-        <span class="session-title">CURRENT SESSION</span>
-        <span class="badge-active">ACTIVE</span>
-    </div>
-""", unsafe_allow_html=True)
-
-# Empty State
-if len(st.session_state.messages) == 0:
-    st.markdown("<div class='welcome-title'>How can I assist your research today?</div>", unsafe_allow_html=True)
-    st.markdown("<div class='welcome-subtitle'>I can analyze documents, summarize findings, or answer<br>specific questions about your archived data.</div>", unsafe_allow_html=True)
+# Initialize session state
+if "messages" not in st.session_state:
+    st.session_state.messages = []
 
-# Chat Container
+# Display chat history
 for msg in st.session_state.messages:
-    is_user = msg["role"] == "user"
-    avatar = None if is_user else "🧠"
-    
-    with st.chat_message(msg["role"], avatar=avatar):
+    with st.chat_message(msg["role"]):
         st.markdown(msg["content"])
-
-
-# ============================================================================
-# UI: CHAT INPUT
-# ============================================================================
-user_input = st.chat_input("Ask a question...")
+        if "sources" in msg and msg.get("sources"):
+            with st.expander("📚 Show Sources"):
+                for idx, src in enumerate(msg["sources"]):
+                    score = src.get('score', 0)
+                    st.caption(f"**Source {idx+1}** [Relevance: {score:.2%}]")
+                    st.markdown(f"**From:** `{src.get('source', 'Unknown')}`")
+                    content = src.get('content', '')
+                    if len(content) > 500:
+                        st.markdown(f"> {content[:500]}...")
+                    else:
+                        st.markdown(f"> {content}")
+
+# Chat input
+user_input = st.chat_input("Ask a question about your documents...")
 
 if user_input:
-    # 1. Show user message
+    # Add user message to history
     st.session_state.messages.append({"role": "user", "content": user_input})
-    with st.chat_message("user", avatar=None):
+    
+    # Display user message
+    with st.chat_message("user"):
         st.markdown(user_input)
     
-    # 2. Process and show assistant response
-    with st.chat_message("assistant", avatar="🧠"):
+    # Get assistant response
+    with st.chat_message("assistant"):
         placeholder = st.empty()
         full_response = ""
+        sources = []
         
         response, error = safe_api_call(
-            "POST", 
-            "/query", 
-            json={"question": user_input}, 
+            "POST",
+            "/query",
+            json={"question": user_input},
             stream=True
         )
         
         if error:
-            placeholder.markdown(error)
-            full_response = error
+            error_msg = error
+            st.error(error_msg)
+            full_response = error_msg
         elif response:
             try:
                 for line in response.iter_lines():
@@ -348,18 +180,49 @@ if user_input:
                         try:
                             decoded_line = line.decode('utf-8')
                             data = json.loads(decoded_line)
-                            if data.get("type") == "token":
+                            
+                            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
                 
-                # Render final response (Replace bullets with checkmarks to match mockup)
-                final_text = full_response.replace("- ", "✅ ")
-                placeholder.markdown(final_text)
-                            
+                placeholder.markdown(full_response)
+                
+                # Display sources if available
+                if sources:
+                    with st.expander("📚 Show Sources"):
+                        for idx, src in enumerate(sources):
+                            score = src.get('score', 0)
+                            st.caption(f"**Source {idx+1}** [Relevance: {score:.2%}]")
+                            st.markdown(f"**From:** `{src.get('source', 'Unknown')}`")
+                            content = src.get('content', '')
+                            if len(content) > 500:
+                                st.markdown(f"> {content[:500]}...")
+                            else:
+                                st.markdown(f"> {content}")
+            
             except Exception as e:
-                placeholder.markdown(f"❌ Error: {str(e)}")
-                full_response = f"Error: {str(e)}"
-        
-    st.session_state.messages.append({"role": "assistant", "content": full_response})
+                error_msg = f"❌ Error processing response: {str(e)}"
+                placeholder.markdown(error_msg)
+                st.error(error_msg)
+                full_response = error_msg
+    
+    # Save assistant message
+    st.session_state.messages.append({
+        "role": "assistant",
+        "content": full_response,
+        "sources": sources
+    })
+
+# Footer
+st.divider()
+st.markdown(
+    "<div style='text-align: center; color: var(--color-text-secondary); font-size: 0.85em;'>"
+    "DocuMind - Enterprise RAG Chatbot | "
+    "<a href='https://github.com/TejeshNaiduKona/DocuMind' target='_blank'>GitHub</a>"
+    "</div>",
+    unsafe_allow_html=True
+)