import gradio as gr import requests from typing import List, Dict import time import os # Disable Gradio analytics for HF Spaces os.environ["GRADIO_ANALYTICS_ENABLED"] = "False" # === CONFIG === API_BASE = "https://sahana31-rag-backend.hf.space" # === CLIENT === class DwaniClient: def __init__(self, base_url=API_BASE): self.base_url = base_url.rstrip('/') def upload_file(self, file) -> dict: """Upload a file (HF temp file compatible)""" files = { "file": (file.name.split("/")[-1], open(file.name, "rb"), "application/pdf") } resp = requests.post(f"{self.base_url}/files/upload", files=files, timeout=300) 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=60) resp.raise_for_status() return resp.json() def list_files(self) -> List[dict]: resp = requests.get(f"{self.base_url}/files/", timeout=60) 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=300) resp.raise_for_status() return resp.json() # === GLOBALS === client = DwaniClient() uploaded_files = {} chat_history: List[Dict] = [] selected_files = [] # === HELPERS === def poll_file_status(file_id: str, max_wait=60): """Wait for file processing""" for _ in range(max_wait): 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: time.sleep(2) return {'status': 'timeout'}, False def create_file_list(): """Display all files with status""" 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) # === UI FUNCTIONS === def upload_multiple(files): """Handle multiple PDF uploads""" if not files: return "No files selected", gr.update(choices=[]), "No files uploaded" global uploaded_files status_msgs = [] for file in files: try: # Upload each file result = client.upload_file(file) file_id = result['file_id'] filename = result['filename'] uploaded_files[file_id] = { 'filename': filename, 'status': 'pending', 'file_id': file_id } # Poll for completion status, success = poll_file_status(file_id) if success: uploaded_files[file_id]['status'] = 'completed' status_msgs.append(f"✅ {filename} - READY") else: uploaded_files[file_id]['status'] = 'failed' status_msgs.append(f"❌ {filename} - FAILED") except Exception as e: status_msgs.append(f"❌ {file.name} - ERROR: {str(e)}") # Update choices for only completed files 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() def refresh_files(): """Refresh file list from backend""" try: files = client.list_files() global uploaded_files uploaded_files.clear() for f in files: uploaded_files[f['file_id']] = f choices = [(f['filename'], f['file_id']) for f in files if f['status'] == 'completed'] return create_file_list(), gr.update(choices=choices) except: return "Refresh failed", gr.update() def update_selected_files(files): """Update selected files""" global selected_files selected_files = files or [] return len(selected_files) def format_chat_response(result): """Format response with sources""" answer = result['answer'] if result.get('sources'): sources = "\n\n**📚 Sources:**\n" for i, src in enumerate(result['sources'][:5], 1): sources += f"{i}. **{src['filename']}** (Page {src['page']})\n" sources += f" > {src['excerpt'][:120]}...\n\n" return answer + sources return answer def send_message(message, history): """Send chat message""" global chat_history, selected_files if not message.strip(): return history, "" if not selected_files: return history, "⚠️ Please select files first!" user_message = {"role": "user", "content": message} assistant_message = {"role": "assistant", "content": "Thinking..."} new_history = history + [user_message, assistant_message] try: api_messages = chat_history + [user_message] result = client.chat(selected_files, api_messages) chat_history.append(user_message) chat_history.append({"role": "assistant", "content": result['answer']}) full_response = format_chat_response(result) final_history = history + [user_message, {"role": "assistant", "content": full_response}] return final_history, "" except Exception as e: error_response = {"role": "assistant", "content": f"❌ Error: {str(e)}"} return new_history[:-1] + [error_response], f"Error: {str(e)}" def clear_chat(): global chat_history chat_history = [] return [] # === GRADIO UI === with gr.Blocks(title="Dwani.ai", theme=gr.themes.Soft()) as demo: gr.Markdown("# 📚 Dwani.ai - Document Chat") gr.Markdown("**Upload multiple PDFs → Chat with page-accurate citations**") with gr.Row(): with gr.Column(scale=1): gr.Markdown("## 📤 Upload Multiple PDFs") file_input = gr.File( label="Select PDFs (Ctrl+Click for multiple)", file_types=[".pdf"], file_count="multiple" ) upload_btn = gr.Button("🚀 Upload & Process All", variant="primary") status_output = gr.Markdown("Ready to upload...") refresh_btn = gr.Button("🔄 Refresh List") files_display = gr.Markdown("No files uploaded") with gr.Column(scale=2): gr.Markdown("## 📋 File Manager") file_checkboxes = gr.CheckboxGroup( label="Select documents to chat with:", choices=[], value=[], info="Only completed files appear here" ) file_count = gr.Number(label="Files selected", value=0, interactive=False) with gr.Row(): gr.Markdown("## 💬 Chat with Documents") chatbot = gr.Chatbot( label="Ask questions about your documents", height=500, avatar_images=("user.png", "assistant.png") ) with gr.Row(): msg_input = gr.Textbox( label="Your question", placeholder="e.g., What are the payment terms? When does the contract expire?", scale=4 ) send_btn = gr.Button("Send", variant="primary", scale=1) with gr.Row(): clear_btn = gr.Button("🗑️ New Chat", variant="secondary") # === EVENTS === upload_btn.click( upload_multiple, inputs=file_input, outputs=[status_output, file_checkboxes, files_display] ) 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 ) # === LAUNCH FOR HF SPACES === if __name__ == "__main__": demo.launch(show_error=True)