Update app.py

app.py

@@ -1,114 +1,273 @@
+
 import gradio as gr
 import requests
+from typing import List, Dict
 import time
+import os
 
-API_BASE = "https://sahana31-rag-backend.hf.space"
-
+# Disable Gradio analytics for HF Spaces
+os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"
 
-class BackendClient:
-    def __init__(self, base):
-        self.base = base.rstrip("/")
+# === CONFIG ===
+API_BASE = "https://sahana31-rag-backend.hf.space"
 
-    def health(self):
-        r = requests.get(f"{self.base}/health", timeout=10)
-        r.raise_for_status()
+# === 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()
 
-    def upload(self, file_obj):
-        r = requests.post(
-            f"{self.base}/files/upload",
-            files={"file": (file_obj.name, file_obj, "application/pdf")},
-            timeout=120
-        )
-        r.raise_for_status()
-        return r.json()
-
-    def status(self, file_id):
-        r = requests.get(f"{self.base}/files/{file_id}", timeout=20)
-        r.raise_for_status()
-        return r.json()
-
-    def chat(self, file_ids, messages):
-        r = requests.post(
-            f"{self.base}/chat-with-document",
-            json={"file_ids": file_ids, "messages": messages},
-            timeout=180
-        )
-        r.raise_for_status()
-        return r.json()
 
+# === GLOBALS ===
+client = DwaniClient()
+uploaded_files = {}
+chat_history: List[Dict] = []
+selected_files = []
 
-client = BackendClient(API_BASE)
+# === 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)
 
-def upload_files(files, store):
+# === UI FUNCTIONS ===
+def upload_multiple(files):
+    """Handle multiple PDF uploads"""
     if not files:
-        return "No files", [], store
+        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()
 
-    client.health()
+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()
 
-    status = []
-    choices = []
+def update_selected_files(files):
+    """Update selected files"""
+    global selected_files
+    selected_files = files or []
+    return len(selected_files)
 
-    for f in files:
-        res = client.upload(f)
-        fid = res["file_id"]
-        name = res["filename"]
+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
 
-        for _ in range(60):
-            s = client.status(fid)
-            if s["status"] == "completed":
-                store[fid] = name
-                choices.append((name, fid))
-                status.append(f"✅ {name}")
-                break
-            time.sleep(2)
-
-    return "\n".join(status), choices, store
-
-
-def chat_fn(msg, history, selected, chat_state):
-    if not selected:
-        history.append((msg, "Select a document first"))
+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)}"
 
-    chat_state.append({"role": "user", "content": msg})
-    history.append((msg, "Thinking..."))
-
-    res = client.chat(selected, chat_state)
-    answer = res["answer"]
-
-    chat_state.append({"role": "assistant", "content": answer})
-    history[-1] = (msg, answer)
-
-    return history, ""
-
-
-with gr.Blocks() as demo:
-    gr.Markdown("# 📄 RAG Document Chat")
-
-    files_state = gr.State({})
-    chat_state = gr.State([])
+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():
-        files = gr.File(file_types=[".pdf"], file_count="multiple")
-        upload_btn = gr.Button("Upload")
-        status = gr.Markdown()
-        selector = gr.CheckboxGroup(label="Documents")
-
-    bot = gr.Chatbot()
-    msg = gr.Textbox()
-    send = gr.Button("Send")
-
+        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_files,
-        inputs=[files, files_state],
-        outputs=[status, selector, files_state],
+        upload_multiple,
+        inputs=file_input,
+        outputs=[status_output, file_checkboxes, files_display]
     )
-
-    send.click(
-        chat_fn,
-        inputs=[msg, bot, selector, chat_state],
-        outputs=[bot, msg],
+    
+    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
     )
 
-demo.queue()
+# === LAUNCH FOR HF SPACES ===
+if __name__ == "__main__":
+    demo.launch(show_error=True)