Update app.py

app.py

@@ -1,41 +1,69 @@
-# app.py
 import gradio as gr
 import requests
 import pandas as pd
 import base64
 from io import BytesIO
+import json
 
 # 🔧 CONFIGURE: Replace with your EC2 public IP
 API_URL = "http://18.191.172.85:5000/ask"
-# Example: "http://3.8.123.45:5000/ask"
+HEALTH_URL = "http://18.191.172.85:5000/health"
+
+def test_connection():
+    try:
+        response = requests.get(HEALTH_URL, timeout=10)
+        if response.status_code == 200:
+            return f"✅ Connection successful! Status: {response.status_code}\nResponse: {response.text}"
+        else:
+            return f"⚠️ Connection returned status: {response.status_code}\nResponse: {response.text}"
+    except requests.exceptions.ConnectionError:
+        return "❌ Connection failed. The server might be down or not accessible from Hugging Face."
+    except requests.exceptions.Timeout:
+        return "⏱️ Connection timed out."
+    except Exception as e:
+        return f"🚨 Unexpected error: {str(e)}"
 
 def query_database(question):
     if not question.strip():
-        return "", None, None, "Please enter a question."
+        return "", None, None, "⚠️ Please enter a question."
     
     try:
-        response = requests.post(API_URL, json={"question": question}, timeout=30)
-        result = response.json()
+        response = requests.post(
+            API_URL, 
+            json={"question": question}, 
+            timeout=60
+        )
+        
+        try:
+            result = response.json()
+        except json.JSONDecodeError:
+            return "", None, None, f"❌ Invalid response from server: {response.text}"
+        
+        if response.status_code != 200:
+            error_msg = result.get("error", f"HTTP {response.status_code}")
+            return "", None, None, f"❌ Server error: {error_msg}"
         
         if "error" in result:
             return "", None, None, f"❌ {result['error']}"
         
         sql = result.get("sql", "N/A")
         rows = result.get("result", [])
-        chart_data = result.get("chart")  # base64 PNG
+        chart_data = result.get("chart")
         
         df = pd.DataFrame(rows) if rows else pd.DataFrame()
         
-        # Decode chart
         chart_image = None
         if chart_data:
-            image_bytes = base64.b64decode(chart_data)
-            chart_image = BytesIO(image_bytes)
+            try:
+                image_bytes = base64.b64decode(chart_data)
+                chart_image = BytesIO(image_bytes)
+            except Exception as e:
+                print(f"Error decoding chart: {e}")
         
-        return sql, df, chart_image, ""
+        return sql, df, chart_image, "✅ Query completed successfully!"
     
     except requests.exceptions.ConnectionError:
-        return "", None, None, "❌ Connection failed. Is your API running?"
+        return "", None, None, "❌ Connection failed. Is your API running and accessible from Hugging Face?"
     except requests.exceptions.Timeout:
         return "", None, None, "⏱️ Request timed out. Try a simpler question."
     except Exception as e:
@@ -56,8 +84,8 @@ theme = gr.themes.Default(
 with gr.Blocks(theme=theme, title="Enterprise SQL Assistant") as demo:
     gr.HTML("""
     <div style="text-align: center; padding: 20px;">
-        <h1 style="color: #3B82F6;">📊 Enterprise SQL Assistant</h1>
-        <p style="color: #9CA3AF; font-size: 1.1em;">
+        <h1 style="color: #3B82F6; margin-bottom: 5px;">📊 Enterprise SQL Assistant</h1>
+        <p style="color: #9CA3AF; font-size: 1.1em; margin-top: 0;">
             Ask questions about your data. Get SQL, results, and insights.
         </p>
     </div>
@@ -73,6 +101,10 @@ with gr.Blocks(theme=theme, title="Enterprise SQL Assistant") as demo:
             )
             submit_btn = gr.Button("🚀 Generate SQL & Results", variant="primary")
             
+            gr.Markdown("### 🔧 Connection Test")
+            test_btn = gr.Button("Test Connection to EC2")
+            connection_status = gr.Textbox(label="Connection Status", interactive=False)
+            
             gr.Markdown("### 💡 Example Queries")
             gr.Examples(
                 examples=[
@@ -80,7 +112,10 @@ with gr.Blocks(theme=theme, title="Enterprise SQL Assistant") as demo:
                     ["What is the total transaction amount?"],
                     ["Show members with their account balances"],
                     ["Which member has the highest balance?"],
-                    ["Show transaction trends over time"]
+                    ["Show transaction trends over time"],
+                    ["Count of members by status"],
+                    ["Show me the first 10 rows"],
+                    ["What is the average age of members?"]
                 ],
                 inputs=question_input,
                 label=""
@@ -110,6 +145,11 @@ with gr.Blocks(theme=theme, title="Enterprise SQL Assistant") as demo:
                     )
     
     # Events
+    test_btn.click(
+        fn=test_connection,
+        outputs=[connection_status]
+    )
+    
     submit_btn.click(
         fn=query_database,
         inputs=question_input,