Update app.py

app.py

@@ -6,6 +6,8 @@ import sqlite3
 import pandas as pd
 from typing import List, Optional
 import re
+from datetime import datetime, timedelta
+import random
 
 # Pydantic Models
 class ValidationStatus(BaseModel):
@@ -30,9 +32,14 @@ def generate_sample_data(user_query: str, groq_api_key: str) -> dict:
     try:
         client = Groq(api_key=groq_api_key)
         
+        # Get current date for context
+        today = datetime.now().strftime('%Y-%m-%d')
+        
         # Request to generate table schema and sample data
         schema_prompt = f"""Based on this query: "{user_query}"
-        
+
+Current date: {today}
+
 Generate a realistic database schema with sample data. Return ONLY valid JSON with this structure:
 {{
   "tables": [
@@ -44,18 +51,26 @@ Generate a realistic database schema with sample data. Return ONLY valid JSON wi
       ],
       "sample_data": [
         {{"column_name": value, ...}},
-        ...at least 10-15 rows
+        ...at least 15-20 rows
       ]
     }}
   ]
 }}
 
-Make the data realistic and relevant to the query. Include enough variety to make the query results meaningful."""
+IMPORTANT INSTRUCTIONS FOR REALISTIC DATA:
+1. For DATE columns: Use dates in format 'YYYY-MM-DD'. Include dates from the last 60 days to ensure some fall within "last 30 days"
+2. For queries mentioning "last X days": Generate at least 50% of dates within that timeframe
+3. For queries with amount/price filters (e.g., "over $500"): Ensure at least 40% of records meet the criteria
+4. For queries with thresholds: Create data both above AND below the threshold
+5. Use realistic names, emails, and values
+6. Make sure there's enough variety in the data to produce meaningful query results
+
+Example date range: from {(datetime.now() - timedelta(days=60)).strftime('%Y-%m-%d')} to {today}"""
 
         response = client.chat.completions.create(
             model="moonshotai/kimi-k2-instruct-0905",
             messages=[
-                {"role": "system", "content": "You are a database expert. Generate realistic table schemas and sample data. Return ONLY valid JSON, no markdown formatting."},
+                {"role": "system", "content": "You are a database expert. Generate realistic table schemas and sample data that will produce meaningful query results. Return ONLY valid JSON, no markdown formatting."},
                 {"role": "user", "content": schema_prompt}
             ],
             temperature=0.7
@@ -68,10 +83,87 @@ Make the data realistic and relevant to the query. Include enough variety to mak
         content = re.sub(r'```\s*$', '', content)
         
         schema_data = json.loads(content)
+        
+        # Post-process: Enhance data to ensure query results
+        schema_data = enhance_sample_data(schema_data, user_query)
+        
         return schema_data
     except Exception as e:
         raise Exception(f"Error generating sample data: {str(e)}")
 
+def enhance_sample_data(schema_data: dict, user_query: str) -> dict:
+    """Enhance sample data to ensure queries return results"""
+    
+    # Detect if query mentions time period
+    time_keywords = {
+        'last 30 days': 30,
+        'last 60 days': 60,
+        'last 7 days': 7,
+        'last week': 7,
+        'last month': 30,
+        'last quarter': 90,
+        'last year': 365
+    }
+    
+    days_back = None
+    for keyword, days in time_keywords.items():
+        if keyword in user_query.lower():
+            days_back = days
+            break
+    
+    # Detect amount/value thresholds
+    amount_pattern = r'\$?(\d+)'
+    amount_match = re.search(r'over \$?(\d+)', user_query.lower())
+    threshold_amount = int(amount_match.group(1)) if amount_match else None
+    
+    for table in schema_data['tables']:
+        enhanced_data = []
+        original_data = table['sample_data']
+        
+        # Find date and amount columns
+        date_cols = [col['name'] for col in table['columns'] if col['type'] == 'DATE']
+        amount_cols = [col['name'] for col in table['columns'] if 'amount' in col['name'].lower() or 'price' in col['name'].lower() or 'total' in col['name'].lower()]
+        
+        for row in original_data:
+            new_row = row.copy()
+            
+            # Enhance date fields to be within the time period
+            if days_back and date_cols:
+                for date_col in date_cols:
+                    if date_col in new_row:
+                        # Generate random date within the period
+                        random_days = random.randint(0, days_back)
+                        new_date = (datetime.now() - timedelta(days=random_days)).strftime('%Y-%m-%d')
+                        new_row[date_col] = new_date
+            
+            # Enhance amount fields to exceed threshold
+            if threshold_amount and amount_cols:
+                for amount_col in amount_cols:
+                    if amount_col in new_row:
+                        # 60% of records above threshold, 40% below
+                        if random.random() < 0.6:
+                            new_row[amount_col] = round(random.uniform(threshold_amount * 1.1, threshold_amount * 3), 2)
+                        else:
+                            new_row[amount_col] = round(random.uniform(threshold_amount * 0.3, threshold_amount * 0.9), 2)
+            
+            enhanced_data.append(new_row)
+        
+        # Add more rows if needed (ensure at least 15 rows)
+        while len(enhanced_data) < 15:
+            # Duplicate and modify existing rows
+            template_row = enhanced_data[len(enhanced_data) % len(original_data)].copy()
+            
+            # Modify IDs to be unique
+            for col in table['columns']:
+                if 'id' in col['name'].lower() and col['type'] == 'INTEGER':
+                    template_row[col['name']] = len(enhanced_data) + 1
+            
+            enhanced_data.append(template_row)
+        
+        table['sample_data'] = enhanced_data
+    
+    return schema_data
+
 def create_tables_in_db(schema_data: dict) -> sqlite3.Connection:
     """Create SQLite tables and populate with sample data"""
     conn = sqlite3.connect(':memory:')
@@ -114,14 +206,14 @@ def generate_sql_query(user_query: str, groq_api_key: str, schema_info: str) ->
 
 User Request: {user_query}
 
-Generate a SQL query that works with the above schema."""
+Generate a SQL query that works with the above schema. Use SQLite-compatible syntax."""
 
         response = client.chat.completions.create(
             model="moonshotai/kimi-k2-instruct-0905",
             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 SQLite.",
+                    "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 SQLite. For date operations, use SQLite functions like date('now') and datetime().",
                 },
                 {"role": "user", "content": enhanced_query},
             ],
@@ -190,9 +282,9 @@ def process_query(user_query: str, groq_api_key: str):
         # Display sample data
         sample_tables_html = []
         for table in schema_data['tables']:
-            df_sample = pd.DataFrame(table['sample_data'][:5])  # Show first 5 rows
-            sample_tables_html.append(f"<h4>Sample Data from '{table['table_name']}' (first 5 rows):</h4>")
-            sample_tables_html.append(df_sample.to_html(index=False, border=1))
+            df_sample = pd.DataFrame(table['sample_data'][:10])  # Show first 10 rows
+            sample_tables_html.append(f"<h4>Sample Data from '{table['table_name']}' (first 10 rows):</h4>")
+            sample_tables_html.append(df_sample.to_html(index=False, border=1, classes='table table-striped'))
         
         # Step 3: Generate SQL query
         output_log.append("### Step 3: Generating SQL Query")
@@ -221,11 +313,13 @@ def process_query(user_query: str, groq_api_key: str):
         result_df = execute_sql_query(conn, sql_generation.query)
         
         if len(result_df) == 0:
-            output_log.append("ℹ️ Query executed successfully but returned 0 rows")
-            result_html = "<p>No results found</p>"
+            output_log.append("⚠️ Query executed successfully but returned 0 rows")
+            output_log.append("This might happen if the sample data doesn't match the query criteria.")
+            result_html = "<p><i>No results found. The query executed successfully but no data matched the criteria.</i></p>"
         else:
             output_log.append(f"✅ Query executed successfully! Returned {len(result_df)} row(s)\n")
-            result_html = result_df.to_html(index=False, border=1)
+            result_html = f"<h4>Query Results ({len(result_df)} rows):</h4>"
+            result_html += result_df.to_html(index=False, border=1, classes='table table-striped')
         
         conn.close()
         
@@ -239,32 +333,56 @@ def process_query(user_query: str, groq_api_key: str):
         error_msg = f"❌ Error: {str(e)}"
         return error_msg, "", "", "", ""
 
+# Custom CSS for better table styling
+custom_css = """
+.table {
+    width: 100%;
+    border-collapse: collapse;
+    margin: 10px 0;
+}
+.table th {
+    background-color: #f0f0f0;
+    font-weight: bold;
+    padding: 8px;
+    text-align: left;
+    border: 1px solid #ddd;
+}
+.table td {
+    padding: 8px;
+    border: 1px solid #ddd;
+}
+.table-striped tbody tr:nth-child(odd) {
+    background-color: #f9f9f9;
+}
+"""
+
 # Gradio Interface
-with gr.Blocks(title="SQL Query Generator & Executor", theme=gr.themes.Soft()) as app:
+with gr.Blocks(title="SQLGenie - AI SQL Query Generator", theme=gr.themes.Ocean(), css=custom_css) as app:
     gr.Markdown("""
