Update app.py

app.py

@@ -1,13 +1,8 @@
-import gradio as gr
 from groq import Groq
 from pydantic import BaseModel
 import json
-import sqlite3
-import re
-from datetime import datetime, timedelta
-import random
+import gradio as gr
 
-# Pydantic models for structured output
 class ValidationStatus(BaseModel):
     is_valid: bool
     syntax_errors: list[str]
@@ -19,186 +14,20 @@ class SQLQueryGeneration(BaseModel):
     estimated_complexity: str
     execution_notes: list[str]
     validation_status: ValidationStatus
+    table_schema: str
+    sample_data: str
+    execution_results: str
+    optimization_notes: list[str]
 
-def generate_sample_data(query, tables_used):
-    """Generate sample data based on the query and tables used"""
-    conn = sqlite3.connect(':memory:')
-    cursor = conn.cursor()
-    
-    sample_data = {}
-    
-    # Generate data based on common table patterns
-    if 'customers' in tables_used:
-        cursor.execute('''
-            CREATE TABLE customers (
-                customer_id INTEGER PRIMARY KEY,
-                name TEXT,
-                email TEXT
-            )
-        ''')
-        
-        customers = [
-            (1, 'Alice Johnson', 'alice@example.com'),
-            (2, 'Bob Smith', 'bob@example.com'),
-            (3, 'Carol Williams', 'carol@example.com'),
-            (4, 'David Brown', 'david@example.com'),
-            (5, 'Eve Davis', 'eve@example.com')
-        ]
-        
-        cursor.executemany('INSERT INTO customers VALUES (?, ?, ?)', customers)
-        sample_data['customers'] = customers
-    
-    if 'orders' in tables_used:
-        cursor.execute('''
-            CREATE TABLE orders (
-                order_id INTEGER PRIMARY KEY,
-                customer_id INTEGER,
-                total_amount REAL,
-                order_date TEXT
-            )
-        ''')
-        
-        today = datetime.now()
-        orders = [
-            (101, 1, 600, (today - timedelta(days=10)).strftime('%Y-%m-%d')),
-            (102, 1, 450, (today - timedelta(days=5)).strftime('%Y-%m-%d')),
-            (103, 2, 1200, (today - timedelta(days=15)).strftime('%Y-%m-%d')),
-            (104, 3, 300, (today - timedelta(days=20)).strftime('%Y-%m-%d')),
-            (105, 3, 800, (today - timedelta(days=2)).strftime('%Y-%m-%d')),
-            (106, 4, 550, (today - timedelta(days=7)).strftime('%Y-%m-%d')),
-            (107, 5, 1500, (today - timedelta(days=12)).strftime('%Y-%m-%d'))
-        ]
-        
-        cursor.executemany('INSERT INTO orders VALUES (?, ?, ?, ?)', orders)
-        sample_data['orders'] = orders
-    
-    if 'products' in tables_used:
-        cursor.execute('''
-            CREATE TABLE products (
-                product_id INTEGER PRIMARY KEY,
-                product_name TEXT,
-                price REAL,
-                category TEXT
-            )
-        ''')
-        
-        products = [
-            (1, 'Laptop', 999.99, 'Electronics'),
-            (2, 'Mouse', 29.99, 'Electronics'),
-            (3, 'Keyboard', 79.99, 'Electronics'),
-            (4, 'Monitor', 299.99, 'Electronics'),
-            (5, 'Desk', 199.99, 'Furniture')
-        ]
-        
-        cursor.executemany('INSERT INTO products VALUES (?, ?, ?, ?)', products)
-        sample_data['products'] = products
-    
-    if 'employees' in tables_used:
-        cursor.execute('''
-            CREATE TABLE employees (
-                employee_id INTEGER PRIMARY KEY,
-                name TEXT,
-                department TEXT,
-                salary REAL
-            )
-        ''')
-        
-        employees = [
-            (1, 'John Doe', 'Engineering', 85000),
-            (2, 'Jane Smith', 'Marketing', 75000),
-            (3, 'Mike Johnson', 'Sales', 70000),
-            (4, 'Sarah Williams', 'Engineering', 90000),
-            (5, 'Tom Brown', 'HR', 65000)
-        ]
-        
-        cursor.executemany('INSERT INTO employees VALUES (?, ?, ?, ?)', employees)
-        sample_data['employees'] = employees
-    
-    return conn, cursor, sample_data
-
-def execute_sql_query(cursor, query):
-    """Execute the SQL query and return results"""
+def generate_sql_query(api_key, user_query):
+    """Generate SQL query from natural language using GROQ API"""
     try:
