Spaces:
Sleeping
Sleeping
| """ | |
| Dwani.ai Gradio UI - FIXED File Upload + Backend Health Check | |
| """ | |
| import gradio as gr | |
| import requests | |
| from typing import List, Dict | |
| from datetime import datetime | |
| import time | |
| import json | |
| # BACKEND MUST BE LIVE - Test this first! | |
| API_BASE = "https://gaganyatri-rag-backend.hf.space" | |
| class DwaniClient: | |
| def __init__(self, base_url=API_BASE): | |
| self.base_url = base_url.rstrip('/') | |
| def test_backend(self): | |
| """Test if backend is alive""" | |
| try: | |
| resp = requests.get(f"{self.base_url}/", timeout=5) | |
| return resp.status_code == 200, resp.json().get('status', 'Unknown') | |
| except: | |
| return False, "Backend not responding" | |
| def upload_file(self, file_obj) -> dict: | |
| """Upload single file object (Gradio format)""" | |
| files = {'file': (file_obj.name, file_obj, 'application/pdf')} | |
| resp = requests.post(f"{self.base_url}/files/upload", | |
| files=files, timeout=30) | |
| resp.raise_for_status() | |
| return resp.json() | |
| def get_file_status(self, file_id: str) -> dict: | |
| resp = requests.get(f"{self.base_url}/files/{file_id}", timeout=10) | |
| resp.raise_for_status() | |
| return resp.json() | |
| def list_files(self) -> List[dict]: | |
| resp = requests.get(f"{self.base_url}/files/", timeout=10) | |
| resp.raise_for_status() | |
| return resp.json() | |
| def chat(self, file_ids: List[str], messages: List[Dict]) -> dict: | |
| payload = {"file_ids": file_ids, "messages": messages} | |
| resp = requests.post(f"{self.base_url}/chat-with-document", | |
| json=payload, timeout=60) | |
| resp.raise_for_status() | |
| return resp.json() | |
| client = DwaniClient() | |
| uploaded_files = {} | |
| chat_history = [] | |
| selected_files = [] | |
| def test_connection(): | |
| """Test backend connection""" | |
| is_alive, status = client.test_backend() | |
| if is_alive: | |
| return f"✅ Backend LIVE: {status}", gr.update(visible=True) | |
| else: | |
| return f"❌ Backend DOWN: {status}\n\nFIX: Check https://huggingface.co/spaces/Sahana31/RAG_backend", gr.update(visible=False) | |
| def poll_file_status(file_id: str, max_wait=120): | |
| """Poll until file ready""" | |
| for i in range(max_wait//2): | |
| try: | |
| status = client.get_file_status(file_id) | |
| if status['status'] == 'completed': | |
| return status, True | |
| if status['status'] == 'failed': | |
| return status, False | |
| time.sleep(2) | |
| except Exception as e: | |
| time.sleep(2) | |
| return {'status': 'timeout'}, False | |
| def upload_multiple(files): | |
| """Handle Gradio file list""" | |
| global uploaded_files | |
| if not files: | |
| return "No files selected", gr.update(choices=[]), create_file_list(), gr.update() | |
| status_msgs = ["**📤 UPLOADING...**"] | |
| for file_obj in files: | |
| try: | |
| # Upload file | |
| result = client.upload_file(file_obj) | |
| file_id = result['file_id'] | |
| filename = result['filename'] | |
| uploaded_files[file_id] = { | |
| 'filename': filename, | |
| 'status': 'pending', | |
| 'file_id': file_id | |
| } | |
| # Poll status | |
| final_status, success = poll_file_status(file_id) | |
| if success: | |
| uploaded_files[file_id]['status'] = 'completed' | |
| status_msgs.append(f"✅ {filename}") | |
| else: | |
| uploaded_files[file_id]['status'] = final_status['status'] | |
| status_msgs.append(f"❌ {filename} - {final_status['status']}") | |
| except requests.exceptions.HTTPError as e: | |
| status_msgs.append(f"❌ {file_obj.name} - HTTP {e.response.status_code}") | |
| except Exception as e: | |
| status_msgs.append(f"❌ {file_obj.name} - {str(e)}") | |
| # Only completed files for chat | |
| choices = [(info['filename'], info['file_id']) for info in uploaded_files.values() | |
| if info['status'] == 'completed'] | |
| return "\n".join(status_msgs), gr.update(choices=choices), create_file_list(), gr.update() | |
| def refresh_files(): | |
| try: | |
| files = client.list_files() | |
| global uploaded_files | |
| uploaded_files.clear() | |
| for f in files: | |
| uploaded_files[f['file_id']] = { | |
| 'filename': f['filename'], | |
| 'status': f['status'], | |
| 'file_id': f['file_id'] | |
| } | |
| choices = [(f['filename'], f['file_id']) for f in files if f['status'] == 'completed'] | |
| return create_file_list(), gr.update(choices=choices) | |
| except Exception as e: | |
| return f"❌ Refresh failed: {str(e)}", gr.update() | |
| def create_file_list(): | |
| if not uploaded_files: | |
| return "**No files uploaded**" | |
| lines = ["**📁 Your Files:**"] | |
| for info in uploaded_files.values(): | |
| emoji = {'completed': '✅', 'processing': '🔄', 'pending': '⏳', 'failed': '❌'}.get(info['status'], '❓') | |
| lines.append(f"{emoji} {info['filename']} ({info['status']})") | |
| return "\n".join(lines) | |
| def update_selected_files(files): | |
| global selected_files | |
| selected_files = files or [] | |
| return len(selected_files) | |
| def send_message(message, history): | |
| global chat_history, selected_files | |
| if not message.strip(): | |
| return history, "" | |
| if not selected_files: | |
| return history, "⚠️ Select files first!" | |
| # Add user message | |
| user_msg = {"role": "user", "content": message} | |
| thinking_msg = {"role": "assistant", "content": "🔄 Thinking..."} | |
| new_history = history + [user_msg, thinking_msg] | |
| try: | |
| # Send to backend | |
| api_messages = chat_history + [user_msg] | |
| result = client.chat(selected_files, api_messages) | |
| # Update history | |
| chat_history.extend([user_msg, {"role": "assistant", "content": result['answer']}]) | |
| # Format response | |
| response = format_chat_response(result) | |
| final_history = history + [user_msg, {"role": "assistant", "content": response}] | |
| return final_history, "" | |
| except Exception as e: | |
| error_msg = f"❌ Error: {str(e)}" | |
| return new_history[:-1] + [{"role": "assistant", "content": error_msg}], error_msg | |
| def format_chat_response(result): | |
| answer = result['answer'] | |
| if result.get('sources'): | |
| sources = "\n\n**📚 Sources:**\n" | |
| for i, src in enumerate(result['sources'][:3], 1): | |
| sources += f"{i}. **{src['filename']}** (Page {src.get('page', '?')})\n" | |
| excerpt = src.get('excerpt', '')[:100] + "..." | |
| sources += f" > {excerpt}\n" | |
| return answer + sources | |
| return answer | |
| def clear_chat(): | |
| global chat_history | |
| chat_history = [] | |
| return [] | |
| # === GRADIO UI === | |
| with gr.Blocks(title="Dwani.ai RAG") as demo: | |
| gr.Markdown("# 📚 Dwani.ai - Document Chatbot") | |
| gr.Markdown("**Upload PDFs → Chat with page-accurate citations**") | |
| # Backend status | |
| with gr.Row(): | |
| status_display = gr.Markdown("🟡 Testing backend...") | |
| test_btn = gr.Button("🔍 Test Backend", variant="secondary") | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| gr.Markdown("## 📤 Upload PDFs") | |
| file_input = gr.File( | |
| label="Select multiple PDFs", | |
| file_types=[".pdf"], | |
| file_count="multiple" | |
| ) | |
| upload_btn = gr.Button("🚀 Upload & Process", variant="primary") | |
| status_output = gr.Markdown("Ready...") | |
| refresh_btn = gr.Button("🔄 Refresh") | |
| files_display = gr.Markdown("No files") | |
| with gr.Column(scale=2): | |
| gr.Markdown("## 📋 Document Manager") | |
| file_checkboxes = gr.CheckboxGroup( | |
| label="Chat with these files:", | |
| choices=[], | |
| value=[], | |
| ) | |
| file_count = gr.Number(label="Selected", interactive=False) | |
| with gr.Row(): | |
| gr.Markdown("## 💬 Chat") | |
| chatbot = gr.Chatbot(height=500) | |
| with gr.Row(): | |
| msg_input = gr.Textbox(placeholder="Ask about your documents...", scale=4) | |
| send_btn = gr.Button("Send", variant="primary", scale=1) | |
| clear_btn = gr.Button("🗑️ New Chat", variant="secondary") | |
| # Events | |
| test_btn.click(test_connection, outputs=[status_display, upload_btn]) | |
| upload_btn.click( | |
| upload_multiple, | |
| inputs=file_input, | |
| outputs=[status_output, file_checkboxes, files_display, file_count] | |
| ) | |
| refresh_btn.click(refresh_files, outputs=[files_display, file_checkboxes]) | |
| file_checkboxes.change(update_selected_files, inputs=file_checkboxes, outputs=file_count) | |
| send_btn.click(send_message, inputs=[msg_input, chatbot], outputs=[chatbot, msg_input]) | |
| msg_input.submit(send_message, inputs=[msg_input, chatbot], outputs=[chatbot, msg_input]) | |
| clear_btn.click(clear_chat, outputs=chatbot) | |
| if __name__ == "__main__": | |
| demo.launch(server_name="0.0.0.0", server_port=7860) | |