app.py

from groq import Groq
from pydantic import BaseModel
import json
import gradio as gr
import pandas as pd

class ValidationStatus(BaseModel):
    is_valid: bool
    syntax_errors: list[str]

class SQLQueryGeneration(BaseModel):
    query: str
    query_type: str
    tables_used: list[str]
    estimated_complexity: str
    execution_notes: list[str]
    validation_status: ValidationStatus
    table_schema: str
    sample_data: str
    execution_results: str
    optimization_notes: list[str]

def parse_execution_results_to_dataframe(execution_results):
    """Convert text-based table results to pandas DataFrame"""
    try:
        lines = execution_results.strip().split('\n')
        if len(lines) < 3:
            return None
        
        # Extract header
        header_line = lines[0]
        headers = [col.strip() for col in header_line.split('|')]
        
        # Extract data rows (skip separator line)
        data_rows = []
        for line in lines[2:]:
            if line.strip() and not line.strip().startswith('-'):
                row = [cell.strip() for cell in line.split('|')]
                if len(row) == len(headers):
                    data_rows.append(row)
        
        if data_rows:
            df = pd.DataFrame(data_rows, columns=headers)
            return df
        return None
    except Exception as e:
        print(f"Error parsing results: {e}")
        return None

def generate_sql_query(api_key, user_query):
    """Generate SQL query from natural language using GROQ API"""
    try:
        if not api_key:
            return "Error: Please enter your GROQ API key", "", "", "", None, ""
        
        if not user_query:
            return "Error: Please enter a query description", "", "", "", None, ""
        
        client = Groq(api_key=api_key)
        
        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.
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 in a pipe-delimited table format

IMPORTANT: The execution_results field must contain a properly formatted table with:
- Header row with column names separated by pipes (|)
- A separator row with dashes
- Data rows with values separated by pipes (|)

Example format:
column1 | column2 | column3
--------|---------|--------
value1  | value2  | value3
value4  | value5  | value6

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 (in pipe-delimited table format)
- Any relevant notes about assumptions made or query optimization suggestions""",
                },
                {
                    "role": "user", 
                    "content": user_query
                },
            ],
            response_format={
                "type": "json_schema",
                "json_schema": {
                    "name": "sql_query_generation",
                    "schema": SQLQueryGeneration.model_json_schema()
                }
            }
        )
        
        sql_query_generation = SQLQueryGeneration.model_validate(
            json.loads(response.choices[0].message.content)
        )
        
        # 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:
            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)}"""
        
        # Convert execution results to DataFrame
        results_df = parse_execution_results_to_dataframe(sql_query_generation.execution_results)
        
        return (
            sql_query_generation.query,
            metadata,
            sql_query_generation.table_schema,
            sql_query_generation.sample_data,
            results_df,
            validation_text
        )
        
    except Exception as e:
        error_msg = f"Error: {str(e)}"
        return error_msg, "", "", "", None, ""

# Create Gradio interface
with gr.Blocks(title="SQL Query Generator", theme=gr.themes.Ocean()) 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",
                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="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
            )
    
    with gr.Row():
        with gr.Column():
            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.Dataframe(
            label="📊 Execution Results",
            headers=None,
            datatype="str",
            row_count=10,
            col_count=None,
            wrap=True,
            interactive=False
        )
    
    generate_btn.click(
        fn=generate_sql_query,
        inputs=[api_key_input, query_input],
        outputs=[
            sql_output,
            metadata_output,
            schema_output,
            sample_data_output,
            execution_output,
            validation_output
        ]
    )
    
    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 in Excel-style table format
        - Optimization suggestions
        """
    )

if __name__ == "__main__":
    demo.launch(share=True)