Update app.py

app.py

@@ -4,6 +4,7 @@ import pandas as pd
 import base64
 from io import BytesIO
 import json
+import time
 
 # 🔧 CONFIGURE: Replace with your EC2 public IP
 API_URL = "http://18.191.172.85:5000/ask"
@@ -11,7 +12,7 @@ HEALTH_URL = "http://18.191.172.85:5000/health"
 
 def test_connection():
     try:
-        response = requests.get(HEALTH_URL, timeout=10)
+        response = requests.get(HEALTH_URL, timeout=30)  # Increased to 30 seconds
         if response.status_code == 200:
             return f"✅ Connection successful! Status: {response.status_code}\nResponse: {response.text}"
         else:
@@ -19,37 +20,64 @@ def test_connection():
     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."
+        return "⏱️ Connection timed out after 30 seconds."
     except Exception as e:
         return f"🚨 Unexpected error: {str(e)}"
 
+def test_chart_generation():
+    try:
+        response = requests.get("http://18.191.172.85:5000/test-chart", timeout=60)  # Increased to 60 seconds
+        if response.status_code == 200:
+            result = response.json()
+            if 'chart' in result:
+                image_bytes = base64.b64decode(result['chart'])
+                chart_image = BytesIO(image_bytes)
+                return "✅ Test chart generated successfully!", chart_image
+            else:
+                return f"❌ Error: {result.get('message', 'Unknown error')}", None
+        else:
+            return f"❌ HTTP Error: {response.status_code}", None
+    except Exception as e:
+        return f"❌ Connection error: {str(e)}", None
+
 def query_database(question):
     if not question.strip():
-        return "", None, None, "⚠️ Please enter a question."
+        return "", None, None, "⚠️ Please enter a question.", ""
     
     try:
+        start_time = time.time()
+        print(f"Sending request to: {API_URL}")
+        print(f"Question: {question}")
+        
         response = requests.post(
             API_URL, 
             json={"question": question}, 
-            timeout=60
+            timeout=300  # Increased timeout to 300 seconds (5 minutes)
         )
         
+        elapsed_time = time.time() - start_time
+        print(f"Response status: {response.status_code}")
+        print(f"Response time: {elapsed_time:.2f} seconds")
+        
         try:
             result = response.json()
         except json.JSONDecodeError:
-            return "", None, None, f"❌ Invalid response from server: {response.text}"
+            return "", None, None, f"❌ Invalid response from server: {response.text}", f"Request time: {elapsed_time:.2f}s"
         
         if response.status_code != 200:
             error_msg = result.get("error", f"HTTP {response.status_code}")
-            return "", None, None, f"❌ Server error: {error_msg}"
+            return "", None, None, f"❌ Server error: {error_msg}", f"Request time: {elapsed_time:.2f}s"
         
         if "error" in result:
-            return "", None, None, f"❌ {result['error']}"
+            return "", None, None, f"❌ {result['error']}", f"Request time: {elapsed_time:.2f}s"
         
         sql = result.get("sql", "N/A")
         rows = result.get("result", [])
         chart_data = result.get("chart")
         
+        print(f"Chart data received: {'Yes' if chart_data else 'No'}")
+        print(f"Chart data length: {len(chart_data) if chart_data else 0}")
+        
         df = pd.DataFrame(rows) if rows else pd.DataFrame()
         
         chart_image = None
@@ -57,17 +85,22 @@ def query_database(question):
             try:
                 image_bytes = base64.b64decode(chart_data)
                 chart_image = BytesIO(image_bytes)
+                print("Chart decoded successfully")
             except Exception as e:
                 print(f"Error decoding chart: {e}")
         
-        return sql, df, chart_image, "✅ Query completed successfully!"
+        details = f"Request time: {elapsed_time:.2f}s\n"
+        details += f"Rows returned: {len(rows)}\n"
+        details += f"Chart generated: {'Yes' if chart_data else 'No'}"
+        
+        return sql, df, chart_image, "✅ Query completed successfully!", details
     
-    except requests.exceptions.ConnectionError:
-        return "", None, None, "❌ Connection failed. Is your API running and accessible from Hugging Face?"
+    except requests.exceptions.ConnectionError as e:
+        return "", None, None, f"❌ Connection failed: {str(e)}", "Connection error"
     except requests.exceptions.Timeout:
-        return "", None, None, "⏱️ Request timed out. Try a simpler question."
+        return "", None, None, "⏱️ Request timed out after 300 seconds. Try a simpler question.", "Timeout error"
     except Exception as e:
-        return "", None, None, f"🚨 Unexpected error: {str(e)}"
+        return "", None, None, f"🚨 Unexpected error: {str(e)}", f"Error: {str(e)}"
 
 # 🎨 Theme: Enterprise Dark Blue
 theme = gr.themes.Default(
@@ -105,6 +138,11 @@ with gr.Blocks(theme=theme, title="Enterprise SQL Assistant") as demo:
             test_btn = gr.Button("Test Connection to EC2")
             connection_status = gr.Textbox(label="Connection Status", interactive=False)
             
+            gr.Markdown("### 🎨 Chart Generation Test")
+            test_chart_btn = gr.Button("Test Chart Generation")
+            chart_test_status = gr.Textbox(label="Chart Test Status", interactive=False)
+            test_chart_output = gr.Image(label="Test Chart", type="pil")
+            
             gr.Markdown("### 💡 Example Queries")
             gr.Examples(
                 examples=[
@@ -113,7 +151,7 @@ with gr.Blocks(theme=theme, title="Enterprise SQL Assistant") as demo:
                     ["Show members with their account balances"],
                     ["Which member has the highest balance?"],
                     ["Show transaction trends over time"],
-                    ["Count of members by status"],
+                    ["Count of members by status"],  # This should generate a chart
                     ["Show me the first 10 rows"],
                     ["What is the average age of members?"]
                 ],
@@ -143,6 +181,9 @@ with gr.Blocks(theme=theme, title="Enterprise SQL Assistant") as demo:
                         height=400,
                         elem_classes="chart-container"
                     )
+            
+            gr.Markdown("### 📊 Request Details")
+            request_details = gr.Textbox(label="Request Details", interactive=False, lines=4)
     
     # Events
     test_btn.click(
@@ -150,13 +191,17 @@ with gr.Blocks(theme=theme, title="Enterprise SQL Assistant") as demo:
         outputs=[connection_status]
     )
     
+    test_chart_btn.click(
+        fn=test_chart_generation,
+        outputs=[chart_test_status, test_chart_output]
+    )
+    
     submit_btn.click(
         fn=query_database,
         inputs=question_input,
-        outputs=[sql_output, results_output, chart_output, status]
+        outputs=[sql_output, results_output, chart_output, status, request_details]
     )
 
 # Launch
 if __name__ == "__main__":
-    demo.launch()
-    
+    demo.launch()