Update app.py

app.py

@@ -7,6 +7,7 @@ import json
 import time
 import traceback
 from PIL import Image
+
 # 🔧 CONFIGURE: Your Flask API URL (Public URL)
 FLASK_API_URL = "http://3.16.57.66:5000"
 
@@ -20,7 +21,6 @@ def query_database(question, dashboard_mode=False, chart_type=None):
         print("=== NEW QUERY REQUEST ===")
         print(f"Time: {time.strftime('%Y-%m-%d %H:%M:%S')}")
         print(f"Question: {question}")
-        print(f"Dashboard mode: {dashboard_mode}")
         print(f"Chart type: {chart_type}")
         print("="*50)
         
@@ -30,10 +30,6 @@ def query_database(question, dashboard_mode=False, chart_type=None):
             "visualize": True
         }
         
-        # Add dashboard_mode if specified
-        if dashboard_mode:
-            payload["dashboard_mode"] = dashboard_mode
-        
         # Add chart_type if specified
         if chart_type and chart_type != "auto":
             payload["viz_type"] = chart_type
@@ -59,7 +55,7 @@ def query_database(question, dashboard_mode=False, chart_type=None):
         if not response.text:
             error_msg = "Empty response from server"
             print(f"ERROR: {error_msg}")
-            return "", None, None, "⚠️ Please enter a question.", "", f"❌ {error_msg}", f"Request time: {elapsed_time:.2f}s", "Empty response", "", "", "", ""
+            return "", None, None, "⚠️ Please enter a question.", "", "", "", "", "", ""
         
         # Try to parse JSON response
         try:
@@ -69,7 +65,7 @@ def query_database(question, dashboard_mode=False, chart_type=None):
             error_msg = f"Invalid JSON response: {str(e)}"
             print(f"ERROR: {error_msg}")
             print(f"Response text (first 1000 chars): {response.text[:1000]}")
-            return "", None, None, "⚠️ Please enter a question.", "", f"❌ {error_msg}", f"Request time: {elapsed_time:.2f}s", "JSON parse error", "", "", "", ""
+            return "", None, None, "⚠️ Please enter a question.", "", "", "", "", "", ""
         
         # Check HTTP status
         if response.status_code != 200:
@@ -80,21 +76,24 @@ def query_database(question, dashboard_mode=False, chart_type=None):
             if response.status_code == 500:
                 error_details = f"Server Error (500): {error_msg}\n"
                 error_details += f"Response: {json.dumps(result, indent=2)}"
-                return "", None, None, "⚠️ Please enter a question.", "", f"❌ Server error: {error_msg}", f"Request time: {elapsed_time:.2f}s", error_details, "", "", "", ""
+                return "", None, None, "⚠️ Please enter a question.", "", "", "", error_details, "", "", ""
             
-            return "", None, None, "⚠️ Please enter a question.", "", f"❌ Server error: {error_msg}", f"Request time: {elapsed_time:.2f}s", "HTTP error", "", "", "", ""
+            return "", None, None, "⚠️ Please enter a question.", "", "", "", f"HTTP error: {error_msg}", "", "", ""
         
-        # Extract data
+        # Extract data - Updated to match new response structure
         sql = result.get("sql", "")
+        chart = result.get("chart", {})  # Changed from "visualization" to "chart"
+        chart_generated = result.get("chart_generated", False)  # New field
+        row_count = result.get("row_count", 0)
+        
+        # Extract result data if available (for backward compatibility)
         rows = result.get("result", [])
-        visualization = result.get("visualization", {})
-        row_count = result.get("row_count", len(rows))
         
         print(f"\n=== EXTRACTED DATA ===")
         print(f"SQL: {sql[:100] if sql else 'None'}...")
-        print(f"Rows count: {row_count}")
-        print(f"Visualization received: {'Yes' if visualization else 'No'}")
-        print(f"Visualization details: {visualization}")
+        print(f"Chart generated: {chart_generated}")
+        print(f"Row count: {row_count}")
+        print(f"Chart data: {chart}")
         
         # Create DataFrame
         df = pd.DataFrame(rows) if rows else pd.DataFrame()
@@ -110,21 +109,22 @@ def query_database(question, dashboard_mode=False, chart_type=None):
         show_image = True
         show_html = False
         
-        if visualization:
+        # Only process chart if chart_generated is True
+        if chart_generated and chart:
             try:
-                if isinstance(visualization, dict):
+                if isinstance(chart, dict):
                     # Check if it's an error response
-                    if "error" in visualization:
-                        chart_error = visualization["error"]
+                    if "error" in chart:
+                        chart_error = chart["error"]
                         print(f"Chart error: {chart_error}")
                     else:
                         # Extract the base64 image string
