Update app.py

app.py

@@ -6,71 +6,10 @@ from io import BytesIO
 import json
 import time
 import traceback
+from PIL import Image
 
-# 🔧 CONFIGURE: Replace with your EC2 public IP
-API_URL = "http://3.14.153.135:5000/ask"
-HEALTH_URL = "http://3.14.153.135:5000/health"
-SIMPLE_TEST_URL = "http://3.14.153.135:5000/simple-test"
-DEBUG_URL = "http://3.14.153.135:5000/debug"
-
-def test_connection():
-    try:
-        print("=== TESTING CONNECTION ===")
-        response = requests.get(HEALTH_URL, timeout=30)
-        print(f"Connection test status: {response.status_code}")
-        print(f"Connection test response: {response.text}")
-        
-        if response.status_code == 200:
-            return f"✅ Connection successful! Status: {response.status_code}\nResponse: {response.text}"
-        else:
-            return f"⚠️ Connection returned status: {response.status_code}\nResponse: {response.text}"
-    except Exception as e:
-        error_msg = f"Connection error: {str(e)}"
-        print(f"Connection error: {error_msg}")
-        print(f"Traceback: {traceback.format_exc()}")
-        return f"❌ {error_msg}"
-
-def test_simple_endpoint():
-    try:
-        print("=== TESTING SIMPLE ENDPOINT ===")
-        response = requests.post(
-            SIMPLE_TEST_URL,
-            json={"question": "test question"},
-            timeout=30
-        )
-        
-        print(f"Simple test status: {response.status_code}")
-        print(f"Simple test response: {response.text}")
-        
-        if response.status_code == 200:
-            result = response.json()
-            return f"✅ Simple test successful! SQL: {result.get('sql', 'N/A')}", None
-        else:
-            return f"❌ Simple test failed: {response.status_code} - {response.text}", None
-    except Exception as e:
-        error_msg = f"Simple test error: {str(e)}"
-        print(f"Simple test error: {error_msg}")
-        print(f"Traceback: {traceback.format_exc()}")
-        return f"❌ {error_msg}", None
-
-def test_debug_endpoint():
-    try:
-        print("=== TESTING DEBUG ENDPOINT ===")
-        response = requests.post(
-            DEBUG_URL,
-            json={"test": "data"},
-            timeout=30
-        )
-        
-        print(f"Debug test status: {response.status_code}")
-        print(f"Debug test response: {response.text}")
-        
-        return f"Debug response: {response.text}", None
-    except Exception as e:
-        error_msg = f"Debug test error: {str(e)}"
-        print(f"Debug test error: {error_msg}")
-        print(f"Traceback: {traceback.format_exc()}")
-        return f"❌ {error_msg}", None
+# 🔧 CONFIGURE: Replace with your Django server URL
+DJANGO_API_URL = "http://your-django-server.com/api/openai-prompt/"
 
 def query_database(question):
     if not question.strip():
@@ -85,22 +24,16 @@ def query_database(question):
         print("="*50)
         
         # Prepare the request
-        headers = {
-            'Content-Type': 'application/json',
-            'User-Agent': 'HuggingFace-Space/1.0'
-        }
-        payload = {"question": question}
+        payload = {"prompt": question}
         
-        print(f"Request URL: {API_URL}")
-        print(f"Request headers: {headers}")
+        print(f"Request URL: {DJANGO_API_URL}")
         print(f"Request payload: {payload}")
         
         # Make the request
         print("Sending request...")
         response = requests.post(
-            API_URL, 
-            json=payload, 
-            headers=headers,
+            DJANGO_API_URL, 
+            json=payload,
             timeout=300
         )
         
@@ -108,8 +41,6 @@ def query_database(question):
         print(f"\n=== RESPONSE RECEIVED ===")
         print(f"Response time: {elapsed_time:.2f} seconds")
         print(f"Response status: {response.status_code}")