-        # Convert MySQL/PostgreSQL specific functions to SQLite
-        sqlite_query = query.replace('DATE_SUB(NOW(), INTERVAL 30 DAY)', 
-                                     f"date('now', '-30 days')")
-        sqlite_query = sqlite_query.replace('NOW()', "date('now')")
-        
-        cursor.execute(sqlite_query)
-        results = cursor.fetchall()
-        columns = [description[0] for description in cursor.description]
-        return results, columns, None
-    except Exception as e:
-        return None, None, str(e)
-
-def format_sample_tables(sample_data):
-    """Format sample tables as HTML for display"""
-    html = "<div style='margin: 20px 0;'>"
-    
-    for table_name, data in sample_data.items():
-        html += f"<h3>📊 Sample {table_name} Table</h3>"
-        html += "<table style='border-collapse: collapse; width: 100%; margin-bottom: 20px;'>"
-        
-        if table_name == 'customers':
-            html += "<tr style='background-color: #f0f0f0;'><th style='border: 1px solid #ddd; padding: 8px;'>customer_id</th><th style='border: 1px solid #ddd; padding: 8px;'>name</th><th style='border: 1px solid #ddd; padding: 8px;'>email</th></tr>"
-            for row in data:
-                html += f"<tr><td style='border: 1px solid #ddd; padding: 8px;'>{row[0]}</td><td style='border: 1px solid #ddd; padding: 8px;'>{row[1]}</td><td style='border: 1px solid #ddd; padding: 8px;'>{row[2]}</td></tr>"
-        
-        elif table_name == 'orders':
-            html += "<tr style='background-color: #f0f0f0;'><th style='border: 1px solid #ddd; padding: 8px;'>order_id</th><th style='border: 1px solid #ddd; padding: 8px;'>customer_id</th><th style='border: 1px solid #ddd; padding: 8px;'>total_amount</th><th style='border: 1px solid #ddd; padding: 8px;'>order_date</th></tr>"
-            for row in data:
-                html += f"<tr><td style='border: 1px solid #ddd; padding: 8px;'>{row[0]}</td><td style='border: 1px solid #ddd; padding: 8px;'>{row[1]}</td><td style='border: 1px solid #ddd; padding: 8px;'>{row[2]}</td><td style='border: 1px solid #ddd; padding: 8px;'>{row[3]}</td></tr>"
-        
-        elif table_name == 'products':
-            html += "<tr style='background-color: #f0f0f0;'><th style='border: 1px solid #ddd; padding: 8px;'>product_id</th><th style='border: 1px solid #ddd; padding: 8px;'>product_name</th><th style='border: 1px solid #ddd; padding: 8px;'>price</th><th style='border: 1px solid #ddd; padding: 8px;'>category</th></tr>"
-            for row in data:
-                html += f"<tr><td style='border: 1px solid #ddd; padding: 8px;'>{row[0]}</td><td style='border: 1px solid #ddd; padding: 8px;'>{row[1]}</td><td style='border: 1px solid #ddd; padding: 8px;'>{row[2]}</td><td style='border: 1px solid #ddd; padding: 8px;'>{row[3]}</td></tr>"
+        if not api_key:
+            return "Error: Please enter your GROQ API key", "", "", "", "", ""
         
-        elif table_name == 'employees':
-            html += "<tr style='background-color: #f0f0f0;'><th style='border: 1px solid #ddd; padding: 8px;'>employee_id</th><th style='border: 1px solid #ddd; padding: 8px;'>name</th><th style='border: 1px solid #ddd; padding: 8px;'>department</th><th style='border: 1px solid #ddd; padding: 8px;'>salary</th></tr>"
-            for row in data:
-                html += f"<tr><td style='border: 1px solid #ddd; padding: 8px;'>{row[0]}</td><td style='border: 1px solid #ddd; padding: 8px;'>{row[1]}</td><td style='border: 1px solid #ddd; padding: 8px;'>{row[2]}</td><td style='border: 1px solid #ddd; padding: 8px;'>{row[3]}</td></tr>"
+        if not user_query:
+            return "Error: Please enter a query description", "", "", "", "", ""
         
