Capstone 2: Autonomous Text-to-SQL Analyst Agent

+---------------------------------------------------------------------------------+
|                           PROJECT 2 ARCHITECTURE                                |
+---------------------------------------------------------------------------------+
|  Question: "Top 5 sales reps by deal volume in Q3"                              |
|                                                                                 |
|  1. Metadata Selector -> Retrieves relevant table DDLs                          |
|  2. ReAct Planner     -> Generates candidate SQL query                          |
|  3. AST Safety Guard  -> Validates non-mutating SELECT with sqlglot             |
|  4. Executor          -> Runs on SQLite / PostgreSQL                            |
|  5. Self-Correction   -> Auto-fixes syntax errors if engine rejects query       |
|  6. Synthesizer       -> Explains results with bullet points                    |
+---------------------------------------------------------------------------------+

Objective

Build a robust, self-healing Text-to-SQL agent that translates complex natural language queries into safe, highly optimized SQL, automatically handling schema misunderstandings.

Complete Implementation

import sqlite3
import sqlglot
from sqlglot.expressions import Select, Drop, Delete, Update
from openai import OpenAI

client = OpenAI()

def validate_safe_sql(query: str) -> str:
    parsed = sqlglot.parse_one(query)
    for expr in parsed.find_all((Drop, Delete, Update)):
        raise PermissionError(f"Mutating statement detected: {type(expr)}")
    if not isinstance(parsed, Select):
        raise ValueError("Only SELECT queries allowed.")
    if not parsed.args.get("limit"):
        parsed = parsed.limit(100)
    return str(parsed)

def execute_agent_text_to_sql(db_conn: sqlite3.Connection, schema_ddl: str, user_question: str) -> str:
    system_prompt = f"""
You are an expert SQL engineer. Given this database schema, write a valid SQLite SQL query to answer the user question.
SCHEMA:
{schema_ddl}

Output ONLY the raw SQL statement. Do not format with markdown blocks.
"""
    # Step 1: Generate initial query
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_question}
        ],
        temperature=0.0
    )
    raw_sql = response.choices[0].message.content.strip().replace("```sql", "").replace("```", "")

    # Step 2: Validate safety
    safe_sql = validate_safe_sql(raw_sql)

    # Step 3: Execute with retry loop
    cursor = db_conn.cursor()
    try:
        cursor.execute(safe_sql)
        rows = cursor.fetchall()
        columns = [desc[0] for desc in cursor.description]
        return f"SUCCESS:
Query: {safe_sql}
Columns: {columns}
Data: {rows}"
    except Exception as e:
        # Self-correction step
        fix_prompt = f"""
The following query generated for SQLite failed with error: {str(e)}
FAILED QUERY: {safe_sql}
SCHEMA: {schema_ddl}
Please output the corrected SQL query only.
"""
        fix_res = client.chat.completions.create(
            model="gpt-4o-mini",
            messages=[{"role": "user", "content": fix_prompt}]
        )
        corrected_sql = validate_safe_sql(fix_res.choices[0].message.content.strip())
        cursor.execute(corrected_sql)
        rows = cursor.fetchall()
        return f"RECOVERED:
Query: {corrected_sql}
Data: {rows}"