Update app.py

app.py

@@ -7,13 +7,12 @@ import json
 import time
 import traceback
 from PIL import Image
-
 # 🔧 CONFIGURE: Your Flask API URL (Public URL)
-FLASK_API_URL = "http://18.219.185.9:5000"
+FLASK_API_URL = "http://52.15.175.95:5000"
 
-def query_database(question):
+def query_database(question, dashboard_mode=False, chart_type=None):
     if not question.strip():
-        return "", None, None, "⚠️ Please enter a question.", "", "", "", ""
+        return "", None, None, "⚠️ Please enter a question.", "", "", "", "", ""
     
     try:
         start_time = time.time()
@@ -21,10 +20,19 @@ def query_database(question):
         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)
         
         # Prepare the request
-        payload = {"question": question}
+        payload = {
+            "question": question,
+            "dashboard": dashboard_mode
+        }
+        
+        # Add chart_type if specified
+        if chart_type:
+            payload["chart_type"] = chart_type
         
         print(f"Request URL: {FLASK_API_URL}/ask")
         print(f"Request payload: {payload}")
@@ -47,7 +55,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,24 +65,26 @@ 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", "")
         rows = result.get("result", [])
         chart_data = result.get("chart")
         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}")
         
         # Create DataFrame
         df = pd.DataFrame(rows) if rows else pd.DataFrame()
@@ -83,67 +93,77 @@ def query_database(question):
         # Process chart - UPDATED TO HANDLE NEW STRUCTURED FORMAT
         chart_image = None
         chart_title = ""
-        chart_type = ""
+        chart_type_result = ""
+        chart_error = None
         
-        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
+        if chart_data:
             try:
-                image_bytes = base64.b64decode(chart_data)
-                chart_image = Image.open(BytesIO(image_bytes))
-                print("Chart decoded successfully (fallback)")
+                if isinstance(chart_data, dict):
+                    # Check if it's an error response
+                    if "error" in chart_data:
+                        chart_error = chart_data["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")
+                        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", "")
+                        print(f"Chart type: {chart_type_result}")
+                        print(f"Chart title: {chart_title}")
+                elif 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 (fallback)")
+                    except Exception as e:
+                        print(f"Error decoding chart (fallback): {e}")
             except Exception as e:
-                print(f"Error decoding chart (fallback): {e}")
+                print(f"Error processing chart: {e}")
+                chart_error = f"Chart processing error: {str(e)}"
         
         # Prepare details
-        details = f"Request time: {elapsed_time:.2f}s\n"
+        details = f"Request time: {elapsed_seconds:.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_type_result:
+            details += f"Chart type: {chart_type_result}\n"
         if chart_title:
             details += f"Chart title: {chart_title}\n"
+        if dashboard_mode:
+            details += f"Dashboard mode: Enabled\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", chart_type, chart_title
+        return sql, df, chart_image, f"✅ Query completed successfully", details, "Success", chart_type_result, chart_title, chart_error
     
     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:
@@ -168,6 +188,32 @@ def check_health():
     except Exception as e:
         return f"❌ Health check failed: {str(e)}", "error"
 
+def get_schema():
+    try:
+        print("Fetching database schema...")
+        response = requests.get(f"{FLASK_API_URL}/schema", timeout=10)
+        
+        if response.status_code == 200:
+            schema_data = response.json()
+            schema = schema_data.get("schema", {})
+            tables = schema_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 += "\n"
+            
+            return schema_text, "success"
+        else:
+            return f"❌ API returned status {response.status_code}", "error"
+    except Exception as e:
+        return f"❌ Failed to fetch schema: {str(e)}", "error"
+
 # 🎨 Theme: Enterprise Dark Blue
 theme = gr.themes.Default(
     primary_hue="blue",
@@ -238,6 +284,14 @@ with gr.Blocks(theme=theme, title="Enterprise SQL Assistant", css="""
         padding: 10px;
         margin-bottom: 10px;
     }
+    .schema-container {
+        background-color: #1e293b;
+        border: 1px solid #334155;
+        border-radius: 8px;
+        padding: 15px;
+        max-height: 400px;
+        overflow-y: auto;
+    }
     /* Fix input text color */
     .gradio-container input[type="text"], 
     .gradio-container textarea {
@@ -266,6 +320,24 @@ with gr.Blocks(theme=theme, title="Enterprise SQL Assistant", css="""
                 label="Natural Language Query",
                 lines=3
             )
+            
+            with gr.Row():
+                dashboard_checkbox = gr.Checkbox(
+                    label="Dashboard Mode",
+                    value=False,
+                    info="Generate multiple charts (only primary will be shown)"
+                )
+                
+                chart_type_dropdown = gr.Dropdown(
+                    label="Chart Type (Optional)",
+                    choices=[
+                        "auto", "time_series", "correlation", "histogram", 
+                        "bar", "scatter", "pie", "table"
+                    ],
+                    value="auto",
+                    info="Force a specific chart type"
+                )
+            
             submit_btn = gr.Button("🚀 Generate SQL & Results", variant="primary", size="lg")
             
             gr.Markdown("### 💡 Example Queries")
@@ -281,9 +353,14 @@ with gr.Blocks(theme=theme, title="Enterprise SQL Assistant", css="""
                     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")
-            health_output = gr.Markdown(label="API Status")
+            with gr.Accordion("Advanced Options", open=False):
+                gr.Markdown("### 🩺 System Information")
+                with gr.Row():
+                    health_btn = gr.Button("Check API Health", variant="secondary", size="sm")
+                    schema_btn = gr.Button("Get Database Schema", variant="secondary", size="sm")
+                
+                health_output = gr.Markdown(label="API Status")
+                schema_output = gr.Markdown(label="Database Schema", elem_classes="schema-container")
         
         with gr.Column(scale=2):
             gr.Markdown("### 🤖 AI Assistant Response")
@@ -316,6 +393,7 @@ with gr.Blocks(theme=theme, title="Enterprise SQL Assistant", css="""
                             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")
+                            chart_error_output = gr.Markdown(label="Chart Error", elem_classes="chart-metadata")
             
             gr.Markdown("### 📊 Request Details")
             request_details = gr.Textbox(label="Request Details", interactive=False, lines=6)
@@ -335,22 +413,26 @@ with gr.Blocks(theme=theme, title="Enterprise SQL Assistant", css="""
         outputs=[health_output, error_details]
     )
     
+    # Schema button
+    schema_btn.click(
+        fn=get_schema,
+        inputs=[],
+        outputs=[schema_output, error_details]
+    )
+    
     # Example query buttons
     for example_btn in [example1, example2, example3, example4, example5, example6, example7, example8]:
         example_btn.click(
             fn=set_example_query,
             inputs=[example_btn],
             outputs=[question_input]
-        ).then(
-            fn=query_database,
-            inputs=question_input,
-            outputs=[sql_output, results_output, chart_output, chatbot_output, request_details, error_details, chart_type_output, chart_title_output]
         )
     
+    # Submit button
     submit_btn.click(
         fn=query_database,
-        inputs=question_input,
-        outputs=[sql_output, results_output, chart_output, chatbot_output, request_details, error_details, chart_type_output, chart_title_output]
+        inputs=[question_input, dashboard_checkbox, chart_type_dropdown],
+        outputs=[sql_output, results_output, chart_output, chatbot_output, request_details, error_details, chart_type_output, chart_title_output, chart_error_output]
     )
 
 # Launch