-        html += "</table>"
-    
-    html += "</div>"
-    return html
-
-def format_execution_result(results, columns):
-    """Format SQL execution results as HTML table"""
-    if not results:
-        return "<p>No results found.</p>"
-    
-    html = "<div style='margin: 20px 0;'>"
-    html += "<h3>✅ SQL Execution Result (Final Output Table)</h3>"
-    html += "<table style='border-collapse: collapse; width: 100%;'>"
-    
-    # Header
-    html += "<tr style='background-color: #4CAF50; color: white;'>"
-    for col in columns:
-        html += f"<th style='border: 1px solid #ddd; padding: 8px;'>{col}</th>"
-    html += "</tr>"
-    
-    # Rows
-    for row in results:
-        html += "<tr>"
-        for cell in row:
-            html += f"<td style='border: 1px solid #ddd; padding: 8px;'>{cell}</td>"
-        html += "</tr>"
-    
-    html += "</table></div>"
-    return html
-
-def process_query(api_key, user_query):
-    """Main function to process natural language query"""
-    if not api_key:
-        return "❌ Please enter your Groq API key", "", "", ""
-    
-    if not user_query:
-        return "❌ Please enter a query", "", "", ""
-    
-    try:
-        # Step 1: Generate SQL using Groq
         client = Groq(api_key=api_key)
         
         response = client.chat.completions.create(
@@ -206,9 +35,24 @@ def process_query(api_key, user_query):
             messages=[
                 {
                     "role": "system",
-                    "content": "You are a SQL expert. Generate structured SQL queries from natural language descriptions with proper syntax validation and metadata. Use standard SQL syntax compatible with MySQL.",
+                    "content": """You are a SQL expert. Generate structured SQL queries from natural language descriptions with proper syntax validation and metadata.
+After generating the SQL query, you must:
+1. Create a sample SQL table schema based on the natural language description, including all necessary columns with appropriate data types
+2. Populate the table with realistic sample data that demonstrates the query's functionality
+3. Execute the generated SQL query against the sample table
+4. Display the SQL table structure and data clearly
+5. Show the query execution results
+Always present your response in this order:
+- Generated SQL query with syntax explanation
+- Table schema (CREATE TABLE statement)
+- Sample data (INSERT statements or table visualization)
+- Query execution results
+- Any relevant notes about assumptions made or query optimization suggestions""",
+                },
+                {
+                    "role": "user", 
+                    "content": user_query
                 },
-                {"role": "user", "content": user_query},
             ],
             response_format={
                 "type": "json_schema",
@@ -223,125 +67,145 @@ def process_query(api_key, user_query):
             json.loads(response.choices[0].message.content)
         )
         
-        # Step 2: Format the structured output
-        step1_output = f"""
-## 🎯 Step 1: Understand User Intent
-
-**User Query:** "{user_query}"
-
-**Identified Components:**
-- **Tables:** {', '.join(sql_query_generation.tables_used)}
-- **Query Type:** {sql_query_generation.query_type}
-- **Complexity:** {sql_query_generation.estimated_complexity}
-"""
-        
-        step2_output = f"""
-## 🔧 Step 2: Generate Structured SQL
-
-```json
-{json.dumps(sql_query_generation.model_dump(), indent=2)}
-```
-
-**Generated SQL Query:**
-```sql
-{sql_query_generation.query}
-```
-"""
-        
-        # Step 3: Generate sample data
-        conn, cursor, sample_data = generate_sample_data(
-            sql_query_generation.query, 
-            sql_query_generation.tables_used
-        )
-        
-        step3_output = f"""
-## 📊 Step 3: Auto-Generate Sample Database Tables
-
-{format_sample_tables(sample_data)}
-"""
-        
-        # Step 4: Execute query
-        results, columns, error = execute_sql_query(cursor, sql_query_generation.query)
-        
-        if error:
-            step4_output = f"""
-## ⚠️ Step 4: SQL Execution
-
-**Error:** {error}
-
-**Note:** The query might use database-specific functions. The sample execution uses SQLite.
-"""
+        # Format validation status
+        validation_text = f"Valid: {sql_query_generation.validation_status.is_valid}\n"
+        if sql_query_generation.validation_status.syntax_errors:
+            validation_text += "Errors:\n" + "\n".join(
+                f"- {error}" for error in sql_query_generation.validation_status.syntax_errors
+            )
         else:
-            step4_output = f"""
-## 🚀 Step 4: Execute Generated SQL on Sample Tables
-
-**Applied Conditions:**
-{chr(10).join([f"- {note}" for note in sql_query_generation.execution_notes])}
-
-{format_execution_result(results, columns)}
-
-**Total Rows Returned:** {len(results)}
-"""
-        
-        conn.close()
-        
-        return step1_output, step2_output, step3_output, step4_output
+            validation_text += "No syntax errors found"
+        
+        # Format metadata
+        metadata = f"""Query Type: {sql_query_generation.query_type}
+Tables Used: {', '.join(sql_query_generation.tables_used)}
+Complexity: {sql_query_generation.estimated_complexity}
+
+Execution Notes:
+{chr(10).join(f"- {note}" for note in sql_query_generation.execution_notes)}
+
+Optimization Notes:
+{chr(10).join(f"- {note}" for note in sql_query_generation.optimization_notes)}"""
+        
+        return (
+            sql_query_generation.query,
+            metadata,
+            sql_query_generation.table_schema,
+            sql_query_generation.sample_data,
+            sql_query_generation.execution_results,
+            validation_text
+        )
         
     except Exception as e:
-        error_msg = f"❌ **Error:** {str(e)}"
-        return error_msg, "", "", ""
+        error_msg = f"Error: {str(e)}"
+        return error_msg, "", "", "", "", ""
 
 # Create Gradio interface
-with gr.Blocks(title="Natural Language to SQL Query Executor", theme=gr.themes.Soft()) as app:
-    gr.Markdown("""
-    # 🔍 Natural Language to SQL Query Executor
-    
-    Convert natural language queries to SQL, auto-generate sample data, and execute queries!
-    
-    **Example queries to try:**
-    - "Find all customers who made orders over $500 in the last 30 days, show their name, email, and total order amount"
-    - "Get all employees in the Engineering department with salary above 80000"
-    - "Show top 5 products by price"
-    """)
+with gr.Blocks(title="SQL Query Generator", theme=gr.themes.Soft()) as demo:
+    gr.Markdown(
+        """
+        # 🗄️ Natural Language to SQL Query Generator
+        Convert your natural language descriptions into structured SQL queries with validation and execution results.
+        """
+    )
     
     with gr.Row():
         with gr.Column():
             api_key_input = gr.Textbox(
-                label="🔑 Groq API Key",
-                placeholder="Enter your Groq API key here...",
-                type="password"
+                label="GROQ API Key",
+                type="password",
+                placeholder="Enter your GROQ API key here...",
+                info="Your API key is not stored and only used for this session"
             )
+            
             query_input = gr.Textbox(
-                label="💬 Natural Language Query",
-                placeholder="Enter your query in plain English...",
+                label="Natural Language Query",
+                placeholder="e.g., Find all the students who scored more than 90 out of 100",
+                lines=3,
+                value="Find all the students who scored more than 90 out of 100"
+            )
+            
+            generate_btn = gr.Button("Generate SQL Query", variant="primary", size="lg")
+            
+            gr.Examples(
+                examples=[
+                    ["Find all the students who scored more than 90 out of 100"],
+                    ["Get the top 5 customers by total purchase amount"],
+                    ["List all employees hired in the last 6 months"],
+                    ["Find products with price between $50 and $100"],
+                    ["Show average salary by department"]
+                ],
+                inputs=query_input,
+                label="Example Queries"
+            )
+    
+    with gr.Row():
+        with gr.Column():
+            sql_output = gr.Code(
+                label="Generated SQL Query",
+                language="sql",
+                lines=5
+            )
+            
+            metadata_output = gr.Textbox(
+                label="Query Metadata",
+                lines=8
+            )
+            
+            validation_output = gr.Textbox(
+                label="Validation Status",
                 lines=3
             )
-            submit_btn = gr.Button("🚀 Generate & Execute SQL", variant="primary", size="lg")
     
     with gr.Row():
         with gr.Column():
-            step1_output = gr.Markdown(label="Step 1: Understanding")
-            step2_output = gr.Markdown(label="Step 2: SQL Generation")
-            step3_output = gr.HTML(label="Step 3: Sample Data")
-            step4_output = gr.HTML(label="Step 4: Execution Results")
+            schema_output = gr.Code(
+                label="Table Schema",
+                language="sql",
+                lines=8
+            )
+        
+        with gr.Column():
+            sample_data_output = gr.Code(
+                label="Sample Data",
+                language="sql",
+                lines=8
+            )
+    
+    with gr.Row():
+        execution_output = gr.Textbox(
+            label="Execution Results",
+            lines=10
+        )
     
-    submit_btn.click(
-        fn=process_query,
+    generate_btn.click(
+        fn=generate_sql_query,
         inputs=[api_key_input, query_input],
-        outputs=[step1_output, step2_output, step3_output, step4_output]
+        outputs=[
+            sql_output,
+            metadata_output,
+            schema_output,
+            sample_data_output,
+            execution_output,
+            validation_output
+        ]
     )
     
-    gr.Markdown("""
-    ---
-    ### 📝 How it works:
-    1. **Understand Intent:** Analyzes your natural language query
-    2. **Generate SQL:** Creates structured SQL with metadata
-    3. **Create Sample Data:** Auto-generates realistic sample tables
-    4. **Execute & Display:** Runs the query and shows results
-    
-    ### 🔗 Get your Groq API key:
-    Visit [console.groq.com](https://console.groq.com) to get your free API key!
-    """)
+    gr.Markdown(
+        """
+        ---
+        ### How to use:
+        1. Enter your GROQ API key (get one from [console.groq.com](https://console.groq.com))
+        2. Type your natural language query description
+        3. Click "Generate SQL Query" to see the results
+        
+        The app will provide:
+        - A validated SQL query
+        - Table schema and sample data
+        - Execution results
+        - Optimization suggestions
+        """
+    )
 
 if __name__ == "__main__":
-    app.launch()
+    demo.launch(share=True)