improved UI/UX

app.py

@@ -1,6 +1,7 @@
 """
-RAG Chat Interface for Hugging Face Space — Gradio UI layer.
-All pure logic and backend HTTP lives in utils.py for testability.
+RAG Knowledge Base — Gradio UI layer.
+Two-column layout: documents panel (left) + chat (right).
+All pure logic lives in utils.py.
 """
 import gradio as gr
 import os
@@ -10,14 +11,12 @@ import requests
 from typing import Generator, Tuple
 
 from utils import (
-    BACKEND_URLS,
     MAX_FILE_SIZE,
     load_sample_data,
     get_backend_url,
     format_file_size,
     extract_text_from_file,
     resolve_file_path,
-    check_health,
     fetch_indexed_docs,
     upload_text_to_backend,
     upload_pasted_text,
@@ -27,7 +26,7 @@ from utils import (
 )
 
 
-# ---------- Sample Data ----------
+# ---------- Sample data ----------
 _SAMPLE_DATA = load_sample_data()
 SAMPLE_LABELS: dict    = _SAMPLE_DATA["labels"]
 SAMPLE_DOCUMENTS: dict = {int(k): v for k, v in _SAMPLE_DATA["documents"].items()}
@@ -36,31 +35,32 @@ SAMPLE_QUESTIONS: list = _SAMPLE_DATA["sample_questions"]
 
 # ---------- Sample-specific helpers ----------
 
-def load_selected_samples(selected_labels: list) -> Tuple[bool, str, int, str]:
-    if not selected_labels:
-        return False, None, None, "⚠️ Please select at least one report."
-    years = sorted([SAMPLE_LABELS[lbl] for lbl in selected_labels])
-    parts = []
-    for year in years:
-        separator = "=" * 60
-        parts.append(f"{separator}\nANNUAL REPORT {year}\n{separator}\n{SAMPLE_DOCUMENTS[year]}")
-    combined = "\n\n".join(parts)
-    if len(years) == 1:
-        filename = f"annual_report_{years[0]}.txt"
-    else:
-        filename = f"annual_reports_{'_'.join(str(y) for y in years)}.txt"
-    return upload_text_to_backend(combined, filename)
-
-
-def build_preview_text(selected_labels: list) -> str:
-    if not selected_labels:
-        return ""
-    years = sorted([SAMPLE_LABELS[lbl] for lbl in selected_labels])
-    parts = []
-    for year in years:
-        separator = "=" * 60
-        parts.append(f"{separator}\nANNUAL REPORT {year}\n{separator}\n{SAMPLE_DOCUMENTS[year]}")
-    return "\n\n".join(parts)
+def upload_sample(label: str) -> Tuple[bool, str, int, str]:
+    """Upload ONE sample report as its own document."""
+    if label not in SAMPLE_LABELS:
+        return False, None, None, "⚠️ Unknown sample."
+    year = SAMPLE_LABELS[label]
+    text = SAMPLE_DOCUMENTS[year]
+    filename = f"annual_report_{year}.txt"
+    return upload_text_to_backend(text, filename)
+
+
+def upload_all_samples() -> Tuple[bool, list, str]:
+    """Upload every sample report as a separate document. Returns (any_success, doc_list, status_msg)."""
+    successes = []
+    failures = []
+    for label in SAMPLE_LABELS.keys():
+        ok, doc_id, _cc, _msg = upload_sample(label)
+        if ok:
+            successes.append(label)
+        else:
+            failures.append(label)
+    docs = fetch_indexed_docs()
+    if successes and not failures:
+        return True, docs, f"✅ Indexed all {len(successes)} sample reports."
+    if successes and failures:
+        return True, docs, f"⚠️ Indexed {len(successes)}, failed {len(failures)}."
+    return False, docs, "❌ Failed to index sample reports."
 
 
 def warmup_with_status(start_time: float = None) -> Generator[Tuple[str, bool, str, float], None, None]:
@@ -81,10 +81,10 @@ def warmup_with_status(start_time: float = None) -> Generator[Tuple[str, bool, s
     threading.Thread(target=do_request, daemon=True).start()
 
     STEPS = [
-        (0, "⏳ Initializing GPU model (Qwen 2.5 3B)..."),
-        (3, "⏳ Initializing GPU model (Qwen 2.5 3B)...\n✅ Embedding model ready (BGE-small)"),
-        (6, "⏳ Initializing GPU model (Qwen 2.5 3B)...\n✅ Embedding model ready (BGE-small)\n⏳ Loading index from volume..."),
-        (9, "✅ GPU model loaded (Qwen 2.5 3B)\n✅ Embedding model ready (BGE-small)\n⏳ Loading index from volume..."),
+        (0, "⏳ Booting GPU container..."),
+        (3, "⏳ Loading embedding model (BGE-small)..."),
+        (6, "⏳ Loading LLM weights (Qwen 2.5 3B)..."),
+        (9, "⏳ Restoring index from volume..."),
     ]
 
     while not result["done"]:
@@ -93,50 +93,105 @@ def warmup_with_status(start_time: float = None) -> Generator[Tuple[str, bool, s
         for threshold, msg in STEPS:
             if elapsed >= threshold:
                 step_msg = msg
-        yield f"{step_msg}\n\n*({elapsed:.1f}s elapsed)*", False, f"⏳ Warming up... ({elapsed:.1f}s)", start_time
+        yield f"{step_msg}\n\n*({elapsed:.1f}s)*", False, f"⏳ Warming up... ({elapsed:.1f}s)", start_time
         time.sleep(1)
 
     elapsed = time.time() - start_time
     if result["success"]:
-        final = (
-            "✅ GPU model loaded (Qwen 2.5 3B)\n"
-            "✅ Embedding model ready (BGE-small)\n"
-            "✅ Index volume ready\n\n"
-            f"**✅ Backend ready in {elapsed:.1f}s**"
-        )
-        yield final, True, f"✅ Ready ({elapsed:.1f}s)", start_time
+        yield f"✅ Backend ready in **{elapsed:.1f}s**", True, f"✅ Ready", start_time
     elif result["error"]:
-        yield f"❌ Cannot connect: {result['error']} ({elapsed:.1f}s)", False, f"❌ Failed ({elapsed:.1f}s)", start_time
+        yield f"❌ Cannot connect: {result['error']}", False, f"❌ Failed", start_time
     else:
-        yield f"❌ Backend returned status {result['status_code']} ({elapsed:.1f}s)", False, f"❌ Failed ({elapsed:.1f}s)", start_time
+        yield f"❌ Backend returned {result['status_code']}", False, f"❌ Failed", start_time
+
+
+# ---------- Display helpers ----------
+
+def render_doc_chips(docs: list, active: str = "(all)") -> str:
+    """Render the indexed docs as a markdown chip list."""
+    if not docs:
+        return (
+            "<div style='padding:1em;text-align:center;color:#888;border:1px dashed #ccc;border-radius:6px'>"
+            "📭 <em>No documents yet</em><br>"
+            "<small>Add documents below to get started.</small>"
+            "</div>"
+        )
+    badge = lambda name, hot: (
+        f"<span style='display:inline-block;padding:4px 10px;margin:3px;border-radius:12px;"
+        f"background:{'#2563eb' if hot else '#eef2ff'};"
+        f"color:{'white' if hot else '#3730a3'};"
+        f"font-size:0.85em;border:1px solid {'#2563eb' if hot else '#c7d2fe'}'>"
+        f"📄 {name}</span>"
+    )
+    chips = "".join(badge(d, d == active) for d in docs)
+    count_line = f"<div style='margin-bottom:6px;color:#666;font-size:0.85em'><strong>{len(docs)}</strong> document{'s' if len(docs) != 1 else ''} indexed:</div>"
+    return f"<div>{count_line}{chips}</div>"
+
+
+def scope_options(docs: list) -> list:
+    """Dropdown choices: search-all + per-doc."""
+    return ["🔍 Search across all documents"] + [f"📄 {d}" for d in docs]
+
+
+def scope_value_to_doc_id(value: str) -> str:
+    """Translate dropdown selection back to backend doc_id (or '' for search-all)."""
+    if not value or value.startswith("🔍"):
+        return ""
+    return value.removeprefix("📄 ").strip()
+
+
+# ---------- Custom CSS ----------
+
+CUSTOM_CSS = """
+.gradio-container { max-width: 1280px !important; margin: 0 auto; }
+#header-banner { background: linear-gradient(90deg, #4f46e5 0%, #7c3aed 100%); color:white; padding: 18px 24px; border-radius: 8px; margin-bottom: 14px; }
+#header-banner h1 { margin: 0; font-size: 1.6em; }
+#header-banner .sub { opacity: 0.85; font-size: 0.95em; margin-top: 4px; }
+.doc-panel { background: #fafafa; border: 1px solid #e5e7eb; border-radius: 8px; padding: 14px; }
+.chat-panel { background: white; border: 1px solid #e5e7eb; border-radius: 8px; padding: 14px; }
+.section-title { font-weight: 600; color: #374151; margin-bottom: 8px; }
+.quick-q-row button { background: #f3f4f6 !important; border: 1px solid #e5e7eb !important; color: #374151 !important; font-size: 0.85em !important; }
+.quick-q-row button:hover { background: #e5e7eb !important; }
+.status-badge { display: inline-block; padding: 2px 8px; border-radius: 10px; font-size: 0.8em; font-weight: 500; }
+footer { display: none !important; }
+"""
 
 
 # ---------- Main App ----------
 
 def create_app() -> gr.Blocks:
-    with gr.Blocks() as demo:
+    theme = gr.themes.Soft(
+        primary_hue="indigo",
+        secondary_hue="violet",
+        neutral_hue="slate",
+    )
+
+    with gr.Blocks(theme=theme, css=CUSTOM_CSS, title="RAG Knowledge Base") as demo:
         page             = gr.State("warmup")
         document_indexed = gr.State(False)
         indexed_docs     = gr.State([])
-        active_doc       = gr.State(None)
 
-        # ── Warmup page ──────────────────────────────────────────────────────
-        with gr.Column(elem_id="warmup-page") as warmup_page:
-            gr.Markdown(
+        # ─── Warmup ────────────────────────────────────────────────────────
+        with gr.Column(visible=True) as warmup_page:
+            gr.HTML(
                 """
-                <div style="text-align: center; margin-top: 15vh;">
-                    <h1 style="font-size: 3em; margin-bottom: 0.5em;">📚 RAG Chat Interface</h1>
-                    <p style="font-size: 1.2em; color: #666;">Powered by Qwen 2.5 & Modal</p>
-                    <br>
-                    <p style="color: #555;">Click the button below to warm up the backend and get started.</p>
+                <div style='text-align:center;margin-top:12vh;padding:2em'>
+                    <div style='font-size:3em;margin-bottom:0.3em'>📚</div>
+                    <h1 style='font-size:2.2em;margin:0;color:#4f46e5'>RAG Knowledge Base</h1>
+                    <p style='color:#6b7280;font-size:1.1em;margin:0.5em 0 2em'>
+                        Ask natural-language questions across your documents.<br>
+                        Powered by Qwen 2.5 + LlamaIndex on Modal.
+                    </p>
                 </div>
                 """
             )
             with gr.Row():
-                with gr.Column():
-                    warmup_btn = gr.Button("🔥 Warm Up Backend", variant="primary", size="lg")
+                with gr.Column(scale=1): pass
+                with gr.Column(scale=2):
+                    warmup_btn = gr.Button("🚀 Start", variant="primary", size="lg")
+                    warmup_status = gr.Markdown("")
+                with gr.Column(scale=1): pass
 
-            warmup_status = gr.Markdown("")
             ready = gr.State(False)
             start_time_state = gr.State(0.0)
 
@@ -149,280 +204,249 @@ def create_app() -> gr.Blocks:
                 outputs=[warmup_status, ready, warmup_btn, start_time_state],
             )
 
-        # ── Chat page ────────────────────────────────────────────────────────
-        with gr.Column(elem_id="chat-page", visible=False) as chat_page:
-            gr.Markdown(
+        # ─── Main app ──────────────────────────────────────────────────────
+        with gr.Column(visible=False) as chat_page:
+
+            gr.HTML(
                 """
-                <div style="text-align: center; margin-bottom: 1em;">
-                    <h1>📚 RAG Chat Interface</h1>
+                <div id='header-banner'>
+                    <h1>📚 RAG Knowledge Base</h1>
+                    <div class='sub'>Ask questions across multiple documents — answers cite their sources.</div>
                 </div>
                 """
             )
 
-            def respond(message, history, is_indexed, current_doc):
+            with gr.Row():
+                # ─── LEFT: Documents Panel ─────────────────────────────────
+                with gr.Column(scale=2, elem_classes=["doc-panel"]):
+                    gr.Markdown("#### 📁 Documents")
+                    doc_list_md = gr.Markdown(render_doc_chips([]))
+                    refresh_btn = gr.Button("🔄 Refresh list", size="sm")
+
+                    with gr.Accordion("➕ Add documents", open=False):
+                        with gr.Tabs():
+                            # — Tab A: samples —
+                            with gr.Tab("Samples"):
+                                gr.Markdown(
+                                    "<small style='color:#666'>Demo annual reports — index all to try cross-document queries.</small>"
+                                )
+                                load_all_samples_btn = gr.Button(
+                                    "📋 Load all 4 sample reports",
+                                    variant="primary",
+                                    size="sm",
+                                )
+                                samples_status = gr.Markdown("")
+
+                            # — Tab B: paste —
+                            with gr.Tab("Paste"):
+                                paste_input = gr.Textbox(
+                                    label="Text content",
+                                    lines=8,
+                                    placeholder="Paste content here…",
+                                )
+                                paste_filename = gr.Textbox(
+                                    label="Document name",
+                                    value="pasted.txt",
+                                    lines=1,
+                                )
+                                load_paste_btn = gr.Button("Index text", variant="primary", size="sm")
+                                paste_status = gr.Markdown("")
+
+                            # — Tab C: file —
+                            with gr.Tab("File"):
+                                file_input = gr.File(
+                                    label="Drop a .txt, .pdf, or .docx file",
+                                    file_types=[".txt", ".pdf", ".docx"],
+                                    file_count="single",
+                                )
+                                gr.Markdown(
+                                    "<small style='color:#888'>Files are parsed locally; only extracted text is sent to the backend.</small>"
+                                )
+                                upload_status = gr.Markdown("")
+
+                # ─── RIGHT: Chat Panel ─────────────────────────────────────
+                with gr.Column(scale=5, elem_classes=["chat-panel"]):
+                    gr.Markdown("#### 💬 Ask a question")
+
+                    with gr.Row():
+                        scope_dropdown = gr.Dropdown(
+                            label="Search scope",
+                            choices=scope_options([]),
+                            value="🔍 Search across all documents",
+                            interactive=True,
+                            scale=3,
+                        )
+
+                    chatbot = gr.Chatbot(
+                        height=440,
+                        show_label=False,
+                        value=[],
+                        type="messages",
+                        avatar_images=(None, "https://api.dicebear.com/7.x/bottts/svg?seed=rag"),
+                        placeholder="<div style='text-align:center;color:#9ca3af;padding:2em'><div style='font-size:2.5em'>💡</div><strong>Index a document, then ask anything.</strong><br><small>Try: \"Compare revenue across all years\" or \"Who replaced Marcus Webb as CFO?\"</small></div>",
+                    )
+
+                    with gr.Row():
+                        question_input = gr.Textbox(
+                            placeholder="Type your question and press Enter…",
+                            lines=1,
+                            scale=9,
+                            show_label=False,
+                            container=False,
+                        )
+                        send_btn = gr.Button("Send ▶", scale=1, variant="primary", size="sm")
+
+                    gr.Markdown("<small style='color:#888'>Quick questions:</small>")
+                    with gr.Row(elem_classes=["quick-q-row"]):
+                        q_btns = [gr.Button(q, size="sm") for q in SAMPLE_QUESTIONS]
+
+            # ─── Chat handler ──────────────────────────────────────────────
+            def respond(message, history, is_indexed, scope_choice):
                 history = history or []
                 if not message or not message.strip():
                     yield history, message
                     return
                 user_msg = {"role": "user", "content": message}
                 if not is_indexed:
-                    yield history + [user_msg, {"role": "assistant", "content": "⚠️ Please index a document first (see section 1)."}], ""
+                    yield history + [user_msg, {"role": "assistant", "content": "⚠️ No documents indexed yet. Add one in the panel on the left."}], ""
                     return
+
+                doc_id = scope_value_to_doc_id(scope_choice)
                 yield history + [user_msg, {"role": "assistant", "content": "🤔 Thinking..."}], ""
+
                 final_answer = "🤔 Thinking..."
-                for chunk in query_backend(message, current_doc):
+                for chunk in query_backend(message, doc_id):
                     final_answer = chunk
                     yield history + [user_msg, {"role": "assistant", "content": final_answer}], ""
 
-            def _after_index(success, doc_id, char_count, source_text, error_msg=""):
-                """
-                Shared post-index handler.
-                Returns: (indexed, docs, active_doc, index_status, preview_box, dropdown_update)
-                """
-                if success:
-                    docs = fetch_indexed_docs()
-                    display = update_index_display(docs)
-                    preview = source_text or ""
-                    return (
-                        True,
-                        docs,
-                        doc_id,
-                        display,
-                        preview,
-                        gr.update(choices=docs, value=doc_id),
-                    )
-                else:
-                    display = error_msg or "📄 *No document indexed yet*"
-                    return False, gr.update(), gr.update(), display, "", gr.update()
-
-            # ── Section 1: Select files to index ──
-            with gr.Accordion("📥 1. Select files to index", open=True):
-                with gr.Tabs():
-                    with gr.Tab("📋 Sample reports"):
-                        sample_selector = gr.CheckboxGroup(
-                            choices=list(SAMPLE_LABELS.keys()),
-                            value=[],
-                            label="Pick one or more sample annual reports",
-                            interactive=True,
-                        )
-                        load_samples_btn = gr.Button("🚀 Index Selected Reports", variant="primary")
-                        samples_status = gr.Markdown("")
-
-                    with gr.Tab("✏️ Paste text"):
-                        paste_input = gr.Textbox(
-                            label="Paste text here",
-                            lines=10,
-                            max_lines=20,
-                            placeholder="Paste any text content you want to ask questions about…",
-                        )
-                        paste_filename = gr.Textbox(
-                            label="Filename (optional)",
-                            value="pasted.txt",
-                            lines=1,
-                        )
-                        load_paste_btn = gr.Button("🚀 Index Pasted Text", variant="primary")
-                        paste_status = gr.Markdown("")
-
-                    with gr.Tab("📁 Upload file"):
-                        file_input = gr.File(
-                            label="Upload .txt, .pdf, or .docx",
-                            file_types=[".txt", ".pdf", ".docx"],
-                            file_count="single",
-                        )
-                        gr.Markdown(
-                            "<small style='color:#888'>PDF and DOCX are parsed locally; "
-                            "only the extracted text is sent to the backend. "
-                            "Indexing starts automatically once a file is selected.</small>"
-                        )
-                        upload_status = gr.Markdown("")
-
-            # ── Active Document Selector ──
-            with gr.Row():
-                active_doc_dropdown = gr.Dropdown(
-                    label="📂 Active document for queries",
-                    choices=[],
-                    value=None,
-                    interactive=True,
-                    scale=5,
-                )
-                refresh_docs_btn = gr.Button("🔄 Refresh", scale=1, min_width=80)
-
-            # ── Section 2: Preview indexed text ──
-            with gr.Accordion("📖 2. Preview indexed text", open=True):
-                index_status = gr.Markdown("📄 *No documents indexed yet*")
-                preview_box = gr.Textbox(
-                    value="",
-                    lines=18,
-                    max_lines=18,
-                    interactive=False,
-                    label="Content",
-                    placeholder="Index a document in section 1 to see its content here.",
+            # ─── Post-index handler (shared across tabs) ───────────────────
+            def _refresh_panel():
+                docs = fetch_indexed_docs()
+                return (
+                    len(docs) > 0,
+                    docs,
+                    render_doc_chips(docs, scope_value_to_doc_id("🔍 Search across all documents") or "(all)"),
+                    gr.update(choices=scope_options(docs), value="🔍 Search across all documents"),
                 )
 
-            # ── Section 3: Chat ──
-            with gr.Accordion("💬 3. Ask questions", open=True):
-                gr.Markdown("<small style='color:#888'>Quick questions:</small>")
-                with gr.Row():
-                    q_btns = [gr.Button(q, size="sm") for q in SAMPLE_QUESTIONS]
+            REFRESH_OUTPUTS = [document_indexed, indexed_docs, doc_list_md, scope_dropdown]
 
-                chatbot = gr.Chatbot(height=420, show_label=False, value=[])
+            refresh_btn.click(fn=_refresh_panel, outputs=REFRESH_OUTPUTS)
 
-                with gr.Row():
-                    question_input = gr.Textbox(
-                        placeholder="Type your question and press Enter…",
-                        lines=1,
-                        scale=9,
-                        show_label=False,
-                        container=False,
-                    )
-                    send_btn = gr.Button("Send ▶", scale=1, variant="primary")
-
-            INDEX_OUTPUTS = [
-                document_indexed,
-                indexed_docs,
-                active_doc,
-                index_status,
-                preview_box,
-                active_doc_dropdown,
-            ]
-
-            SAMPLES_LABEL = "🚀 Index Selected Reports"
-            PASTE_LABEL   = "🚀 Index Pasted Text"
-
-            def _busy(label):
-                return gr.update(value=label, interactive=False)
-
-            def _ready_btn(label):
-                return gr.update(value=label, interactive=True)
-
-            def _samples_action(selected_labels):
-                yield _busy("📖 Preparing reports..."), "📖 Preparing reports…", gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update()
-                if not selected_labels:
-                    result = _after_index(False, None, None, "", "⚠️ Please select at least one report.")
-                    yield (_ready_btn(SAMPLES_LABEL), "⚠️ Please select at least one report.") + result
-                    return
-                yield _busy("📤 Uploading & indexing..."), "📤 Uploading & indexing…", gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update()
-                success, doc_id, char_count, msg = load_selected_samples(selected_labels)
-                source_text = build_preview_text(selected_labels) if success else ""
-                result = _after_index(success, doc_id, char_count, source_text, msg)
-                final_status = "✅ Done." if success else msg
-                yield (_ready_btn(SAMPLES_LABEL), final_status) + result
-
-            load_samples_btn.click(
+            # ─── Tab actions ───────────────────────────────────────────────
+            def _samples_action():
+                yield "📤 Indexing 4 sample reports..."
+                _ok, _docs, msg = upload_all_samples()
+                yield msg
+
+            load_all_samples_btn.click(
                 fn=_samples_action,
-                inputs=[sample_selector],
-                outputs=[load_samples_btn, samples_status] + INDEX_OUTPUTS,
+                outputs=[samples_status],
+            ).then(
+                fn=_refresh_panel,
+                outputs=REFRESH_OUTPUTS,
             )
 
             def _paste_action(text, filename):
-                yield _busy("📖 Preparing text..."), "📖 Preparing text…", gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update()
                 if not text or not text.strip():
-                    result = _after_index(False, None, None, "", "⚠️ Please paste some text first.")
-                    yield (_ready_btn(PASTE_LABEL), "⚠️ Please paste some text first.") + result
+                    yield "⚠️ Please paste some text first."
                     return
-                yield _busy("📤 Uploading & indexing..."), "📤 Uploading & indexing…", gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update()
+                yield "📤 Indexing text..."
                 success, doc_id, char_count, msg = upload_pasted_text(text, filename)
-                result = _after_index(success, doc_id, char_count, text if success else "", msg)
-                final_status = "✅ Done." if success else msg
-                yield (_ready_btn(PASTE_LABEL), final_status) + result
+                yield msg
 
             load_paste_btn.click(
                 fn=_paste_action,
                 inputs=[paste_input, paste_filename],
-                outputs=[load_paste_btn, paste_status] + INDEX_OUTPUTS,
+                outputs=[paste_status],
+            ).then(
+                fn=_refresh_panel,
+                outputs=REFRESH_OUTPUTS,
             )
 
             def _upload_action(file_obj):
-                yield "📖 Parsing file…", gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update()
+                yield "📖 Parsing file…"
                 path = resolve_file_path(file_obj)
                 if not path:
-                    result = _after_index(False, None, None, "", "❌ No file selected.")
-                    yield ("❌ No file selected.",) + result
+                    yield "❌ No file selected."
                     return
                 if not os.path.exists(path):
-                    result = _after_index(False, None, None, "", "❌ File not found.")
-                    yield ("❌ File not found.",) + result
+                    yield "❌ File not found."
                     return
                 file_size = os.path.getsize(path)
                 if file_size > MAX_FILE_SIZE:
-                    msg = f"❌ File too large. Max {format_file_size(MAX_FILE_SIZE)}."
-                    result = _after_index(False, None, None, "", msg)
-                    yield (msg,) + result
+                    yield f"❌ File too large. Max {format_file_size(MAX_FILE_SIZE)}."
                     return
                 text_content, err = extract_text_from_file(path)
                 if err or not text_content or not text_content.strip():
-                    err_msg = err or "❌ File is empty or no extractable text found."
-                    result = _after_index(False, None, None, "", err_msg)
-                    yield (err_msg,) + result
+                    yield err or "❌ File is empty or unreadable."
                     return
-                yield "📤 Uploading & indexing…", gr.update(), gr.update(), gr.update(), gr.update(), gr.update(), gr.update()
+                yield "📤 Uploading & indexing…"
                 base = os.path.basename(path)
                 stem, _ext = os.path.splitext(base)
                 filename = f"{stem}.txt"
                 success, doc_id, char_count, msg = upload_text_to_backend(text_content, filename)
-                result = _after_index(success, doc_id, char_count, text_content if success else "", msg)
-                final_status = f"✅ Indexed **{doc_id}** ({char_count:,} chars)." if success else msg
-                yield (final_status,) + result
+                yield (
+                    f"✅ Indexed **{doc_id}** ({char_count:,} chars)."
+                    if success else msg
+                )
 
             file_input.upload(
                 fn=_upload_action,
                 inputs=[file_input],
-                outputs=[upload_status] + INDEX_OUTPUTS,
+                outputs=[upload_status],
             ).then(
                 fn=lambda: gr.update(value=None),
                 outputs=[file_input],
+            ).then(
+                fn=_refresh_panel,
+                outputs=REFRESH_OUTPUTS,
             )
 
-            active_doc_dropdown.change(
-                fn=lambda d: d,
-                inputs=[active_doc_dropdown],
-                outputs=[active_doc],
-            )
-
-            refresh_docs_btn.click(
-                fn=lambda: gr.update(choices=fetch_indexed_docs()),
-                outputs=[active_doc_dropdown],
-            )
-
+            # ─── Chat wiring ───────────────────────────────────────────────
             for btn, q in zip(q_btns, SAMPLE_QUESTIONS):
                 btn.click(
                     fn=lambda x=q: x,
                     outputs=[question_input],
                 ).then(
                     fn=respond,
-                    inputs=[question_input, chatbot, document_indexed, active_doc],
+                    inputs=[question_input, chatbot, document_indexed, scope_dropdown],
                     outputs=[chatbot, question_input],
                 )
 
             send_btn.click(
                 fn=respond,
-                inputs=[question_input, chatbot, document_indexed, active_doc],
+                inputs=[question_input, chatbot, document_indexed, scope_dropdown],
                 outputs=[chatbot, question_input],
             )
             question_input.submit(
                 fn=respond,
-                inputs=[question_input, chatbot, document_indexed, active_doc],
+                inputs=[question_input, chatbot, document_indexed, scope_dropdown],
                 outputs=[chatbot, question_input],
             )
 
+        # Page switch
         ready.change(
             fn=lambda r: (gr.update(visible=not r), gr.update(visible=r), "chat" if r else "warmup"),
             inputs=[ready],
             outputs=[warmup_page, chat_page, page],
         )
 
-        demo.load(
-            fn=load_index_state,
-            outputs=[document_indexed, indexed_docs, active_doc, index_status],
-        ).then(
-            fn=lambda docs, active: gr.update(choices=docs, value=active),
-            inputs=[indexed_docs, active_doc],
-            outputs=[active_doc_dropdown],
-        )
+        # Load state on first page render
+        def _initial_load():
+            is_indexed, docs, _active, _display = load_index_state()
+            return (
+                is_indexed,
+                docs,
+                render_doc_chips(docs, "(all)"),
+                gr.update(choices=scope_options(docs), value="🔍 Search across all documents"),
+            )
+
+        demo.load(fn=_initial_load, outputs=REFRESH_OUTPUTS)
 
     return demo
 
 
 if __name__ == "__main__":
     app = create_app()
-    app.launch(server_name="0.0.0.0", server_port=7860, theme=gr.themes.Soft())
+    app.launch(server_name="0.0.0.0", server_port=7860)

utils.py

@@ -219,9 +219,16 @@ def query_backend(message: str, doc_id: Optional[str] = None) -> Generator[str,
         response.raise_for_status()
         data = response.json()
         answer = clean_answer(data.get("answer", "⚠️ Received an empty response."))
+
+        # Append sources (from unified-index backend)
+        sources = data.get("sources") or []
+        if sources:
+            src_chips = " ".join(f"`{s}`" for s in sources)
+            answer += f"\n\n<sub style='color:#666'>📎 Sources: {src_chips}</sub>"
+
         elapsed = data.get("elapsed")
         if elapsed is not None:
-            answer += f"\n\n<sub style='color:#888'>⏱️ {elapsed:.1f} seconds</sub>"
+            answer += f"\n\n<sub style='color:#888'>⏱️ {elapsed:.1f}s</sub>"
         yield answer
     except requests.exceptions.Timeout:
         yield "❌ Query timed out. Please try again."