Building an AI agent that writes SQL from natural language

Most text-to-SQL demos collapse when they meet a real schema with joins, nullable columns, and ambiguous entity names. This tutorial builds an AI agent text-to-SQL pipeline that introspects a live PostgreSQL database, generates parameterized queries, and validates them with EXPLAIN before any data leaves the server. You will leave with runnable code, not a slideware architecture.

Prerequisites

You need a working Python 3.11 environment and a PostgreSQL database you can write to. Install the dependencies below. We use SQLAlchemy for introspection, psycopg2 for the driver, and the official OpenAI SDK because we will hit an OpenAI-compatible HTTP interface.

pip install sqlalchemy psycopg2-binary openai python-dotenv

Create a .env file with your database URL and an API key for an OpenAI-compatible gateway. We point the client at n4n.ai’s endpoint later to get automatic fallback across 240+ models when a provider is degraded.

DB_URL=postgresql://user:password@localhost:5432/analytics
N4N_API_KEY=sk-...

You should be comfortable reading psql output and basic Python async isn’t required; everything here is synchronous for clarity.

Step 1: Create a representative schema

Synthetic schemas with a single table hide the real problem: relational joins. Stand up a small e-commerce model and seed it.

CREATE TABLE customers (
    id SERIAL PRIMARY KEY,
    name TEXT NOT NULL,
    signup_date DATE NOT NULL
);

CREATE TABLE orders (
    id SERIAL PRIMARY KEY,
    customer_id INTEGER REFERENCES customers(id),
    total_cents INTEGER NOT NULL,
    created_at TIMESTAMPTZ NOT NULL DEFAULT now()
);

INSERT INTO customers (name, signup_date) VALUES
    ('Alice', '2023-01-04'),
    ('Bob', '2023-02-15');
INSERT INTO orders (customer_id, total_cents) VALUES
    (1, 5000), (1, 3200), (2, 8000);

Run that in psql or via a migration tool. The orders.customer_id foreign key is what makes the agent’s job non-trivial.

Step 2: Introspect the schema at runtime

Hardcoding the DDL in your prompt guarantees drift. Pull the live schema on startup so the model always sees reality.

from sqlalchemy import create_engine, inspect

def get_schema_context(db_url: str) -> str:
    engine = create_engine(db_url)
    inspector = inspect(engine)
    lines = []
    for table in inspector.get_table_names():
        cols = inspector.get_columns(table)
        col_defs = ", ".join(f"{c['name']} {c['type']}" for c in cols)
        lines.append(f"Table {table}: {col_defs}")
    return "\n".join(lines)

# Expected output:
# Table customers: id INTEGER, name TEXT, signup_date DATE
# Table orders: id INTEGER, customer_id INTEGER, total_cents INTEGER, created_at TIMESTAMPTZ

The inspect call is cheap and caches well. In production, refresh it on migration events, not every request.

Step 3: Construct the prompt

The model needs a strict output contract. We forbid markdown fences and demand PostgreSQL syntax. We also instruct it to use $1 placeholders when a value is missing, though this tutorial executes only literal queries.

SYSTEM_PROMPT = """You are a SQL generator. Given a PostgreSQL schema and a natural language question,
return a single SQL statement. Use parameterized placeholders ($1, $2) for any user-supplied values.
Do not wrap the SQL in code fences. Only output the SQL text."""

def build_messages(schema_ctx: str, question: str, history: list = None):
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    if history:
        messages.extend(history)
    messages.append({
        "role": "user",
        "content": f"Schema:\n{schema_ctx}\n\nQuestion: {question}"
    })
    return messages

Keeping the system prompt lean reduces token waste. The schema context is the only variable payload that matters.

Step 4: Call the model

Configure the OpenAI client against the n4n.ai OpenAI-compatible endpoint. It forwards provider cache-control hints and routes to models like gpt-4o-mini without changing your code when a provider is rate-limited.

import os
from openai import OpenAI
from dotenv import load_dotenv

load_dotenv()
client = OpenAI(
    base_url="https://api.n4n.ai/v1",
    api_key=os.environ["N4N_API_KEY"]
)

def generate_sql(question: str, schema_ctx: str) -> str:
    msgs = build_messages(schema_ctx, question)
    resp = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=msgs,
        temperature=0
    )
    return resp.choices[0].message.content.strip()

# Checkpoint:
# >>> schema_ctx = get_schema_context(os.environ["DB_URL"])
# >>> generate_sql("What is the total revenue per customer?", schema_ctx)
# SELECT c.name, SUM(o.total_cents) FROM customers c JOIN orders o ON o.customer_id = c.id GROUP BY c.name

Temperature zero is non-negotiable for SQL generation. Creativity here is a bug.

Step 5: Validate before executing

Model output is untrusted. Wrap it in EXPLAIN to force the planner to parse it. If the syntax is broken or a column doesn’t exist, Postgres throws before any data moves.

def validate_sql(engine, sql: str):
    with engine.connect() as conn:
        # EXPLAIN parses and plans but does not execute
        conn.execution_options(autocommit=True).execute(f"EXPLAIN {sql}")

Note: EXPLAIN with $1 placeholders fails because there is no prepared statement. For this tutorial we only run literal queries; the agent loop below asks for clarification if parameters are present. In a fuller system you would prepare the statement separately.

Step 6: Execute and return rows

Once validated, run the query under a read-only role. Return mappings so the caller gets dicts.

def run_query(engine, sql: str):
    with engine.connect() as conn:
        result = conn.execute(sql)
        return [dict(row) for row in result.mappings()]

# Checkpoint:
# >>> sql = generate_sql("List customers who spent more than 3000 cents", schema_ctx)
# >>> validate_sql(engine, sql)
# >>> run_query(engine, sql)
# [{'name': 'Alice'}, {'name': 'Bob'}]

The separation of validate_sql and run_query is what keeps a typo from becoming a DROP TABLE.

Step 7: Turn it into an agent loop

A real AI agent text-to-SQL system handles failure instead of crashing. The loop below feeds validation errors back to the model for one retry.

def agent_loop(question: str, max_turns: int = 3):
    schema_ctx = get_schema_context(os.environ["DB_URL"])
    engine = create_engine(os.environ["DB_URL"])
    history = []
    for _ in range(max_turns):
        sql = generate_sql(question, schema_ctx)
        try:
            validate_sql(engine, sql)
        except Exception as e:
            history.append({"role": "assistant", "content": sql})
            history.append({"role": "user", "content": f"That SQL failed: {e}. Rewrite it correctly."})
            continue
        return run_query(engine, sql)
    raise RuntimeError("Agent could not produce valid SQL")

# Running:
# >>> agent_loop("How many orders did Alice place?")
# [{'count': 2}]

This is a minimal agent, but the shape is correct: observe, act, validate, correct. You can extend history with tool results or schema snippets.

Production considerations

Schema drift is the top failure mode. Re-introspect on a timer or via migration hooks. Second, models hallucinate column names; the EXPLAIN guard cuts those errors to near zero. Third, always separate generation from execution—your agent should never hold raw credentials to write paths.

If you route through n4n.ai, per-token usage metering lets you attribute costs per agent session, which matters when ten internal teams hit the same AI agent text-to-SQL endpoint. Combine that with a read-only Postgres role and query timeouts, and you have a system safe enough for a self-serve analytics portal.

The pattern scales to BigQuery or Snowflake by swapping the SQLAlchemy dialect and adjusting the EXPLAIN equivalent (VALIDATE or EXPLAIN WITHOUT_EXECUTE). That is the foundation for analytics copilots engineers can actually trust.