Update app.py

app.py

@@ -8,12 +8,12 @@ import time
 import traceback
 from PIL import Image
 
-# 🔧 CONFIGURE: Your Flask API URL
+# 🔧 CONFIGURE: Your Flask API URL (Public URL)
 FLASK_API_URL = "http://18.219.185.9:5000"
 
 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()
@@ -47,7 +47,7 @@ def query_database(question):
         if not response.text:
             error_msg = "Empty response from server"
             print(f"ERROR: {error_msg}")
-            return "", None, None, f"❌ {error_msg}", f"Request time: {elapsed_time:.2f}s", "Empty response"
+            return "", None, None, f"❌ {error_msg}", f"Request time: {elapsed_time:.2f}s", "Empty response", "", ""
         
         # Try to parse JSON response
         try:
@@ -57,13 +57,13 @@ def query_database(question):
             error_msg = f"Invalid JSON response: {str(e)}"
             print(f"ERROR: {error_msg}")
             print(f"Response text (first 500 chars): {response.text[:500]}")
-            return "", None, None, f"❌ {error_msg}", f"Request time: {elapsed_time:.2f}s", "JSON parse error"
+            return "", None, None, f"❌ {error_msg}", f"Request time: {elapsed_time:.2f}s", "JSON parse error", "", ""
         
         # Check HTTP status
         if response.status_code != 200:
             error_msg = result.get("error", f"HTTP {response.status_code}")
             print(f"ERROR: HTTP status {response.status_code}: {error_msg}")
-            return "", None, None, f"❌ Server error: {error_msg}", f"Request time: {elapsed_time:.2f}s", "HTTP error"
+            return "", None, None, f"❌ Server error: {error_msg}", f"Request time: {elapsed_time:.2f}s", "HTTP error", "", ""
         
         # Extract data
         sql = result.get("sql", "")
@@ -80,46 +80,70 @@ def query_database(question):
         df = pd.DataFrame(rows) if rows else pd.DataFrame()
         print(f"DataFrame shape: {df.shape}")
         
-        # Process chart
+        # Process chart - UPDATED TO HANDLE NEW STRUCTURED FORMAT
         chart_image = None
-        if chart_data:
+        chart_title = ""
+        chart_type = ""
+        
+        if chart_data and isinstance(chart_data, dict):
+            try:
+                # Extract the base64 image string from the structured chart data
+                chart_image_b64 = chart_data.get("image")
+                if chart_image_b64:
+                    image_bytes = base64.b64decode(chart_image_b64)
+                    chart_image = Image.open(BytesIO(image_bytes))
+                    print("Chart decoded successfully")
+                
+                # Extract chart metadata
+                chart_title = chart_data.get("title", "")
+                chart_type = chart_data.get("type", "")
+                print(f"Chart type: {chart_type}")
+                print(f"Chart title: {chart_title}")
+            except Exception as e:
+                print(f"Error decoding chart: {e}")
+        elif chart_data and isinstance(chart_data, str):
+            # Fallback for backward compatibility
             try:
                 image_bytes = base64.b64decode(chart_data)
                 chart_image = Image.open(BytesIO(image_bytes))
-                print("Chart decoded successfully")
+                print("Chart decoded successfully (fallback)")
             except Exception as e:
-                print(f"Error decoding chart: {e}")
+                print(f"Error decoding chart (fallback): {e}")
         
         # Prepare details
         details = f"Request time: {elapsed_time:.2f}s\n"
         details += f"Status code: {response.status_code}\n"
         details += f"Rows returned: {row_count}\n"
         details += f"Chart generated: {'Yes' if chart_data else 'No'}\n"
+        if chart_type:
+            details += f"Chart type: {chart_type}\n"
+        if chart_title:
+            details += f"Chart title: {chart_title}\n"
         details += f"Response size: {len(response.text)} bytes"
         
         print(f"=== REQUEST COMPLETED SUCCESSFULLY ===")
         
-        return sql, df, chart_image, f"✅ Query completed successfully", details, "Success"
+        return sql, df, chart_image, f"✅ Query completed successfully", details, "Success", chart_type, chart_title
     
     except requests.exceptions.ConnectionError as e:
         error_msg = f"Connection failed: {str(e)}"
         print(f"CONNECTION ERROR: {error_msg}")
         print(f"Traceback: {traceback.format_exc()}")
-        return "", None, None, f"❌ {error_msg}", "Connection error", "Connection error"
+        return "", None, None, f"❌ {error_msg}", "Connection error", "Connection error", "", ""
     except requests.exceptions.Timeout:
         error_msg = "Request timed out after 300 seconds"
         print(f"TIMEOUT ERROR: {error_msg}")
