Update app.py

app.py

@@ -25,13 +25,13 @@ def query_database(question, dashboard_mode=False, chart_type=None):
         print(f"Chart type: {chart_type}")
         print("="*50)
         
-        # Prepare the request
+        # Prepare the request - FIXED: Use correct parameter names
         payload = {
             "question": question,
             "visualize": True  # Always request visualization
         }
         
-        # Add chart_type if specified
+        # Add chart_type if specified - FIXED: Use correct parameter name
         if chart_type and chart_type != "auto":
             payload["viz_type"] = chart_type
         
@@ -89,7 +89,7 @@ def query_database(question, dashboard_mode=False, chart_type=None):
         df = pd.DataFrame(rows) if rows else pd.DataFrame()
         print(f"DataFrame shape: {df.shape}")
         
-        # Process visualization
+        # Process visualization - FIXED: Handle base64 prefix
         chart_image = None
         chart_title = ""
         chart_type_result = ""
@@ -106,6 +106,9 @@ def query_database(question, dashboard_mode=False, chart_type=None):
                         # Extract the base64 image string from the structured visualization data
                         chart_image_b64 = visualization.get("image")
                         if chart_image_b64:
+                            # FIXED: Remove base64 prefix if present
+                            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")
@@ -118,6 +121,9 @@ def query_database(question, dashboard_mode=False, chart_type=None):
                 elif isinstance(visualization, str):
                     # Fallback for backward compatibility
                     try:
+                        # FIXED: Remove base64 prefix if present
+                        if visualization.startswith("data:image/"):
+                            visualization = visualization.split(",")[1]
                         image_bytes = base64.b64decode(visualization)
                         chart_image = Image.open(BytesIO(image_bytes))
                         print("Chart decoded successfully (fallback)")
@@ -172,7 +178,7 @@ def check_health():
         if response.status_code == 200:
             health_data = response.json()
             status = health_data.get('status', 'unknown')
-            tables = health_data.get('tables', [])
+            tables = health_data.get('tables', [])  # FIXED: This is now correct
             
             health_msg = f"✅ API Status: {status.upper()}\n"
             health_msg += f"📊 Tables: {', '.join(tables) if tables else 'None'}\n"
@@ -186,6 +192,7 @@ def check_health():
 def get_schema():
     try:
         print("Fetching database schema...")
+        # FIXED: Use correct endpoint
         response = requests.get(f"{FLASK_API_URL}/tables", timeout=10)
         
         if response.status_code == 200:
@@ -323,7 +330,8 @@ with gr.Blocks(theme=theme, title="Enterprise SQL Assistant", css="""
                 chart_type_dropdown = gr.Dropdown(
                     label="Chart Type (Optional)",
                     choices=[
-                        "auto", "bar", "line", "scatter", "pie", "histogram"
+                        "auto", "bar", "line", "scatter", "pie", "histogram",
+                        "time_series", "correlation"  # FIXED: Added missing chart types
                     ],
                     value="auto",
                     info="Force a specific chart type"