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://52.15.175.95:5000"
 
@@ -27,12 +28,12 @@ def query_database(question, dashboard_mode=False, chart_type=None):
         # Prepare the request
         payload = {
             "question": question,
-            "dashboard": dashboard_mode
+            "visualize": True  # Always request visualization
         }
         
         # Add chart_type if specified
-        if chart_type:
-            payload["chart_type"] = chart_type
+        if chart_type and chart_type != "auto":
+            payload["viz_type"] = chart_type
         
         print(f"Request URL: {FLASK_API_URL}/ask")
         print(f"Request payload: {payload}")
@@ -76,63 +77,61 @@ def query_database(question, dashboard_mode=False, chart_type=None):
         # Extract data
         sql = result.get("sql", "")
         rows = result.get("result", [])
-        chart_data = result.get("chart")
+        visualization = result.get("visualization", {})
         row_count = result.get("row_count", len(rows))
-        elapsed_seconds = result.get("elapsed_seconds", elapsed_time)
         
         print(f"\n=== EXTRACTED DATA ===")
         print(f"SQL: {sql[:100] if sql else 'None'}...")
         print(f"Rows count: {row_count}")
-        print(f"Chart data received: {'Yes' if chart_data else 'No'}")
-        print(f"Elapsed seconds: {elapsed_seconds}")
+        print(f"Visualization received: {'Yes' if visualization else 'No'}")
         
         # Create DataFrame
         df = pd.DataFrame(rows) if rows else pd.DataFrame()
         print(f"DataFrame shape: {df.shape}")
         
-        # Process chart - UPDATED TO HANDLE NEW STRUCTURED FORMAT
+        # Process visualization
         chart_image = None
         chart_title = ""
         chart_type_result = ""
         chart_error = None
         
-        if chart_data:
+        if visualization:
             try:
-                if isinstance(chart_data, dict):
+                if isinstance(visualization, dict):
                     # Check if it's an error response
-                    if "error" in chart_data:
-                        chart_error = chart_data["error"]
+                    if "error" in visualization:
+                        chart_error = visualization["error"]
                         print(f"Chart error: {chart_error}")
                     else:
-                        # Extract the base64 image string from the structured chart data
-                        chart_image_b64 = chart_data.get("image")
+                        # Extract the base64 image string from the structured visualization data
+                        chart_image_b64 = visualization.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_result = chart_data.get("type", "")
+                        chart_title = visualization.get("title", "")
+                        chart_type_result = visualization.get("type", "")
                         print(f"Chart type: {chart_type_result}")
                         print(f"Chart title: {chart_title}")
-                elif isinstance(chart_data, str):
+                elif isinstance(visualization, str):
                     # Fallback for backward compatibility
                     try:
-                        image_bytes = base64.b64decode(chart_data)
+                        image_bytes = base64.b64decode(visualization)
                         chart_image = Image.open(BytesIO(image_bytes))
                         print("Chart decoded successfully (fallback)")
                     except Exception as e:
                         print(f"Error decoding chart (fallback): {e}")
             except Exception as e:
-                print(f"Error processing chart: {e}")
-                chart_error = f"Chart processing error: {str(e)}"
+                print(f"Error processing visualization: {e}")
+                chart_error = f"Visualization processing error: {str(e)}"
         
         # Prepare details
-        details = f"Request time: {elapsed_seconds:.2f}s\n"
+        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"
+        details += f"Chart generated: {'Yes' if visualization else 'No'}\n"
         if chart_type_result:
             details += f"Chart type: {chart_type_result}\n"
         if chart_title:
@@ -173,14 +172,10 @@ def check_health():
         if response.status_code == 200:
             health_data = response.json()
             status = health_data.get('status', 'unknown')
-            model = health_data.get('model', 'unknown')
             tables = health_data.get('tables', [])
-            data_rows = health_data.get('data_rows', 0)
             
             health_msg = f"✅ API Status: {status.upper()}\n"
-            health_msg += f"🤖 Model: {model}\n"
             health_msg += f"📊 Tables: {', '.join(tables) if tables else 'None'}\n"
-            health_msg += f"📈 Data Rows: {data_rows:,}"
             
             return health_msg, "success"
         else:
@@ -191,21 +186,18 @@ def check_health():
 def get_schema():
     try:
         print("Fetching database schema...")
-        response = requests.get(f"{FLASK_API_URL}/schema", timeout=10)
+        response = requests.get(f"{FLASK_API_URL}/tables", timeout=10)
         
         if response.status_code == 200:
-            schema_data = response.json()
-            schema = schema_data.get("schema", {})
-            tables = schema_data.get("tables", [])
+            tables_data = response.json()
+            tables = tables_data.get("tables", [])
             
             schema_text = "## Database Schema\n\n"
             for table in tables:
-                schema_text += f"### {table}\n"
-                if table in schema:
-                    columns = schema[table]
-                    schema_text += "| Column | Type |\n|--------|------|\n"
-                    for col in columns:
-                        schema_text += f"| {col['name']} | {col['type']} |\n"
+                schema_text += f"### {table['name']}\n"
+                schema_text += "| Column |\n|--------|\n"
+                for col in table['columns']:
+                    schema_text += f"| {col} |\n"
                 schema_text += "\n"
             
             return schema_text, "success"
@@ -331,8 +323,7 @@ with gr.Blocks(theme=theme, title="Enterprise SQL Assistant", css="""
                 chart_type_dropdown = gr.Dropdown(
                     label="Chart Type (Optional)",
                     choices=[
-                        "auto", "time_series", "correlation", "histogram", 
-                        "bar", "scatter", "pie", "table"
+                        "auto", "bar", "line", "scatter", "pie", "histogram"
                     ],
                     value="auto",
                     info="Force a specific chart type"
@@ -348,10 +339,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"])  # 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
+                    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 distribution of transaction amounts", elem_classes=["example-btn"])
+                    example8 = gr.Button("Show correlations between numeric fields", elem_classes=["example-btn"])
             
             with gr.Accordion("Advanced Options", open=False):
                 gr.Markdown("### 🩺 System Information")