-    # 🔍 SQL Query Generator & Executor
+    # ⚡ SQLGenie - AI SQL Query Generator & Executor
     
-    This app uses Groq's Kimi-K2 model to:
-    1. Generate realistic sample database tables based on your query
-    2. Generate a structured SQL query from natural language
-    3. Execute the query and show results
+    Transform natural language into SQL queries and see instant results! This app:
+    1. 🎲 Generates realistic sample database tables based on your query
+    2. 🧙 Creates a structured SQL query from natural language using AI
+    3. ⚙️ Executes the query on sample data
+    4. 📊 Shows you the results instantly
     
     ### How to use:
-    1. Enter your Groq API key ([Get one here](https://console.groq.com/keys))
-    2. Enter your query in plain English
-    3. Click "Generate & Execute SQL"
+    1. Enter your Groq API key ([Get one free here](https://console.groq.com/keys))
+    2. Describe what data you want in plain English
+    3. Click "Generate & Execute SQL" and watch the magic happen! ✨
     """)
     
     with gr.Row():
         with gr.Column(scale=2):
             api_key_input = gr.Textbox(
-                label="Groq API Key",
+                label="🔑 Groq API Key",
                 placeholder="Enter your Groq API key here...",
                 type="password"
             )
             
             query_input = gr.Textbox(
-                label="Natural Language Query",
+                label="💬 Natural Language Query",
                 placeholder="Example: Find all customers who made orders over $500 in the last 30 days, show their name, email, and total order amount",
                 lines=3
             )
@@ -302,10 +420,12 @@ with gr.Blocks(title="SQL Query Generator & Executor", theme=gr.themes.Soft()) a
             ["List all products that are out of stock along with their supplier information"],
             ["Show the top 5 employees by total sales in the last quarter"],
             ["Find all students who scored above 85% in Mathematics and their contact details"],
-            ["Get all active users who haven't logged in for more than 60 days"]
+            ["Get all active users who haven't logged in for more than 60 days"],
+            ["Show all transactions above $1000 in the last week with customer details"],
+            ["Find employees in the Engineering department with salary over $80000"]
         ],
         inputs=query_input,
-        label="Example Queries"
+        label="💡 Example Queries - Click to try!"
     )
     
     submit_btn.click(
@@ -313,6 +433,15 @@ with gr.Blocks(title="SQL Query Generator & Executor", theme=gr.themes.Soft()) a
         inputs=[query_input, api_key_input],
         outputs=[process_output, sql_output, sample_data_output, result_output, gr.Textbox(visible=False)]
     )
+    
+    gr.Markdown("""
+    ---
+    ### 🎯 Tips for Best Results:
+    - Be specific about time periods (e.g., "last 30 days", "last quarter")
+    - Mention thresholds clearly (e.g., "over $500", "above 85%")
+    - Specify what fields you want to see (e.g., "show name, email, total")
+    - The app generates realistic sample data automatically to match your query!
+    """)
 
 if __name__ == "__main__":
     app.launch()