Spaces:
Paused
Paused
| """ | |
| Modern Gradio UI for Multimodal RAG system. | |
| Chat and document management with clean, responsive layout. | |
| """ | |
| import os | |
| import time | |
| import fnmatch | |
| import requests | |
| import gradio as gr | |
| from pathlib import Path | |
| def _ping_self(url: str): | |
| """Lightweight GET to keep the HF Space from going to sleep.""" | |
| try: | |
| requests.get(url, timeout=10) | |
| except Exception: | |
| pass | |
| def start_keep_alive_scheduler(space_url: str): | |
| from apscheduler.schedulers.background import BackgroundScheduler | |
| from apscheduler.triggers.interval import IntervalTrigger | |
| scheduler = BackgroundScheduler(timezone="UTC", daemon=True) | |
| scheduler.add_job( | |
| _ping_self, | |
| trigger=IntervalTrigger(minutes=5), | |
| args=[space_url], | |
| id="keep_alive_5", | |
| replace_existing=True, | |
| ) | |
| scheduler.start() | |
| return scheduler | |
| API_BASE = os.environ.get("API_BASE", "http://localhost:8000") | |
| def api_get(path: str, timeout: int = 10): | |
| try: | |
| r = requests.get(f"{API_BASE}{path}", timeout=timeout) | |
| r.raise_for_status() | |
| return r.json() | |
| except Exception as e: | |
| return {"error": str(e)} | |
| def api_post(path: str, json=None, files=None, timeout: int = 60, _retries: int = 3): | |
| import requests.exceptions | |
| for attempt in range(_retries): | |
| try: | |
| r = requests.post(f"{API_BASE}{path}", json=json, files=files, timeout=timeout) | |
| r.raise_for_status() | |
| return r.json() | |
| except requests.exceptions.ConnectionError: | |
| if attempt < _retries - 1: | |
| time.sleep(3) | |
| continue | |
| return {"error": "⚠️ Backend not ready yet — please wait a moment and try again."} | |
| except Exception as e: | |
| return {"error": str(e)} | |
| return {"error": "⚠️ Backend unreachable after retries."} | |
| def api_delete(path: str, timeout: int = 10): | |
| try: | |
| r = requests.delete(f"{API_BASE}{path}", timeout=timeout) | |
| r.raise_for_status() | |
| return r.json() | |
| except Exception as e: | |
| return {"error": str(e)} | |
| def get_status(): | |
| data = api_get("/status") | |
| if not data or "error" in data: | |
| # Always return a tuple of the right length | |
| return [], [], "⚠️ API unavailable", "Unknown", "Unknown" | |
| docs = data.get("documents", []) | |
| files = data.get("data_dir_files", []) | |
| chunks = data.get("total_chunks", 0) | |
| model = data.get("model", "unknown") | |
| device = data.get("device", "CPU") | |
| status_msg = ( | |
| f'✅ <span style="color:#facc15;">' | |
| f'<span style="color:#4ade80;font-weight:bold;">{len(docs)}</span>' | |
| f' document(s) indexed | ' | |
| f'<span style="color:#4ade80;font-weight:bold;">{chunks}</span>' | |
| f' chunks</span>' | |
| ) | |
| return docs, files, status_msg, model, device | |
| def _wait_for_backend(timeout: int = 300) -> bool: | |
| """Poll /status until backend responds or timeout (seconds) elapses.""" | |
| import requests | |
| deadline = time.time() + timeout | |
| while time.time() < deadline: | |
| try: | |
| r = requests.get(f"{API_BASE}/status", timeout=3) | |
| if r.status_code == 200: | |
| return True | |
| except Exception: | |
| pass | |
| time.sleep(3) | |
| return False | |
| def _header_html(model, device): | |
| return ( | |
| "# 🧠 <span style='color:#3b82f6;'>Multimodal RAG</span>\n" | |
| "<span style='color:#7c5cfc;font-size:1.1em;'>Chat with your pdf, word, excel, csv, txt, image, chart, and table documents.</span>\n" | |
| f"<br><span style='color:#ff9800;font-weight:bold;'>| LLM: {model} | Device: {device} |</span>" | |
| " <span style='color:#7c5cfc;font-weight:bold;'>Powered by ChromaDB</span>" | |
| ) | |
| def upload_files(files): | |
| """Generator: yields (upload_status, doc_list, status_text, submit_btn, header_md) | |
| after each phase so the UI updates live as each file is uploaded and indexed. | |
| """ | |
| import urllib.parse | |
| _noop = (gr.update(), gr.update(), gr.update(), gr.update(), gr.update()) | |
| def _emit(html, refresh=False): | |
| if refresh: | |
| docs, _, status_msg, model, device = get_status() | |
| return (html, | |
| gr.update(choices=docs or [], value=None), | |
| status_msg, | |
| gr.update(interactive=True), | |
| _header_html(model, device), | |
| gr.update(choices=docs or [])) | |
| return (html, *_noop) | |
| if not files: | |
| yield _emit("<span style='color:#f87171'>No files selected.</span>") | |
| return | |
| if not _wait_for_backend(timeout=300): | |
| yield _emit("<span style='color:#ff4d4f'>❌ Backend did not start in time. Please refresh and try again.</span>") | |
| return | |
| UPLOAD_TIMEOUT = 1800 # 30 min — large PDFs (e.g. 1000-page codes) take time | |
| messages = [] | |
| for file in files: | |
| path = Path(file.name) | |
| # ── Phase 1: uploading to server ────────────────────────────────────────── | |
| preview = messages + [f"<span style='color:#fbbf24'>⚙ {path.name}: uploading…</span>"] | |
| yield _emit('<br>'.join(preview)) | |
| with open(path, "rb") as f: | |
| resp = api_post( | |
| "/documents/upload", | |
| files={"file": (path.name, f, "application/octet-stream")}, | |
| timeout=60, | |
| ) | |
| if "error" in resp: | |
| messages.append(f"<span style='color:#ff4d4f'>❌ {path.name}: {resp['error']}</span>") | |
| yield _emit('<br>'.join(messages)) | |
| continue | |
| # ── Phase 2: indexing (parsing + embedding) ─────────────────────────────── | |
| # IMPORTANT: we MUST yield on every 2-second poll loop iteration, not only | |
| # on phase changes. HF Space (and many nginx proxies) kill SSE connections | |
| # that go idle for >60 s — yielding continuously keeps the stream alive. | |
| encoded = urllib.parse.quote(path.name, safe="") | |
| deadline = time.time() + UPLOAD_TIMEOUT | |
| poll = {} | |
| status = "processing" | |
| phase_str = "parsing document…" | |
| poll_start = time.time() | |
| while time.time() < deadline: | |
| time.sleep(2) | |
| elapsed = int(time.time() - poll_start) | |
| poll = api_get(f"/documents/upload/status?filename={encoded}", timeout=15) | |
| if "error" in poll: | |
| err_msg = str(poll.get("error", "")) | |
| if "404" in err_msg or "No upload job" in err_msg: | |
| break | |
| # Transient network error — keep alive with last known phase | |
| preview = messages + [f"<span style='color:#fbbf24'>⚙ {path.name}: {phase_str} ({elapsed}s)</span>"] | |
| yield _emit('<br>'.join(preview)) | |
| continue | |
| status = poll.get("status", "processing") | |
| if status == "done": | |
| break | |
| if status == "error": | |
| break | |
| # Update phase label (may change between "parsing" and "embedding N chunks") | |
| phase_str = poll.get("phase", phase_str) or phase_str | |
| # Always yield — keeps HF proxy alive and shows live elapsed time | |
| preview = messages + [f"<span style='color:#fbbf24'>⚙ {path.name}: {phase_str} ({elapsed}s)</span>"] | |
| yield _emit('<br>'.join(preview)) | |
| # ── Phase 3: result for this file ───────────────────────────────────────── | |
| if status == "done": | |
| chunks = poll.get("chunks", "?") | |
| messages.append(f"<span style='color:#4ade80'>✓ {path.name}: indexed ({chunks} chunks)</span>") | |
| elif status == "error": | |
| messages.append(f"<span style='color:#ff4d4f'>❌ {path.name}: {poll.get('message', 'indexing failed')}</span>") | |
| else: | |
| messages.append(f"<span style='color:#f87171'>⚠ {path.name}: timed out — check Refresh</span>") | |
| # Refresh doc list immediately after each file so it appears without waiting | |
| yield _emit('<br>'.join(messages), refresh=True) | |
| # Final refresh in case the last yield already covered it | |
| docs, _, status_msg, model, device = get_status() | |
| yield ( | |
| '<br>'.join(messages), | |
| gr.update(choices=docs or [], value=None), | |
| status_msg, | |
| gr.update(interactive=True), | |
| _header_html(model, device), | |
| gr.update(choices=docs or []), | |
| ) | |
| def delete_document(filenames): | |
| import urllib.parse | |
| if not filenames: | |
| return "<span style='color:#f87171'>Please select at least one document.</span>" | |
| messages = [] | |
| for filename in filenames: | |
| resp = api_delete(f"/documents/{urllib.parse.quote(filename, safe='')}") | |
| if "error" in resp: | |
| messages.append(f"<span style='color:#ff4d4f'>❌ {filename}: {resp['error']}</span>") | |
| else: | |
| messages.append(f"<span style='color:#4ade80'>🗑 {filename}: removed</span>") | |
| return '<br>'.join(messages) | |
| def delete_all_embeddings(): | |
| resp = api_delete("/documents") | |
| if "error" in resp: | |
| return f"<span style='color:#ff4d4f'>❌ {resp['error']}</span>" | |
| return f"<span style='color:#4ade80'>🗑 {resp['message']}</span>" | |
| def get_debug_info(): | |
| import json | |
| data = api_get("/debug", timeout=15) | |
| if "error" in data: | |
| return f"<span style='color:#ff4d4f'>⚠️ {data['error']}</span>" | |
| return f"<pre style='color:#e2e8f0;font-size:0.85em;white-space:pre-wrap;background:#111;padding:10px;border-radius:8px'>{json.dumps(data, indent=2)}</pre>" | |
| def add_url(url: str) -> str: | |
| """ | |
| Start a background crawl, then poll until done (or error). | |
| Returns an HTML status string. | |
| """ | |
| url = url.strip() | |
| if not url: | |
| return "<span style='color:#f87171'>Please enter a URL.</span>" | |
| if not url.startswith(("http://", "https://")): | |
| url = "https://" + url | |
| # Kick off the background crawl | |
| resp = api_post("/documents/url", json={"url": url}, timeout=30) | |
| if "error" in resp: | |
| return f"<span style='color:#ff4d4f'>❌ {resp['error']}</span>" | |
| # Poll until the crawl finishes (max 10 min) | |
| import urllib.parse | |
| encoded = urllib.parse.quote(url, safe="") | |
| deadline = time.time() + 600 | |
| while time.time() < deadline: | |
| time.sleep(5) | |
| status = api_get(f"/documents/url/status?url={encoded}", timeout=10) | |
| if "error" in status: | |
| break | |
| if status.get("status") == "done": | |
| return f"<span style='color:#4ade80'>✅ {status['message']}</span>" | |
| if status.get("status") == "error": | |
| return f"<span style='color:#ff4d4f'>❌ {status.get('message', 'Crawl failed.')}</span>" | |
| return "<span style='color:#f87171'>⚠️ Crawl is taking longer than expected — refresh the document list in a moment.</span>" | |
| def refresh_ui(): | |
| docs, files, status_msg, model, _ = get_status() | |
| has_docs = len(docs) > 0 if docs is not None else False | |
| return gr.update(choices=docs or [], value=None), status_msg, gr.update(interactive=has_docs) | |
| def chat_fn(message, history, n_results, temperature, source_filter=None): | |
| """Send query to API, return complete answer (no character streaming).""" | |
| if not message.strip(): | |
| return history, "" | |
| payload = {"question": message, "n_results": n_results, "temperature": temperature} | |
| if source_filter: | |
| payload["source_filter"] = source_filter | |
| resp = api_post("/query", json=payload, timeout=480) | |
| tokens_user = resp.get("tokens_user", 0) | |
| tokens_assistant = resp.get("tokens_assistant", 0) | |
| if "error" in resp: | |
| answer = f"⚠️ {resp['error']}" | |
| else: | |
| answer = resp.get("answer", "I DON'T KNOW") | |
| sources = resp.get("sources", []) | |
| chunks_used = resp.get("chunks_used", 0) | |
| sql_query = resp.get("sql_query", "") | |
| answer_method = resp.get("answer_method", "rag") | |
| if answer_method == "table_query": | |
| method_note = "🗃️ *Answer generated from structured table query (SQL)*" | |
| else: | |
| method_note = f"🔍 *Answer retrieved from document chunks ({chunks_used} chunks)*" | |
| if sql_query: | |
| answer += f"\n\n**SQL query used:**\n```sql\n{sql_query}\n```" | |
| if sources: | |
| answer += f"\n\n{method_note}\n📄 *Sources: {', '.join(sources)}*" | |
| else: | |
| answer += f"\n\n{method_note}" | |
| history = list(history) if history else [] | |
| if history and isinstance(history[0], tuple): | |
| new_hist = [] | |
| for user, bot in history: | |
| if user is not None: | |
| new_hist.append({"role": "user", "content": user}) | |
| if bot is not None: | |
| new_hist.append({"role": "assistant", "content": bot}) | |
| history = new_hist | |
| history.append({"role": "user", "content": message}) | |
| history.append({"role": "assistant", "content": answer}) | |
| return history, {"tokens_user": tokens_user, "tokens_assistant": tokens_assistant} | |
| def clear_memory(): | |
| resp = api_post("/memory/clear") | |
| stats = api_get("/memory/stats") | |
| if "error" in resp: | |
| return f"⚠️ {resp['error']}", [] | |
| msg = f"🧹 Memory cleared. {stats.get('message_count', 0)} messages in memory." | |
| return msg, [] | |
| import re as _re | |
| def _extract_text(content): | |
| """Handle Gradio 6.x content: str or [{'text': '...', 'type': 'text'}, ...]""" | |
| if isinstance(content, str): | |
| return content | |
| if isinstance(content, list): | |
| return "".join( | |
| item["text"] if isinstance(item, dict) and "text" in item else str(item) | |
| for item in content | |
| ) | |
| return str(content) if content else "" | |
| def get_last_answer(history): | |
| if not history: | |
| return "" | |
| last = history[-1] | |
| if isinstance(last, dict): | |
| return _extract_text(last.get("content", "")) if last.get("role") == "assistant" else "" | |
| return _extract_text(last[1]) if last[1] else "" | |
| def format_chat_history(history): | |
| if not history: | |
| return "" | |
| lines = [] | |
| for msg in history: | |
| if isinstance(msg, dict): | |
| role = "User" if msg.get("role") == "user" else "Assistant" | |
| content = _extract_text(msg.get("content", "")) | |
| else: | |
| role = "User" if msg[0] else "Assistant" | |
| content = _extract_text(msg[0] or msg[1] or "") | |
| if content: | |
| lines.append(f"{role}: {content}") | |
| return "\n\n".join(lines) | |
| _EMOJI_RE = _re.compile( | |
| "[\U0001F300-\U0001F9FF\U00002702-\U000027B0\U000024C2-\U0001F251]+", | |
| flags=_re.UNICODE, | |
| ) | |
| def _clean_for_tts(text: str) -> str: | |
| text = _re.sub(r"<details>.*?</details>", "", text, flags=_re.DOTALL) # remove SQL block | |
| text = _re.sub(r"\n\n[🗃🔍📄][^\n]*", "", text) # remove method/sources lines | |
| text = _re.sub(r"\*+([^*]*)\*+", r"\1", text) # remove bold/italic | |
| text = _re.sub(r"`[^`]*`", "", text) # remove inline code | |
| text = _re.sub(r"#+\s", "", text) # remove headers | |
| text = _EMOJI_RE.sub("", text) # remove emojis | |
| return text.strip() | |
| _copy_counter = [0] | |
| def get_chat_for_copy(history): | |
| _copy_counter[0] += 1 | |
| return f"{_copy_counter[0]}\n{format_chat_history(history)}" | |
| _UI_THEME = gr.themes.Soft() | |
| _UI_CSS = """ | |
| .main-col { max-width: 900px; margin: 0 auto; } | |
| /* Chat window — deep cosmic violet */ | |
| .chatbot-wrap { background: #16082e; border-radius: 12px; border: 1px solid #3d1a7a; } | |
| .gradio-container { background: #10131a; } | |
| /* Sample questions — deep midnight navy */ | |
| .sample-q-panel { | |
| background: #071b2e !important; | |
| border-radius: 12px !important; | |
| padding: 10px 10px 6px 10px !important; | |
| border: 1px solid #1a4d7a !important; | |
| } | |
| /* Search filter — deep crimson/burgundy */ | |
| .filter-panel { | |
| background: #1c0810 !important; | |
| border-radius: 12px !important; | |
| padding: 10px !important; | |
| border: 1px solid #6b1a2e !important; | |
| } | |
| .filter-panel .wrap { | |
| max-height: 72px !important; | |
| overflow-y: auto !important; | |
| flex-wrap: wrap !important; | |
| } | |
| .filter-panel .wrap::-webkit-scrollbar { width: 4px; } | |
| .filter-panel .wrap::-webkit-scrollbar-track { background: #1c0810; } | |
| .filter-panel .wrap::-webkit-scrollbar-thumb { background: #6b1a2e; border-radius: 4px; } | |
| #thinking-indicator { display: none !important; } | |
| /* Question box — deep forest emerald */ | |
| #chat-input-wrap { | |
| position: relative !important; | |
| background: #061c10 !important; | |
| border-radius: 10px !important; | |
| border: 1px solid #1a6b3d !important; | |
| padding: 6px !important; | |
| } | |
| #chat-input-wrap textarea { | |
| background: #061c10 !important; | |
| color: #e2e8f0 !important; | |
| } | |
| #ask-btn { width: 52px !important; min-width: 52px !important; max-width: 52px !important; padding: 0 !important; aspect-ratio: 1; } | |
| .sample-q-btn button { | |
| font-size: 0.75em !important; | |
| padding: 5px 10px !important; | |
| border-radius: 14px !important; | |
| border: 1px solid #1a3a5c !important; | |
| background: #0a2540 !important; | |
| color: #94a3b8 !important; | |
| white-space: normal !important; | |
| overflow: hidden; | |
| text-overflow: ellipsis; | |
| min-height: unset !important; | |
| height: auto !important; | |
| width: 100% !important; | |
| text-align: left !important; | |
| } | |
| .sample-q-btn button:hover { | |
| background: #0e3460 !important; | |
| border-color: #3b82f6 !important; | |
| color: #e2e8f0 !important; | |
| } | |
| """ | |
| SAMPLE_QUESTIONS = [ | |
| "How you can help me?", "How many documents are there?", | |
| "What is the first document about?", "Summarize each doc in max 10 bullet points.", | |
| "List top keywords or concepts mentioned.", "What are the names mentioned in docs?", | |
| "List all the links in the documents?", "What problems or challenges are discussed?", | |
| ] | |
| # ─── Build UI ────────────────────────────────────────────────────────────────── | |
| def build_ui(): | |
| with gr.Blocks(title="Multimodal RAG") as demo: | |
| with gr.Row(): | |
| with gr.Column(elem_classes="main-col"): | |
| header_md = gr.Markdown( | |
| "# 🧠 <span style='color:#3b82f6;'>Multimodal RAG</span>\n" | |
| "<span style='color:#7c5cfc;font-size:1.1em;'>Chat with your pdf, word, excel, csv, txt, image, chart, and table documents.</span>\n" | |
| "<br><span style='color:#ff9800;font-weight:bold;'>| Loading... |</span>", | |
| elem_id="header", | |
| ) | |
| with gr.Tabs(selected=0) as tabs: | |
| with gr.TabItem("💬 Chat"): | |
| thinking_indicator = gr.HTML(value="", elem_id="thinking-indicator") | |
| with gr.Row(): | |
| with gr.Column(scale=5): | |
| chatbot = gr.Chatbot( | |
| label="Chat", | |
| height=580, | |
| elem_classes="chatbot-wrap", | |
| autoscroll=False, | |
| ) | |
| with gr.Column(scale=1): | |
| with gr.Column(elem_classes="sample-q-panel"): | |
| sample_q_btns = [] | |
| for _q in SAMPLE_QUESTIONS: | |
| _btn = gr.Button(_q, size="sm", elem_classes="sample-q-btn") | |
| sample_q_btns.append(_btn) | |
| with gr.Column(elem_classes="filter-panel"): | |
| filter_pattern_tb = gr.Textbox( | |
| placeholder="*.png *.pdf img_*", | |
| show_label=False, | |
| lines=1, | |
| max_lines=1, | |
| ) | |
| source_filter_dd = gr.Dropdown( | |
| choices=[], | |
| label="🔍 Search in (leave empty = all docs)", | |
| multiselect=True, | |
| interactive=True, | |
| ) | |
| with gr.Row(elem_id="chat-action-row"): | |
| msg_input = gr.Textbox( | |
| placeholder="", | |
| show_label=False, | |
| scale=7, | |
| lines=1, | |
| max_lines=1, | |
| autofocus=True, | |
| elem_id="chat-input-wrap", | |
| ) | |
| submit_btn = gr.Button("Ask", elem_id="ask-btn", elem_classes="primary-btn", scale=1, min_width=48) | |
| read_btn = gr.Button("Read", elem_id="read-btn", scale=1) | |
| copy_btn = gr.Button("📋 Copy Chat", elem_id="copy-btn", elem_classes=["btn-copy"], scale=1) | |
| clear_chat_btn = gr.Button("🗑 Clear Chat", elem_id="clear-chat-btn", elem_classes=["btn-clear"], scale=1) | |
| with gr.Row(): | |
| n_results_slider = gr.Slider( | |
| minimum=1, maximum=15, value=8, step=1, | |
| label="Top K (context chunks)", | |
| elem_id="topk-slider", | |
| ) | |
| temperature_slider = gr.Slider( | |
| minimum=0.0, maximum=2.0, value=0.0, step=0.1, | |
| label="Temperature (0 = deterministic)", | |
| elem_id="temperature-slider", | |
| ) | |
| token_stats_text = gr.Markdown( | |
| value="<span style='color:#3b82f6; font-weight:600;'>Tokens sent: 0 Tokens received: 0</span>", | |
| elem_id="token-stats", | |
| visible=True, | |
| ) | |
| memory_stats_text = gr.Markdown(value="", elem_id="memory-stats") | |
| with gr.TabItem("📁 Documents"): | |
| status_text = gr.Markdown(value="⏳ Loading...", elem_classes="status-bar") | |
| file_upload = gr.UploadButton( | |
| label="⬆ Upload & Index Files", | |
| file_count="multiple", | |
| file_types=[".pdf", ".png", ".jpg", ".jpeg", ".tiff", ".docx", ".xlsx", ".csv", ".txt"], | |
| elem_classes="primary-btn", | |
| elem_id="file-upload", | |
| ) | |
| upload_status = gr.HTML(value="", elem_id="upload-status") | |
| with gr.Row(): | |
| url_input = gr.Textbox( | |
| placeholder="https://example.com — crawls up to 2 levels deep + linked PDFs", | |
| show_label=False, | |
| scale=5, | |
| elem_id="url-input", | |
| ) | |
| add_url_btn = gr.Button("🌐 Add URL", scale=1, elem_id="add-url-btn") | |
| doc_list = gr.CheckboxGroup( | |
| choices=[], | |
| label="Indexed Documents", | |
| elem_classes="doc-list", | |
| interactive=True, | |
| ) | |
| doc_list_state = gr.State([]) | |
| with gr.Row(): | |
| delete_btn = gr.Button("🗑 Remove selected", elem_id="delete-btn") | |
| delete_all_btn = gr.Button("🗑 Remove ALL", elem_id="delete-all-btn") | |
| refresh_btn = gr.Button("↻ Refresh list", elem_id="refresh-btn") | |
| reextract_btn = gr.Button("⚙ Re-extract tables & images", elem_id="reextract-btn") | |
| # Confirmation row for Remove ALL | |
| with gr.Row(visible=False) as confirm_row: | |
| gr.Markdown('<span style="font-size:0.95em;color:#f87171;">⚠️ Remove ALL embeddings? This cannot be undone.</span>') | |
| confirm_yes_btn = gr.Button("✔ Yes, remove all", elem_id="confirm-yes-btn") | |
| confirm_no_btn = gr.Button("✖ Cancel", elem_id="confirm-no-btn") | |
| debug_btn = gr.Button("🔍 Debug Info", elem_id="debug-btn") | |
| debug_out = gr.HTML(value="", elem_id="debug-out") | |
| tts_audio_box= gr.Textbox(value="", visible=False, elem_id="tts-ready-box") | |
| copy_box = gr.Textbox(value="", visible=False, elem_id="copy-box") | |
| def refresh_and_update(selected_docs=None): | |
| docs, files, status_msg, model, device = get_status() | |
| # If selected_docs is provided, keep only those still present | |
| if selected_docs is not None: | |
| still_selected = [d for d in (selected_docs or []) if d in (docs or [])] | |
| else: | |
| still_selected = None | |
| return ( | |
| gr.update(choices=docs or [], value=still_selected), | |
| status_msg, | |
| gr.update(interactive=True), | |
| _header_html(model, device), | |
| gr.update(choices=docs or []), | |
| ) | |
| demo.load( | |
| fn=lambda selected: refresh_and_update(selected), | |
| inputs=[doc_list_state], | |
| outputs=[doc_list, status_text, submit_btn, header_md, source_filter_dd], | |
| ) | |
| # Unlock Web Speech API for mobile (iOS Safari blocks speechSynthesis | |
| # from async callbacks unless speak() is called once in a direct user gesture first) | |
| _JS_UNLOCK_TTS = """() => { | |
| function unlock() { | |
| if (!window.speechSynthesis) return; | |
| var u = new SpeechSynthesisUtterance(''); | |
| window.speechSynthesis.speak(u); | |
| window.speechSynthesis.cancel(); | |
| } | |
| document.addEventListener('click', unlock, {once: true}); | |
| document.addEventListener('touchend', unlock, {once: true}); | |
| }""" | |
| demo.load(fn=None, js=_JS_UNLOCK_TTS) | |
| # Client-side keep-alive: ping /status every 30 s from the browser. | |
| # This sends real HTTP traffic to the HF Space so the container never | |
| # goes idle as long as someone has the tab open. | |
| _JS_KEEPALIVE = """() => { | |
| setInterval(function() { | |
| fetch('/status').catch(function(){}); | |
| }, 30000); | |
| }""" | |
| demo.load(fn=None, js=_JS_KEEPALIVE) | |
| def _apply_filter(pattern: str): | |
| docs, *_ = get_status() | |
| docs = docs or [] | |
| if not pattern or not pattern.strip(): | |
| return gr.update(choices=docs, value=[]) | |
| matched = [] | |
| for pat in pattern.split(): | |
| matched.extend(fnmatch.filter(docs, pat)) | |
| matched = list(dict.fromkeys(matched)) # dedupe, preserve order | |
| return gr.update(choices=docs, value=matched) | |
| filter_pattern_tb.submit(fn=_apply_filter, inputs=[filter_pattern_tb], outputs=[source_filter_dd]) | |
| demo.load( | |
| fn=None, | |
| js="""() => { setTimeout(function(){ window.scrollTo(0, 0); }, 200); }""" | |
| ) | |
| demo.load( | |
| fn=None, | |
| js="""() => { | |
| // ── 0. Inject CSS for user bubbles ── | |
| var _css = document.createElement('style'); | |
| _css.textContent = '' | |
| + '.chatbot-wrap .message-row:not(.bot-row) .message-bubble,' | |
| + '.chatbot-wrap .message-row:not(.bot-row) .bubble-wrap > *,' | |
| + '.chatbot-wrap [data-testid="user"] > div,' | |
| + '.chatbot-wrap .role-user .message,' | |
| + '.chatbot-wrap .user-row .message-bubble' | |
| + ' { background:#3b82f6!important; color:#000!important; border-radius:12px!important; }' | |
| + '.chatbot-wrap .message-row:not(.bot-row) .message-bubble p,' | |
| + '.chatbot-wrap .message-row:not(.bot-row) .message-bubble span,' | |
| + '.chatbot-wrap .message-row:not(.bot-row) .prose,' | |
| + '.chatbot-wrap .message-row:not(.bot-row) .prose p,' | |
| + '.chatbot-wrap [data-testid="user"] p,' | |
| + '.chatbot-wrap .role-user p,' | |
| + '.chatbot-wrap .user-row p,' | |
| + '.chatbot-wrap .user-row span' | |
| + ' { color:#000!important; }'; | |
| document.head.appendChild(_css); | |
| // ── 1. Button gradient colors ── | |
| var READ_BLUE = {bg:'linear-gradient(135deg,#2563eb 0%,#60a5fa 100%)', sh:'0 4px 16px rgba(37,99,235,0.55)'}; | |
| var READ_ORANGE = {bg:'linear-gradient(135deg,#ea580c 0%,#fb923c 100%)', sh:'0 4px 18px rgba(234,88,12,0.6)'}; | |
| var STYLE_RULES = [ | |
| { id:'ask-btn', bg:'linear-gradient(135deg,#7c5cfc 0%,#a78bfa 100%)', sh:'0 4px 18px rgba(124,92,252,0.55)' }, | |
| { id:'file-upload', bg:'linear-gradient(135deg,#0ea5e9 0%,#38bdf8 100%)', sh:'0 4px 18px rgba(14,165,233,0.55)' }, | |
| { id:'copy-btn', bg:'linear-gradient(135deg,#059669 0%,#34d399 100%)', sh:'0 4px 16px rgba(5,150,105,0.5)' }, | |
| { id:'clear-chat-btn',bg:'linear-gradient(135deg,#dc2626 0%,#f87171 100%)', sh:'0 4px 16px rgba(220,38,38,0.5)' }, | |
| { id:'delete-btn', bg:'linear-gradient(135deg,#ef4444 0%,#fca5a5 100%)', sh:'0 4px 16px rgba(239,68,68,0.5)' }, | |
| { id:'delete-all-btn',bg:'linear-gradient(135deg,#7f1d1d 0%,#b91c1c 100%)', sh:'0 4px 18px rgba(127,29,29,0.65)' }, | |
| { id:'confirm-yes-btn',bg:'linear-gradient(135deg,#7f1d1d 0%,#b91c1c 100%)',sh:'0 4px 18px rgba(127,29,29,0.65)' }, | |
| { id:'confirm-no-btn',bg:'linear-gradient(135deg,#374151 0%,#6b7280 100%)', sh:'0 2px 10px rgba(107,114,128,0.4)' }, | |
| { id:'refresh-btn', bg:'linear-gradient(135deg,#4338ca 0%,#818cf8 100%)', sh:'0 4px 16px rgba(67,56,202,0.5)' }, | |
| { id:'add-url-btn', bg:'linear-gradient(135deg,#0d9488 0%,#2dd4bf 100%)', sh:'0 4px 16px rgba(13,148,136,0.5)' }, | |
| { id:'read-btn', bg:READ_BLUE.bg, sh:READ_BLUE.sh }, | |
| ]; | |
| function styleEl(el, bg, sh) { | |
| el.style.setProperty('background', bg, 'important'); | |
| el.style.setProperty('box-shadow', sh, 'important'); | |
| el.style.setProperty('color', '#fff', 'important'); | |
| el.style.setProperty('border', 'none', 'important'); | |
| el.style.setProperty('font-weight', '700', 'important'); | |
| el.style.setProperty('border-radius', '8px', 'important'); | |
| el.style.setProperty('letter-spacing','0.4px','important'); | |
| el.style.setProperty('text-shadow', '0 1px 3px rgba(0,0,0,0.35)', 'important'); | |
| } | |
| var _applyingColors = false; | |
| function applyColors() { | |
| if (_applyingColors) return; | |
| _applyingColors = true; | |
| for (var i = 0; i < STYLE_RULES.length; i++) { | |
| var rule = STYLE_RULES[i]; | |
| var wrap = document.getElementById(rule.id); | |
| if (!wrap) continue; | |
| var el = wrap.querySelector('button') || wrap.querySelector('label') || wrap; | |
| if (rule.id === 'read-btn') { | |
| var c = window._ttsPlaying ? READ_ORANGE : READ_BLUE; | |
| styleEl(el, c.bg, c.sh); | |
| var want = window._ttsPlaying ? 'Stop' : 'Read'; | |
| if (el.textContent.trim() !== want) el.textContent = want; | |
| } else { | |
| styleEl(el, rule.bg, rule.sh); | |
| } | |
| } | |
| _applyingColors = false; | |
| } | |
| setTimeout(applyColors, 150); | |
| setTimeout(applyColors, 700); | |
| setInterval(applyColors, 2000); | |
| var _mt = null; | |
| new MutationObserver(function() { | |
| if (_mt) clearTimeout(_mt); | |
| _mt = setTimeout(applyColors, 80); | |
| }).observe(document.body, { childList: true, subtree: true }); | |
| // ── 2. Click press feedback ── | |
| document.addEventListener('mousedown', function(e) { | |
| var btn = e.target.closest('button'); | |
| if (!btn) return; | |
| btn.style.setProperty('transform', 'scale(0.93)', 'important'); | |
| btn.style.setProperty('opacity', '0.82', 'important'); | |
| function reset() { | |
| btn.style.removeProperty('transform'); | |
| btn.style.removeProperty('opacity'); | |
| btn.removeEventListener('mouseup', reset); | |
| btn.removeEventListener('mouseleave', reset); | |
| } | |
| btn.addEventListener('mouseup', reset); | |
| btn.addEventListener('mouseleave', reset); | |
| }, true); | |
| // ── 3. Scroll to bottom — called once when answer is complete ── | |
| window._scrollChatToBottom = function() { | |
| var chatEl = document.querySelector('.chatbot-wrap'); | |
| if (!chatEl) return; | |
| var scrollEl = null; | |
| var divs = chatEl.querySelectorAll('div'); | |
| for (var i = 0; i < divs.length; i++) { | |
| if (divs[i].scrollHeight > divs[i].clientHeight + 10) scrollEl = divs[i]; | |
| } | |
| if (scrollEl) { | |
| setTimeout(function() { scrollEl.scrollTop = scrollEl.scrollHeight; }, 50); | |
| } | |
| }; | |
| // ── 4. TTS — event delegation on document (immune to re-renders) ── | |
| window._ttsText = null; | |
| window._ttsPlaying = false; | |
| function _getReadBtn() { | |
| var wrap = document.getElementById('read-btn'); | |
| if (!wrap) return null; | |
| return wrap.querySelector('button') || wrap; | |
| } | |
| window._ttsSetBtn = function(playing) { | |
| var btn = _getReadBtn(); | |
| if (!btn) return; | |
| var c = playing ? READ_ORANGE : READ_BLUE; | |
| styleEl(btn, c.bg, c.sh); | |
| btn.textContent = playing ? 'Stop' : 'Read'; | |
| }; | |
| window._ttsToggle = function() { | |
| if (!window.speechSynthesis) return; | |
| if (window._ttsPlaying) { | |
| window.speechSynthesis.cancel(); | |
| window._ttsPlaying = false; | |
| window._ttsSetBtn(false); | |
| } else if (window._ttsText) { | |
| window.speechSynthesis.cancel(); | |
| var utt = new SpeechSynthesisUtterance(window._ttsText); | |
| window._ttsPlaying = true; | |
| window._ttsSetBtn(true); | |
| utt.onend = function() { | |
| window._ttsPlaying = false; | |
| window._ttsSetBtn(false); | |
| }; | |
| window.speechSynthesis.speak(utt); | |
| } | |
| }; | |
| // Event delegation — click anywhere, check if it's the read button | |
| document.addEventListener('click', function(e) { | |
| var wrap = document.getElementById('read-btn'); | |
| if (!wrap) return; | |
| if (wrap.contains(e.target)) { | |
| if (window._ttsToggle) window._ttsToggle(); | |
| } | |
| }, true); | |
| // ── 5. Thinking overlay inside question box ── | |
| function _setupThinkingOverlay() { | |
| var indicator = document.getElementById('thinking-indicator'); | |
| var wrap = document.getElementById('chat-input-wrap'); | |
| if (!indicator || !wrap) { setTimeout(_setupThinkingOverlay, 400); return; } | |
| var ov = document.createElement('div'); | |
| ov.style.cssText = 'position:absolute;top:50%;left:14px;right:60px;transform:translateY(-50%);z-index:100;pointer-events:none;display:none;'; | |
| var line1 = document.createElement('div'); | |
| line1.style.cssText = 'color:#3b82f6;font-weight:600;font-size:1.05em;'; | |
| line1.textContent = '⏳ Thinking...'; | |
| var line2 = document.createElement('div'); | |
| line2.style.cssText = 'color:#9ca3af;font-size:0.88em;margin-top:2px;white-space:nowrap;overflow:hidden;text-overflow:ellipsis;'; | |
| ov.appendChild(line1); | |
| ov.appendChild(line2); | |
| wrap.appendChild(ov); | |
| new MutationObserver(function() { | |
| var span = indicator.querySelector('[data-q]'); | |
| var isThinking = !!indicator.textContent.trim(); | |
| if (isThinking) { | |
| line2.textContent = span ? (span.getAttribute('data-q') || '') : ''; | |
| ov.style.display = 'block'; | |
| } else { | |
| ov.style.display = 'none'; | |
| } | |
| }).observe(indicator, {childList:true, subtree:true, characterData:true}); | |
| } | |
| setTimeout(_setupThinkingOverlay, 600); | |
| }""" | |
| ) | |
| file_upload.upload( | |
| fn=upload_files, | |
| inputs=[file_upload], | |
| outputs=[upload_status, doc_list, status_text, submit_btn, header_md, source_filter_dd], | |
| ) | |
| add_url_btn.click( | |
| fn=lambda url, selected: (add_url(url), *refresh_and_update(selected), gr.update(value="")), | |
| inputs=[url_input, doc_list_state], | |
| outputs=[upload_status, doc_list, status_text, submit_btn, header_md, source_filter_dd, url_input], | |
| ) | |
| url_input.submit( | |
| fn=lambda url, selected: (add_url(url), *refresh_and_update(selected), gr.update(value="")), | |
| inputs=[url_input, doc_list_state], | |
| outputs=[upload_status, doc_list, status_text, submit_btn, header_md, source_filter_dd, url_input], | |
| ) | |
| def _delete_and_refresh(selected): | |
| msg = delete_document(selected) | |
| # Remove deleted docs from selection | |
| docs, *_ = get_status() | |
| remaining = [d for d in (selected or []) if d in (docs or [])] | |
| return (msg, *refresh_and_update(remaining), remaining) | |
| delete_btn.click( | |
| fn=_delete_and_refresh, | |
| inputs=[doc_list_state], | |
| outputs=[upload_status, doc_list, status_text, submit_btn, header_md, source_filter_dd, doc_list_state], | |
| ) | |
| # Show confirmation row when Remove ALL is clicked | |
| delete_all_btn.click( | |
| fn=lambda: gr.update(visible=True), | |
| outputs=[confirm_row], | |
| ) | |
| # Confirm: execute delete, hide confirmation row | |
| confirm_yes_btn.click( | |
| fn=lambda: (gr.update(visible=False), delete_all_embeddings(), *refresh_and_update()), | |
| outputs=[confirm_row, upload_status, doc_list, status_text, submit_btn, header_md, source_filter_dd], | |
| ) | |
| # Cancel: just hide the confirmation row | |
| confirm_no_btn.click( | |
| fn=lambda: gr.update(visible=False), | |
| outputs=[confirm_row], | |
| ) | |
| refresh_btn.click( | |
| fn=lambda selected: refresh_and_update(selected), | |
| inputs=[doc_list_state], | |
| outputs=[doc_list, status_text, submit_btn, header_md, source_filter_dd], | |
| ) | |
| def run_reextract(): | |
| r = api_post("/reextract", timeout=300) | |
| if "error" in r: | |
| return f"❌ Re-extract failed: {r['error']}" | |
| lines = [f"✅ Re-extraction complete:"] | |
| for src, counts in r.get("results", {}).items(): | |
| lines.append(f" • {src} — {counts['tables']} table(s), {counts['images']} image(s)") | |
| return "\n".join(lines) | |
| reextract_btn.click( | |
| fn=run_reextract, | |
| inputs=[], | |
| outputs=[status_text], | |
| ) | |
| debug_btn.click(fn=get_debug_info, outputs=[debug_out]) | |
| # Always update state when user changes selection | |
| doc_list.change( | |
| fn=lambda sel: sel, | |
| inputs=[doc_list], | |
| outputs=[doc_list_state], | |
| ) | |
| # Auto-refresh the document list every 5 s so users see newly indexed | |
| # files appear without manually clicking Refresh. | |
| gr.Timer(value=5).tick( | |
| fn=lambda selected: refresh_and_update(selected), | |
| inputs=[doc_list_state], | |
| outputs=[doc_list, status_text, submit_btn, header_md, source_filter_dd], | |
| ) | |
| def _thinking_html(q): | |
| import html as _html | |
| return f"<span style='color:#facc15;font-weight:600;' data-q='{_html.escape(q, quote=True)}'>⏳ Thinking...</span>" | |
| def on_submit(message, history, n, temp, src_filter): | |
| if not message.strip(): | |
| yield history, "", gr.update(), "", gr.update() | |
| return | |
| history = history or [] | |
| yield gr.update(), "", _thinking_html(message), gr.update(), gr.update() | |
| updated_history, stats = chat_fn(message, history, n, temp, src_filter) | |
| tokens_user = stats.get("tokens_user", 0) if isinstance(stats, dict) else 0 | |
| tokens_assistant = stats.get("tokens_assistant", 0) if isinstance(stats, dict) else 0 | |
| tts_text = _clean_for_tts(get_last_answer(updated_history)) if updated_history else "" | |
| tok_html = f"<span style='color:#3b82f6; font-weight:600;'>Tokens sent: {tokens_user} Tokens received: {tokens_assistant}</span>" | |
| yield updated_history, "", "", tts_text, tok_html | |
| _submit_outputs = [chatbot, msg_input, thinking_indicator, tts_audio_box, token_stats_text] | |
| submit_btn.click( | |
| fn=on_submit, | |
| inputs=[msg_input, chatbot, n_results_slider, temperature_slider, source_filter_dd], | |
| outputs=_submit_outputs, | |
| ) | |
| msg_input.submit( | |
| fn=on_submit, | |
| inputs=[msg_input, chatbot, n_results_slider, temperature_slider, source_filter_dd], | |
| outputs=_submit_outputs, | |
| ) | |
| # Wire each sample question button: load text then submit | |
| for _sq_btn, _sq_text in zip(sample_q_btns, SAMPLE_QUESTIONS): | |
| _sq_btn.click( | |
| fn=lambda q=_sq_text: q, | |
| outputs=[msg_input], | |
| ).then( | |
| fn=on_submit, | |
| inputs=[msg_input, chatbot, n_results_slider, temperature_slider, source_filter_dd], | |
| outputs=_submit_outputs, | |
| ) | |
| def on_clear_chat(): | |
| status, history = clear_memory() | |
| if history and isinstance(history[0], tuple): | |
| new_hist = [] | |
| for user, bot in history: | |
| if user is not None: | |
| new_hist.append({"role": "user", "content": user}) | |
| if bot is not None: | |
| new_hist.append({"role": "assistant", "content": bot}) | |
| history = new_hist | |
| return history, status, "<span style='color:#3b82f6; font-weight:600;'>Tokens sent: 0 Tokens received: 0</span>" | |
| tts_audio_box.change( | |
| fn=None, | |
| inputs=[tts_audio_box], | |
| js="""(val) => { | |
| if (window.speechSynthesis && window.speechSynthesis.speaking) | |
| window.speechSynthesis.cancel(); | |
| window._ttsPlaying = false; | |
| window._ttsText = val || null; | |
| if (window._ttsSetBtn) window._ttsSetBtn(false); | |
| function _scrollNow() { | |
| var chatEl = document.querySelector('.chatbot-wrap'); | |
| if (!chatEl) return; | |
| var scrollEl = null; | |
| var divs = chatEl.querySelectorAll('div'); | |
| for (var i = 0; i < divs.length; i++) { | |
| if (divs[i].scrollHeight > divs[i].clientHeight + 10) scrollEl = divs[i]; | |
| } | |
| if (scrollEl) scrollEl.scrollTop = scrollEl.scrollHeight; | |
| } | |
| setTimeout(_scrollNow, 100); | |
| setTimeout(_scrollNow, 350); | |
| setTimeout(_scrollNow, 700); | |
| }""", | |
| ) | |
| copy_btn.click(fn=get_chat_for_copy, inputs=[chatbot], outputs=[copy_box]) | |
| copy_box.change( | |
| fn=None, | |
| inputs=[copy_box], | |
| js="""(val) => { | |
| const text = val.split('\\n').slice(1).join('\\n'); | |
| if (text.trim()) navigator.clipboard.writeText(text).catch(() => {}); | |
| }""", | |
| ) | |
| clear_chat_btn.click( | |
| fn=on_clear_chat, | |
| outputs=[chatbot, memory_stats_text, token_stats_text], | |
| ) | |
| demo.queue() | |
| return demo | |
| if __name__ == "__main__": | |
| if os.environ.get("SPACE_ID"): | |
| _space_url = "https://irajkoohi-multimodalrag.hf.space" | |
| start_keep_alive_scheduler(_space_url) | |
| ui = build_ui() | |
| ui.launch(server_name="0.0.0.0", server_port=7860, show_error=True, | |
| theme=_UI_THEME, css=_UI_CSS) | |