-                        chart_image_b64 = visualization.get("image")
+                        chart_image_b64 = chart.get("image")
                         if chart_image_b64:
-                            # Handle base64 prefix
-                            if chart_image_b64.startswith("data:image/"):
-                                chart_image_b64 = chart_image_b64.split(",")[1]
                             try:
+                                # Handle base64 prefix
+                                if chart_image_b64.startswith("data:image/"):
+                                    chart_image_b64 = chart_image_b64.split(",")[1]
                                 image_bytes = base64.b64decode(chart_image_b64)
                                 chart_image = Image.open(BytesIO(image_bytes))
                                 print("Chart decoded successfully")
@@ -135,28 +135,20 @@ def query_database(question, dashboard_mode=False, chart_type=None):
                                 print(f"Error decoding chart: {e}")
                                 chart_error = f"Chart decoding error: {str(e)}"
                         
-                        # If no image, check for HTML fallback
-                        if not chart_image and "html" in visualization:
-                            chart_html = visualization["html"]
-                            print("HTML chart fallback available")
-                            chart_format = "html"
-                            show_image = False
-                            show_html = True
-                        
                         # Extract chart metadata
-                        chart_title = visualization.get("title", "")
-                        chart_type_result = visualization.get("type", "")
-                        chart_format = visualization.get("format", chart_format)
+                        chart_title = chart.get("title", "")
+                        chart_type_result = chart.get("type", "")
+                        chart_format = chart.get("format", chart_format)
                         print(f"Chart type: {chart_type_result}")
                         print(f"Chart title: {chart_title}")
                         print(f"Chart format: {chart_format}")
-                elif isinstance(visualization, str):
+                elif isinstance(chart, str):
                     # Fallback for backward compatibility
                     try:
                         # Handle base64 prefix
-                        if visualization.startswith("data:image/"):
-                            visualization = visualization.split(",")[1]
-                        image_bytes = base64.b64decode(visualization)
+                        if chart.startswith("data:image/"):
+                            chart = chart.split(",")[1]
+                        image_bytes = base64.b64decode(chart)
                         chart_image = Image.open(BytesIO(image_bytes))
                         print("Chart decoded successfully (fallback)")
                         chart_format = "png"
@@ -166,22 +158,26 @@ def query_database(question, dashboard_mode=False, chart_type=None):
                         print(f"Error decoding chart (fallback): {e}")
                         chart_error = f"Chart decoding error: {str(e)}"
             except Exception as e:
-                print(f"Error processing visualization: {e}")
-                chart_error = f"Visualization processing error: {str(e)}"
+                print(f"Error processing chart: {e}")
+                chart_error = f"Chart processing error: {str(e)}"
+        elif not chart_generated:
+            print("Chart was not generated")
+            chart_error = "Chart generation was not successful"
+        else:
+            print("No chart data available")
+            chart_error = "No chart data available"
         
         # 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 visualization else 'No'}\n"
+        details += f"Chart generated: {chart_generated}\n"
         if chart_type_result:
             details += f"Chart type: {chart_type_result}\n"
         if chart_title:
             details += f"Chart title: {chart_title}\n"
         if chart_format:
             details += f"Chart format: {chart_format}\n"
-        if dashboard_mode:
-            details += f"Dashboard mode: Enabled\n"
         details += f"Response size: {len(response.text)} bytes"
         
         print(f"=== REQUEST COMPLETED SUCCESSFULLY ===")
@@ -417,12 +413,7 @@ with gr.Blocks(theme=theme, title="Enterprise SQL Assistant", css="""
             )
             
             with gr.Row():
-                dashboard_checkbox = gr.Checkbox(
-                    label="Dashboard Mode",
-                    value=False,
-                    info="Generate multiple charts (only primary will be shown)"
-                )
-                
+                # Removed dashboard_mode as backend no longer supports it
                 chart_type_dropdown = gr.Dropdown(
                     label="Chart Type (Optional)",
                     choices=[
@@ -529,10 +520,10 @@ with gr.Blocks(theme=theme, title="Enterprise SQL Assistant", css="""
             outputs=[question_input]
         )
     
-    # Submit button
+    # Submit button - Updated to remove dashboard_mode parameter
     submit_btn.click(
         fn=query_database,
-        inputs=[question_input, dashboard_checkbox, chart_type_dropdown],
+        inputs=[question_input, chart_type_dropdown],  # Removed dashboard_mode
         outputs=[sql_output, results_output, chart_output, html_output, gr.Number(visible=False, value=1), gr.Number(visible=False, value=0), chatbot_output, request_details, error_details, chart_type_output, chart_title_output, chart_format_output, chart_error_output]
     )