-        print(f"Response headers: {dict(response.headers)}")
-        print(f"Response text length: {len(response.text)} characters")
         print(f"Response text: {response.text}")
         
         # Check if response is empty
@@ -130,26 +61,27 @@ def query_database(question):
         
         # Check HTTP status
         if response.status_code != 200:
-            error_msg = result.get("error", f"HTTP {response.status_code}")
+            error_msg = result.get("message", 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"
         
         # Check for application error
-        if "error" in result:
-            error_msg = result['error']
+        if "status" in result and result["status"] == "error":
+            error_msg = result.get("message", "Unknown error")
             print(f"ERROR: Application error: {error_msg}")
             return "", None, None, f"❌ {error_msg}", f"Request time: {elapsed_time:.2f}s", "Application error"
         
         # Extract data
-        sql = result.get("sql", "N/A")
+        conversational_response = result.get("response", "")
+        sql = result.get("sql", "")
         rows = result.get("result", [])
         chart_data = result.get("chart")
+        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: {len(rows)}")
+        print(f"Rows count: {row_count}")
         print(f"Chart data received: {'Yes' if chart_data else 'No'}")
-        print(f"Chart data length: {len(chart_data) if chart_data else 0}")
         
         # Create DataFrame
         df = pd.DataFrame(rows) if rows else pd.DataFrame()
@@ -160,7 +92,7 @@ def query_database(question):
         if chart_data:
             try:
                 image_bytes = base64.b64decode(chart_data)
-                chart_image = BytesIO(image_bytes)
+                chart_image = Image.open(BytesIO(image_bytes))
                 print("Chart decoded successfully")
             except Exception as e:
                 print(f"Error decoding chart: {e}")
@@ -168,13 +100,13 @@ def query_database(question):
         # Prepare details
         details = f"Request time: {elapsed_time:.2f}s\n"
         details += f"Status code: {response.status_code}\n"
-        details += f"Rows returned: {len(rows)}\n"
+        details += f"Rows returned: {row_count}\n"
         details += f"Chart generated: {'Yes' if chart_data else 'No'}\n"
         details += f"Response size: {len(response.text)} bytes"
         
         print(f"=== REQUEST COMPLETED SUCCESSFULLY ===")
         
-        return sql, df, chart_image, "✅ Query completed successfully!", details, "Success"
+        return sql, df, chart_image, conversational_response, details, "Success"
     
     except requests.exceptions.ConnectionError as e:
         error_msg = f"Connection failed: {str(e)}"
@@ -204,11 +136,64 @@ theme = gr.themes.Default(
     font=["Inter", "sans-serif"]
 ).set(
     body_background_fill="*neutral_950",
-    background_fill_secondary="*neutral_900"
+    background_fill_secondary="*neutral_900",
+    button_primary_background_fill="*primary_600",
+    button_primary_background_fill_hover="*primary_700",
+    button_secondary_background_fill="*neutral_800",
+    button_secondary_background_fill_hover="*neutral_700",
+    block_title_text_color="*primary_400",
+    block_label_text_color="*neutral_300",
+    input_background_fill="*neutral_800",
+    input_border_color="*neutral_700",
+    input_border_width="1px",
+    input_text_color="*neutral_100",
 )
 
 # 🚀 UI Layout
-with gr.Blocks(theme=theme, title="Enterprise SQL Assistant") as demo:
+with gr.Blocks(theme=theme, title="Enterprise SQL Assistant", css="""
+    .example-btn {
+        background-color: #1e40af;
+        color: white;
+        border-radius: 8px;
+        padding: 8px 12px;
+        margin: 4px;
+        font-size: 0.9em;
+        border: none;
+        cursor: pointer;
+        transition: all 0.2s;
+    }
+    .example-btn:hover {
+        background-color: #1d4ed8;
+        transform: translateY(-1px);
+    }
+    .chatbot-container {
+        border: 1px solid #334155;
+        border-radius: 8px;
+        padding: 15px;
+        background-color: #1e293b;
+        min-height: 200px;
+        max-height: 400px;
+        overflow-y: auto;
+    }
+    .status-success {
+        color: #10b981;
+        font-weight: bold;
+    }
+    .status-error {
+        color: #ef4444;
+        font-weight: bold;
+    }
+    .status-warning {
+        color: #f59e0b;
+        font-weight: bold;
+    }
+    .chart-container {
+        border: 1px solid #334155;
+        border-radius: 8px;
+        padding: 10px;
+        background-color: #1e293b;
+    }
+""") as demo:
     gr.HTML("""
     <div style="text-align: center; padding: 20px;">
         <h1 style="color: #3B82F6; margin-bottom: 5px;">📊 Enterprise SQL Assistant</h1>
@@ -226,36 +211,27 @@ with gr.Blocks(theme=theme, title="Enterprise SQL Assistant") as demo:
                 label="Natural Language Query",
                 lines=3
             )
-            submit_btn = gr.Button("🚀 Generate SQL & Results", variant="primary")
-            
-            gr.Markdown("### 🔧 Connection Tests")
-            with gr.Accordion("Connection Tests", open=True):
-                test_btn = gr.Button("Test Connection to EC2")
-                connection_status = gr.Textbox(label="Connection Status", interactive=False)
-                
-                simple_test_btn = gr.Button("Test Simple Endpoint")
-                simple_test_status = gr.Textbox(label="Simple Test Status", interactive=False)
-                
-                debug_test_btn = gr.Button("Test Debug Endpoint")
-                debug_test_status = gr.Textbox(label="Debug Test Status", interactive=False)
+            submit_btn = gr.Button("🚀 Generate SQL & Results", variant="primary", size="lg")
             
             gr.Markdown("### 💡 Example Queries")
-            gr.Markdown("""
-            Try these example queries:
-            
-            - How many members are there?
-            - What is the total transaction amount?
-            - Show members with their account balances
-            - Which member has the highest balance?
-            - Show transaction trends over time
-            - Count of members by status
-            - Show me the first 10 rows
-            - What is the average age of members?
-            """)
+            with gr.Row():
+                with gr.Column():
+                    example1 = gr.Button("How many members are there?", elem_classes=["example-btn"])
+                    example2 = gr.Button("What is the total transaction amount?", elem_classes=["example-btn"])
+                    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"])
         
         with gr.Column(scale=2):
-            gr.Markdown("### 📄 Results")
-            status = gr.Textbox(label="Status", interactive=False)
+            gr.Markdown("### 🤖 AI Assistant Response")
+            chatbot_output = gr.Markdown(
+                label="AI Response",
+                elem_classes=["chatbot-container"]
+            )
             
             with gr.Tabs():
                 with gr.Tab("SQL Query"):
@@ -282,26 +258,27 @@ with gr.Blocks(theme=theme, title="Enterprise SQL Assistant") as demo:
             gr.Markdown("### 🔍 Error Details")
             error_details = gr.Textbox(label="Error Details", interactive=False, lines=4)
     
-    # Events
-    test_btn.click(
-        fn=test_connection,
-        outputs=[connection_status]
-    )
-    
-    simple_test_btn.click(
-        fn=test_simple_endpoint,
-        outputs=[simple_test_status]
-    )
+    # Function to handle example query clicks
+    def set_example_query(example_text):
+        return example_text
     
-    debug_test_btn.click(
-        fn=test_debug_endpoint,
-        outputs=[debug_test_status]
-    )
+    # Events
+    # 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]
+        )
     
     submit_btn.click(
         fn=query_database,
         inputs=question_input,
-        outputs=[sql_output, results_output, chart_output, status, request_details, error_details]
+        outputs=[sql_output, results_output, chart_output, chatbot_output, request_details, error_details]
     )
 
 # Launch