-        return "", None, None, f"⏱️ {error_msg}", "Timeout error", "Timeout error"
+        return "", None, None, f"⏱️ {error_msg}", "Timeout error", "Timeout error", "", ""
     except requests.exceptions.RequestException as e:
         error_msg = f"Request exception: {str(e)}"
         print(f"REQUEST ERROR: {error_msg}")
         print(f"Traceback: {traceback.format_exc()}")
-        return "", None, None, f"❌ {error_msg}", "Request error", "Request error"
+        return "", None, None, f"❌ {error_msg}", "Request error", "Request error", "", ""
     except Exception as e:
         error_msg = f"Unexpected error: {str(e)}"
         print(f"UNEXPECTED ERROR: {error_msg}")
         print(f"Traceback: {traceback.format_exc()}")
-        return "", None, None, f"🚨 {error_msg}", f"Error: {str(e)}", "Unexpected error"
+        return "", None, None, f"🚨 {error_msg}", f"Error: {str(e)}", "Unexpected error", "", ""
 
 def check_health():
     try:
@@ -207,6 +231,13 @@ with gr.Blocks(theme=theme, title="Enterprise SQL Assistant", css="""
         padding: 10px;
         background-color: #1e293b;
     }
+    .chart-metadata {
+        background-color: #1e293b;
+        border: 1px solid #334155;
+        border-radius: 8px;
+        padding: 10px;
+        margin-bottom: 10px;
+    }
     /* Fix input text color */
     .gradio-container input[type="text"], 
     .gradio-container textarea {
@@ -245,10 +276,10 @@ with gr.Blocks(theme=theme, title="Enterprise SQL Assistant", css="""
                     example3 = gr.Button("Show members with their account balances", elem_classes=["example-btn"])
                     example4 = gr.Button("Which member has the highest balance?", elem_classes=["example-btn"])
                 with gr.Column():
-                    example5 = gr.Button("Show transaction trends over time", elem_classes=["example-btn"])
-                    example6 = gr.Button("Count of members by status", elem_classes=["example-btn"])
-                    example7 = gr.Button("Show me the first 10 rows", elem_classes=["example-btn"])
-                    example8 = gr.Button("What is the average age of members?", elem_classes=["example-btn"])
+                    example5 = gr.Button("Show transaction trends over time", elem_classes=["example-btn"])  # Time series
+                    example6 = gr.Button("Count of members by status", elem_classes=["example-btn"])  # Bar chart
+                    example7 = gr.Button("Show distribution of transaction amounts", elem_classes=["example-btn"])  # Histogram
+                    example8 = gr.Button("Show correlations between numeric fields", elem_classes=["example-btn"])  # Heatmap
             
             gr.Markdown("### 🩺 API Health Check")
             health_btn = gr.Button("Check API Health", variant="secondary")
@@ -273,12 +304,18 @@ with gr.Blocks(theme=theme, title="Enterprise SQL Assistant", css="""
                     )
                 
                 with gr.Tab("Visual Insights"):
-                    chart_output = gr.Image(
-                        label="Chart",
-                        type="pil",
-                        height=400,
-                        elem_classes="chart-container"
-                    )
+                    with gr.Row():
+                        with gr.Column(scale=3):
+                            chart_output = gr.Image(
+                                label="Chart",
+                                type="pil",
+                                height=400,
+                                elem_classes="chart-container"
+                            )
+                        with gr.Column(scale=1):
+                            gr.Markdown("### 📊 Chart Information")
+                            chart_type_output = gr.Markdown(label="Chart Type", elem_classes="chart-metadata")
+                            chart_title_output = gr.Markdown(label="Chart Title", elem_classes="chart-metadata")
             
             gr.Markdown("### 📊 Request Details")
             request_details = gr.Textbox(label="Request Details", interactive=False, lines=6)
@@ -307,13 +344,13 @@ with gr.Blocks(theme=theme, title="Enterprise SQL Assistant", css="""
         ).then(
             fn=query_database,
             inputs=question_input,
-            outputs=[sql_output, results_output, chart_output, chatbot_output, request_details, error_details]
+            outputs=[sql_output, results_output, chart_output, chatbot_output, request_details, error_details, chart_type_output, chart_title_output]
         )
     
     submit_btn.click(
         fn=query_database,
         inputs=question_input,
-        outputs=[sql_output, results_output, chart_output, chatbot_output, request_details, error_details]
+        outputs=[sql_output, results_output, chart_output, chatbot_output, request_details, error_details, chart_type_output, chart_title_output]
     )
 
 # Launch