--- title: "SQL Agent" description: "AST-based safety validation (not regex) on generated SQL is the detail that makes this defensible in a real interview β€” shows you think about AI output as..." source: "https://confidentprep.com/paths/ai-augmented-engineering/sql-agent/" path: "AI-Augmented Engineering" repository: "https://github.com/confidentprep14-web/ai-engineering-projects/tree/main/sql-agent" token_estimate: 5064 --- # SQL Agent ## Overview A natural-language-to-SQL agent with `sqlglot` AST-based safety validation (SELECT-only, provably β€” not a fragile regex check), schema introspection, and a self-correcting retry loop that injects SQL errors back into the prompt. Same retry-on-failure pattern as the test generator, applied somewhere a wrong answer can actually hurt β€” a database. ## What to do | Field | Value | |---|---| | Path | 2 β€” AI-Augmented Engineering | | Position | 10 of 12 | | Difficulty | 🟑 Optional Docker | | Estimated time | 4 hours | | AWS cost | None | --- ### Agent Pickup Instructions ```bash # 1. Enter project cd projects/p2-10-sql-agent # 2. Create virtual environment python -m venv .venv && source .venv/bin/activate # 3. Install dependencies pip install -r requirements.txt # 4. Copy env file and add your API key cp .env.example .env # 5. Initialize the sample e-commerce database python scripts/init_db.py # 6. Run the agent interactively python src/main.py --db ecommerce.db # 7. Run the 10-pair test suite python src/main.py --db ecommerce.db --test # 8. Run tests pytest tests/ -v ``` **Done when:** - [ ] Interactive mode translates natural language to SQL and explains the result in plain English - [ ] INSERT/UPDATE/DELETE/DROP are rejected with an error message (not executed) - [ ] Error recovery loop retries up to 2 times, injecting the SQL error into the next prompt - [ ] `--test` runs all 10 NLβ†’SQL pairs and prints pass/fail for each - [ ] All 4 tests pass --- ### What this project is A natural language to SQL agent that introspects a SQLite database schema at startup, translates user questions into SELECT statements using an LLM, validates SQL using the `sqlglot` AST parser (not regex) before execution, executes safe queries, and interprets results in plain English. A retry loop (max 2 attempts) injects SQL error messages back into the LLM prompt for self-correction. A built-in test suite exercises 10 natural language to SQL pairs against a provided e-commerce schema to measure accuracy. --- ### What the learner achieves "I built a natural language to SQL agent with parser-based safety validation (not regex), schema introspection at startup, and a self-correcting retry loop β€” and I can demonstrate it on a real e-commerce schema with 10 test cases." --- ### Folder structure ``` p2-10-sql-agent/ β”œβ”€β”€ src/ β”‚ β”œβ”€β”€ main.py # CLI: --db , --schema , --test β”‚ β”œβ”€β”€ llm.py # Provider-agnostic LLM wrapper β”‚ β”œβ”€β”€ schema_inspector.py # introspect_db, format_schema_for_prompt β”‚ β”œβ”€β”€ sql_validator.py # parse_and_validate (sqlglot AST) β”‚ └── agent.py # nl_to_sql, execute_safe, retry_on_error, interpret_results β”œβ”€β”€ tests/ β”‚ β”œβ”€β”€ test_sql_validator.py β”‚ β”œβ”€β”€ test_agent.py β”‚ └── test_schema_inspector.py β”œβ”€β”€ scripts/ β”‚ └── init_db.py # Creates ecommerce.db with schema + sample data β”œβ”€β”€ .env.example β”œβ”€β”€ requirements.txt └── README.md ``` --- ### E-commerce Schema The database has three tables. Use this schema in all tests and in `init_db.py`. ```sql CREATE TABLE customers ( id INTEGER PRIMARY KEY, name TEXT NOT NULL, email TEXT UNIQUE NOT NULL, country TEXT NOT NULL, created_at TEXT NOT NULL -- ISO datetime string ); CREATE TABLE products ( id INTEGER PRIMARY KEY, name TEXT NOT NULL, category TEXT NOT NULL, price REAL NOT NULL, stock_quantity INTEGER NOT NULL DEFAULT 0 ); CREATE TABLE orders ( id INTEGER PRIMARY KEY, customer_id INTEGER NOT NULL REFERENCES customers(id), product_id INTEGER NOT NULL REFERENCES products(id), quantity INTEGER NOT NULL, total_amount REAL NOT NULL, status TEXT NOT NULL, -- 'pending', 'shipped', 'delivered', 'cancelled' ordered_at TEXT NOT NULL -- ISO datetime string ); ``` Sample data to insert (via `init_db.py`): - 5 customers (mix of US, UK, CA) - 8 products (electronics, books, clothing categories; prices $9.99 to $299.99) - 15 orders (mix of statuses, customers, products) --- ### 10 NLβ†’SQL Test Pairs | # | Natural Language | Expected SQL (exact match not required β€” semantic match) | |---|---|---| | 1 | "Show me all customers from the US" | `SELECT * FROM customers WHERE country = 'US'` | | 2 | "What are the 5 most expensive products?" | `SELECT * FROM products ORDER BY price DESC LIMIT 5` | | 3 | "How many orders are in 'shipped' status?" | `SELECT COUNT(*) FROM orders WHERE status = 'shipped'` | | 4 | "List all products with less than 10 items in stock" | `SELECT * FROM products WHERE stock_quantity < 10` | | 5 | "What is the total revenue from delivered orders?" | `SELECT SUM(total_amount) FROM orders WHERE status = 'delivered'` | | 6 | "Show me orders placed by customer with ID 1" | `SELECT * FROM orders WHERE customer_id = 1` | | 7 | "Which customers have placed more than 2 orders?" | `SELECT customer_id, COUNT(*) as order_count FROM orders GROUP BY customer_id HAVING COUNT(*) > 2` | | 8 | "What is the average order value?" | `SELECT AVG(total_amount) FROM orders` | | 9 | "List customers along with their total number of orders" | `SELECT c.name, COUNT(o.id) FROM customers c LEFT JOIN orders o ON c.id = o.customer_id GROUP BY c.id` | | 10 | "What products have never been ordered?" | `SELECT p.* FROM products p LEFT JOIN orders o ON p.id = o.product_id WHERE o.id IS NULL` | --- ### .env.example ```bash # LLM provider: anthropic | openai | ollama LLM_PROVIDER=anthropic # Model to use LLM_MODEL= # API key for Anthropic ANTHROPIC_API_KEY= # API key for OpenAI OPENAI_API_KEY= # Ollama base URL OLLAMA_BASE_URL=http://localhost:11434 # Default SQLite database path DEFAULT_DB_PATH=ecommerce.db # Maximum SQL generation retries on error MAX_SQL_RETRIES=2 # Maximum rows to display in results table MAX_DISPLAY_ROWS=20 ``` --- ### requirements.txt ``` anthropic==0.30.0 openai==1.35.0 sqlglot==23.3.0 python-dotenv==1.0.1 rich==13.7.1 pytest==8.2.2 pytest-mock==3.14.0 ``` --- ### src/ β€” what to implement #### `src/llm.py` **`get_completion(prompt: str, system: str = "") -> str`** - Standard interface: reads `LLM_PROVIDER` and `LLM_MODEL` from env - Edge cases: raise `RuntimeError` if API key is missing --- #### `src/schema_inspector.py` **`introspect_db(db_path: str) -> dict`** - Inputs: path to SQLite database file - Output: `{"tables": [{"name": str, "columns": [{"name": str, "type": str, "nullable": bool, "pk": bool}], "foreign_keys": [...]}]}` - Behavior: use `sqlite3.connect(db_path)`; query `sqlite_master` for tables; for each table, query `PRAGMA table_info` and `PRAGMA foreign_key_list` - Edge cases: if db file does not exist, raise `FileNotFoundError` **`format_schema_for_prompt(schema: dict) -> str`** - Inputs: schema dict from `introspect_db` - Output: formatted string for inclusion in LLM prompt - Format: ``` Tables: customers (id INTEGER PK, name TEXT, email TEXT UNIQUE, country TEXT, created_at TEXT) products (id INTEGER PK, name TEXT, category TEXT, price REAL, stock_quantity INTEGER) orders (id INTEGER PK, customer_id INTEGER FKβ†’customers.id, product_id INTEGER FKβ†’products.id, quantity INTEGER, total_amount REAL, status TEXT, ordered_at TEXT) ``` --- #### `src/sql_validator.py` **`ValidationResult` (dataclass)** ```python @dataclass class ValidationResult: is_valid: bool error: str | None # Human-readable error if not valid statement_type: str | None # "SELECT", "INSERT", etc. ``` **`parse_and_validate(sql: str) -> ValidationResult`** - Inputs: SQL string - Output: `ValidationResult` - Behavior: - Use `sqlglot.parse_one(sql)` to parse - Check `type(parsed).__name__`: if it is not `Select`, return `ValidationResult(is_valid=False, error=f"Only SELECT queries are allowed. Got: {type(parsed).__name__}.", statement_type=type(parsed).__name__)` - If parsing raises `sqlglot.errors.ParseError`, return `ValidationResult(is_valid=False, error=f"SQL parse error: {e}", statement_type=None)` - If statement type is SELECT, return `ValidationResult(is_valid=True, error=None, statement_type="SELECT")` - Do NOT use string matching or regex β€” use the AST node type **`extract_sql_from_response(response: str) -> str`** - Inputs: raw LLM response string - Output: clean SQL string - Behavior: strip markdown code fences (` ```sql ` and ` ``` `); strip leading/trailing whitespace; if multiple statements, take only the first (split on `;` and take index 0) --- #### `src/agent.py` **`nl_to_sql(question: str, schema_str: str, prior_error: str = "") -> str`** - Inputs: natural language question, formatted schema string, optional prior SQL error - Output: raw LLM response (may include markdown β€” extract separately) - Behavior: - System prompt: "You are a SQLite expert. Given a database schema and a user question, generate a single SELECT statement that answers the question. Return ONLY the SQL query, no explanation." - User prompt: `Schema:\n{schema_str}\n\nQuestion: {question}` - If `prior_error` is non-empty, prepend: `"The previous SQL produced this error: {prior_error}\nGenerate a corrected SQL query."` - Call `get_completion` **`execute_safe(sql: str, db_path: str) -> tuple[list[dict], list[str]]`** - Inputs: validated SQL string, database path - Output: `(rows, column_names)` where `rows` is a list of dicts - Behavior: `sqlite3.connect(db_path)`; `cursor.execute(sql)`; `cursor.fetchall()`; convert to list of dicts using column names from `cursor.description` - Edge cases: runtime SQL errors β†’ raise `RuntimeError(str(e))` β€” caller handles retry **`retry_on_error(question: str, schema_str: str, db_path: str, max_retries: int = 2) -> tuple`** - Inputs: question, schema, db path, max retries - Output: `(rows, column_names, sql_used, attempts)` or raises `RuntimeError` if all retries fail - Behavior: 1. Call `nl_to_sql` with no prior error 2. Call `extract_sql_from_response` 3. Call `parse_and_validate` β€” if invalid, use the validation error as `prior_error`; retry 4. If valid, call `execute_safe` β€” if it raises, use the exception as `prior_error`; retry 5. On success, return results 6. Log `Retry {n}/{max_retries}: {error}` on each retry **`interpret_results(question: str, rows: list[dict], column_names: list[str]) -> str`** - Inputs: original question, query results, column names - Output: plain-English explanation - Behavior: if `rows` is empty, return "No results found." without LLM call; otherwise, format top 5 rows as a simple table in the prompt and ask the LLM to explain what the results mean in context of the question - System prompt: "You are a data analyst. Explain these SQL query results in plain English, in 2-3 sentences. Focus on the answer to the question, not the data format." --- #### `src/main.py` CLI entry point using `argparse`. Arguments: - `--db ` β€” SQLite database path (default from env `DEFAULT_DB_PATH`) - `--schema ` β€” optional pre-formatted schema file (skip introspection) - `--test` β€” run the 10 NLβ†’SQL test pairs and print pass/fail Behavior for interactive mode: 1. Load `.env`; introspect db schema; print `Schema loaded: N tables` 2. Loop: `Question> ` prompt; call `retry_on_error`; print SQL used; print results table; print plain-English interpretation 3. Exit on `quit` or `exit` Behavior for `--test`: 1. Load schema 2. For each of the 10 test pairs: call `retry_on_error`; execute; check that results are non-empty for queries that should return data; print `[PASS]` or `[FAIL]` with the SQL generated 3. Print summary: `Passed N/10` --- ### tests/ β€” what to test #### Test 1 β€” sql_validator blocks INSERT (`test_sql_validator.py`) - Call `parse_and_validate("INSERT INTO customers VALUES (1, 'Alice', 'a@b.com', 'US', '2024-01-01')")` - Assert `is_valid == False` - Assert `"INSERT"` appears in `error` or `statement_type` #### Test 2 β€” sql_validator blocks UPDATE (`test_sql_validator.py`) - Call `parse_and_validate("UPDATE customers SET name='Bob' WHERE id=1")` - Assert `is_valid == False` #### Test 3 β€” sql_validator allows SELECT with JOIN (`test_sql_validator.py`) - Call `parse_and_validate("SELECT c.name, COUNT(o.id) FROM customers c LEFT JOIN orders o ON c.id = o.customer_id GROUP BY c.id")` - Assert `is_valid == True` - Assert `statement_type == "SELECT"` #### Test 4 β€” Retry injects SQL error into context (`test_agent.py`) - Mock `execute_safe` to raise `RuntimeError("no such column: foo")` on first call, succeed on second - Mock `nl_to_sql` to track call arguments - Call `retry_on_error` - Assert `nl_to_sql` was called twice - Assert the second call's `prior_error` argument contains "no such column" #### Test 5 β€” Schema introspection returns table names (`test_schema_inspector.py`) - Create a temp in-memory SQLite db with the e-commerce schema - Call `introspect_db` - Assert `len(schema["tables"]) == 3` - Assert table names include `"customers"`, `"products"`, `"orders"` #### Test 6 β€” Empty results handled without LLM call (`test_agent.py`) - Call `interpret_results("How many?", rows=[], column_names=["count"])` - Assert result == "No results found." - Assert no LLM call was made (use mock to verify `get_completion` was not called) --- ### README.md content ```markdown # SQL Agent Natural language β†’ SQL β†’ plain English. Uses parser-based SQL validation (sqlglot AST) and a self-correcting retry loop. ## Quick start ```bash python -m venv .venv && source .venv/bin/activate pip install -r requirements.txt cp .env.example .env # add your API key ## Initialize the e-commerce database python scripts/init_db.py ## Run interactively python src/main.py --db ecommerce.db ## Run the 10-pair test suite python src/main.py --db ecommerce.db --test ``` ## Safety model Only SELECT statements are allowed. Validation uses the `sqlglot` AST β€” not regex: ```python parsed = sqlglot.parse_one(sql) if not isinstance(parsed, sqlglot.expressions.Select): raise ValueError("Only SELECT allowed") ``` ## E-commerce schema Three tables: `customers`, `products`, `orders` See `scripts/init_db.py` for full schema and sample data. ## Retry loop ``` NL β†’ SQL β†’ validate β†’ execute ↓ invalid/error inject error into prompt β†’ retry (max 2) ``` ## Running tests ```bash pytest tests/ -v ``` ``` --- ### GUIDE.md content ```markdown # Build Guide β€” SQL Agent ## Step 1 β€” Schema introspection At startup, read the schema so the LLM has accurate context: ```python import sqlite3 def introspect_db(db_path): conn = sqlite3.connect(db_path) cursor = conn.cursor() tables = [] cursor.execute("SELECT name FROM sqlite_master WHERE type='table'") for (table_name,) in cursor.fetchall(): cursor.execute(f"PRAGMA table_info({table_name})") columns = [ {"name": row[1], "type": row[2], "nullable": not row[3], "pk": bool(row[5])} for row in cursor.fetchall() ] tables.append({"name": table_name, "columns": columns}) return {"tables": tables} ``` ## Step 2 β€” Parser-based validation (not regex) The key insight: regex cannot parse SQL reliably. A query like `SELECT * FROM (SELECT name FROM t) WHERE ...` contains a subquery that regex might mistake for a simple SELECT. Use the AST: ```python import sqlglot def parse_and_validate(sql): try: parsed = sqlglot.parse_one(sql) except sqlglot.errors.ParseError as e: return ValidationResult(is_valid=False, error=str(e), statement_type=None) stmt_type = type(parsed).__name__ if stmt_type != "Select": return ValidationResult( is_valid=False, error=f"Only SELECT allowed. Got: {stmt_type}", statement_type=stmt_type, ) return ValidationResult(is_valid=True, error=None, statement_type="SELECT") ``` ## Step 3 β€” The generation prompt ``` System: You are a SQLite expert. Given a schema and a question, write a single SELECT query. Return ONLY the SQL β€” no explanation, no markdown. User: Schema: {schema_str} Question: {question} ``` If retrying: prepend `"Previous SQL caused this error: {error}\nWrite a corrected query."` ## Step 4 β€” Interpretation prompt ``` System: You are a data analyst. Explain these SQL results in 2-3 plain-English sentences focused on the answer to the question. User: Question: {question} Results (first 5 rows): {formatted_rows} ``` ## Step 5 β€” Wire the test suite ```python NL_SQL_PAIRS = [ ("Show me all customers from the US", "customers"), ("5 most expensive products", "products"), ... ] for question, _ in NL_SQL_PAIRS: rows, cols, sql, attempts = retry_on_error(question, schema_str, db_path) passed = len(rows) > 0 or "COUNT" in sql.upper() print(f"{'[PASS]' if passed else '[FAIL]'} {question}") ``` ## Debugging tips - If the model generates markdown fences, `extract_sql_from_response` strips them β€” verify this runs before validation - If sqlglot version mismatch causes import errors, pin to `sqlglot==23.3.0` exactly - If schema introspection returns no tables, check that `init_db.py` ran and created the `.db` file ## How to talk about this in an interview **"Why sqlglot instead of regex for SQL validation?"** > SQL can't be reliably parsed with regex. A query like `SELECT name FROM (SELECT * FROM users WHERE type='DROP TABLE')` contains 'DROP TABLE' as a string literal β€” regex would flag it, sqlglot's AST won't. Parser-based validation is accurate; regex-based validation is security theater. **"What does the retry loop do?"** > When the database returns an error β€” like 'no such column: usr_id' β€” I inject that exact error into the next prompt. The model sees: 'your previous SQL caused this error, write a corrected version.' This handles about 80% of first-attempt failures without human intervention. **"How do you prevent SQL injection?"** > Two layers: AST validation (only SELECT allowed) and SQLite's parameterized queries for any user-provided values. But since we're generating SQL from an LLM's output rather than concatenating user strings, the injection surface is different from traditional web apps. ``` --- ### Failure modes to handle | Failure | Where | How to handle | |---|---|---| | LLM returns markdown fences around SQL | `sql_validator.py` | `extract_sql_from_response` strips ` ```sql ` and ` ``` ` before validation | | Generated SQL is invalid (parse error) | `agent.py` | Treat parse error as `prior_error`; retry with error injected | | Generated SQL is valid but fails at runtime | `agent.py` | Catch `RuntimeError` from `execute_safe`; inject error; retry | | All retries fail | `agent.py` | Raise `RuntimeError(f"Could not generate valid SQL after {max_retries} attempts")` | | Database file missing | `schema_inspector.py` | Raise `FileNotFoundError` with helpful message: "Run: python scripts/init_db.py" | | Query returns 0 rows | `agent.py` | Return empty list; `interpret_results` returns "No results found." without LLM call | --- ### The metric this project measures **What is measured:** Test suite accuracy (N/10 pairs passing) and retry rate (attempts needed per question). **Format (stdout):** ``` [PASS] Show me all customers from the US | SQL: SELECT * FROM customers WHERE country = 'US' [PASS] What are the 5 most expensive products? | SQL: SELECT * FROM products ORDER BY price DESC LIMIT 5 ... Passed 9/10 | avg attempts: 1.2 ``` **Target:** At least 8/10 NLβ†’SQL pairs must pass (return non-empty results for data queries, or a valid result for aggregate queries). Average attempts must be logged. This measures both LLM SQL generation quality and retry loop effectiveness together. ### Model source currency - Anthropic model list: https://docs.anthropic.com/en/docs/about-claude/models/overview - OpenAI model list: https://platform.openai.com/docs/models - verifiedOn: 2026-08-07 - `LLM_MODEL` is required and has no implicit default; choose a supported model for the configured provider at setup time. ## Source code A reference implementation of SQL Agent β€” README, guide, source and tests β€” is at https://github.com/confidentprep14-web/ai-engineering-projects/tree/main/sql-agent.