Building a SQL agent with LangChain and tool calling

A langchain sql agent tool calling setup lets you turn natural language into executable queries without hand-rolling prompt templates or regex parsing. This tutorial walks through building one that handles schema discovery, query validation, and error recovery — patterns that survive contact with real databases. You’ll end up with an agent that can answer questions like “show me the top 5 customers by revenue last quarter” against a live PostgreSQL or SQLite instance.

Step 1: Install the right dependencies

LangChain’s SQL tooling lives in langchain-community and langchain-experimental. Pin versions — the agent interfaces shifted between 0.1 and 0.2.

pip install "langchain==0.2.*" "langchain-community==0.2.*" "langchain-experimental==0.2.*" \
  "langchain-openai==0.1.*" sqlalchemy psycopg2-binary python-dotenv

If you’re targeting SQLite for local dev, swap psycopg2-binary for sqlite3 (stdlib). For production PostgreSQL, keep both.

Step 2: Create a database connection with SQLAlchemy

The agent needs a SQLDatabase wrapper, not a raw connection. This gives it dialect-aware introspection.

# db.py
from sqlalchemy import create_engine, text
from sqlalchemy.pool import NullPool
from langchain_community.utilities import SQLDatabase
import os

DATABASE_URL = os.getenv("DATABASE_URL", "sqlite:///./demo.db")

# NullPool avoids connection leaks in short-lived agent runs
engine = create_engine(DATABASE_URL, poolclass=NullPool, future=True)

# Optional: verify connectivity and run a quick migration
with engine.connect() as conn:
    conn.execute(text("SELECT 1"))
    # Create sample tables if they don't exist
    conn.execute(text("""
        CREATE TABLE IF NOT EXISTS customers (
            id INTEGER PRIMARY KEY,
            name TEXT NOT NULL,
            email TEXT UNIQUE,
            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        );
    """))
    conn.execute(text("""
        CREATE TABLE IF NOT EXISTS orders (
            id INTEGER PRIMARY KEY,
            customer_id INTEGER REFERENCES customers(id),
            amount_cents INTEGER NOT NULL,
            currency TEXT DEFAULT 'USD',
            placed_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
        );
    """))
    conn.commit()

db = SQLDatabase(engine, sample_rows_in_table_info=3)

Verify: Run python -c "from db import db; print(db.get_usable_table_names())" — you should see ['customers', 'orders'].

Step 3: Define the tool-calling LLM

Use a model that supports OpenAI-style tool calling. GPT-4o, GPT-4-turbo, and Claude 3.5 Sonnet all work. The agent expects bind_tools to be available.

# llm.py
from langchain_openai import ChatOpenAI
import os

llm = ChatOpenAI(
    model="gpt-4o",
    temperature=0,
    api_key=os.getenv("OPENAI_API_KEY"),
    # If you're routing through a gateway that honors provider cache hints:
    # base_url="https://api.n4n.ai/v1",
)

Verify: python -c "from llm import llm; print(llm.invoke('ping').content)" — returns a response.

Step 4: Build the SQL toolkit with custom tools

The built-in SQLDatabaseToolkit gives you query_sql_db and info_sql_db. For production, wrap them to add guardrails: row limits, read-only enforcement, and explain-plan logging.

# tools.py
from langchain_community.agent_toolkits import SQLDatabaseToolkit
from langchain_community.utilities import SQLDatabase
from langchain_core.tools import Tool
from sqlalchemy import text
from db import db, engine

toolkit = SQLDatabaseToolkit(db=db, llm=None)  # llm not needed for base tools
base_tools = toolkit.get_tools()

# Find the raw query tool
query_tool = next(t for t in base_tools if t.name == "query_sql_db")
info_tool = next(t for t in base_tools if t.name == "info_sql_db")

MAX_ROWS = 100

def safe_query(sql: str) -> str:
    """Execute a SELECT-only query with a hard row limit."""
    stripped = sql.strip().rstrip(";")
    if not stripped.lower().startswith("select"):
        return "Error: Only SELECT statements are allowed."
    # Enforce LIMIT if not present
    if "limit" not in stripped.lower():
        stripped += f" LIMIT {MAX_ROWS}"
    try:
        with engine.connect() as conn:
            result = conn.execute(text(stripped))
            rows = result.fetchall()
            cols = result.keys()
            if not rows:
                return "(0 rows)"
            # Format as markdown table for the LLM
            header = " | ".join(cols)
            separator = " | ".join(["---"] * len(cols))
            body = "\n".join(" | ".join(str(v) for v in row) for row in rows)
            return f"{header}\n{separator}\n{body}\n\n({len(rows)} rows)"
    except Exception as e:
        return f"Error executing query: {e}"

def safe_info(table_names: str) -> str:
    """Schema info for one or more tables (comma-separated)."""
    names = [t.strip() for t in table_names.split(",")]
    return db.get_table_info(table_names=names)

sql_query_tool = Tool(
    name="query_sql_db",
    description=(
        "Execute a SELECT query against the database. "
        "Only SELECT statements allowed. Results capped at 100 rows. "
        "Input: SQL string. Output: markdown table or error."
    ),
    func=safe_query,
)

sql_info_tool = Tool(
    name="info_sql_db",
    description=(
        "Get schema and sample rows for specified tables. "
        "Input: comma-separated table names (e.g., 'customers, orders'). "
        "Output: CREATE TABLE statements with sample data."
    ),
    func=safe_info,
)

tools = [sql_query_tool, sql_info_tool]

Verify: Run a quick REPL test:

from tools import sql_query_tool, sql_info_tool
print(sql_info_tool.invoke("customers"))
print(sql_query_tool.invoke("SELECT * FROM customers LIMIT 2"))

Step 5: Construct the agent with a strict system prompt

The prompt is where most SQL agents fail. Be explicit about: dialect, table discovery flow, error handling, and output format.

# agent.py
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.messages import SystemMessage
from langchain.agents import create_tool_calling_agent, AgentExecutor
from llm import llm
from tools import tools

SYSTEM_PROMPT = """You are a SQL agent for a PostgreSQL database (dialect: postgresql).
Follow this process for every user question:

1. **Discover schema first** — call `info_sql_db` with the relevant table names before writing any query.
   If you don't know which tables, start with `info_sql_db` for all tables.
2. **Write a single SELECT query** — no INSERT, UPDATE, DELETE, DDL, or transactions.
   Use explicit column lists, not `SELECT *`.
   Qualify columns with table names when joining.
   Use `LIMIT 100` (the tool enforces this, but include it anyway).
3. **Execute** — call `query_sql_db` with your SQL.
4. **If the query fails**, read the error, adjust, and retry once. Do not hallucinate fixes.
5. **Answer the user** — summarize results in natural language. Include the SQL you ran.

Rules:
- Never assume column names. Always check `info_sql_db` first.
- Date math: use `CURRENT_DATE - INTERVAL '3 months'` for "last quarter".
- Money is stored in `amount_cents` (integer). Divide by 100.0 for dollars.
- If the user asks for "top N", use `ORDER BY ... DESC LIMIT N`.
- Return only the final answer to the user. Do not emit tool calls in your final response.
"""

prompt = ChatPromptTemplate.from_messages([
    SystemMessage(content=SYSTEM_PROMPT),
    MessagesPlaceholder(variable_name="chat_history", optional=True),
    ("human", "{input}"),
    MessagesPlaceholder(variable_name="agent_scratchpad"),
])

agent = create_tool_calling_agent(llm, tools, prompt)

executor = AgentExecutor(
    agent=agent,
    tools=tools,
    verbose=True,
    max_iterations=6,
    handle_parsing_errors=True,
    return_intermediate_steps=True,
)

Verify: python -c "from agent import executor; print(executor.invoke({'input': 'list all tables'}))" — should return table names without errors.

Step 6: Add a thin CLI for interactive testing

# cli.py
from agent import executor

def main():
    print("SQL Agent ready. Type 'exit' to quit.\n")
    while True:
        try:
            user = input("🗄️  > ").strip()
        except (EOFError, KeyboardInterrupt):
            break
        if user.lower() in {"exit", "quit"}:
            break
        if not user:
            continue
        result = executor.invoke({"input": user})
        print(f"\n{result['output']}\n")

if __name__ == "__main__":
    main()

Run it:

python cli.py

Test queries to verify end-to-end:

Input	Expected behavior
`show me all customers`	Calls `info_sql_db('customers')`, then `SELECT * FROM customers LIMIT 100`
`top 5 customers by total revenue`	Discovers `orders.amount_cents`, joins, aggregates, orders, limits
`how many orders placed in the last 30 days?`	Uses `CURRENT_DATE - INTERVAL '30 days'` on `orders.placed_at`
`delete all orders`	Refuses — tool only allows SELECT

Step 7: Harden for production

Read-only database user

Create a Postgres role with SELECT only:

CREATE ROLE sql_agent_ro LOGIN PASSWORD '...';
GRANT CONNECT ON DATABASE yourdb TO sql_agent_ro;
GRANT USAGE ON SCHEMA public TO sql_agent_ro;
GRANT SELECT ON ALL TABLES IN SCHEMA public TO sql_agent_ro;
ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT SELECT ON TABLES TO sql_agent_ro;

Update DATABASE_URL to use this role.

Query timeout

PostgreSQL statement timeout prevents runaway queries:

# In db.py, when creating the engine:
from sqlalchemy import event

@event.listens_for(engine, "connect")
def set_timeout(dbapi_connection, connection_record):
    cursor = dbapi_connection.cursor()
    cursor.execute("SET statement_timeout = '30s'")
    cursor.close()

Structured logging of intermediate steps

The executor returns intermediate_steps. Log them for debugging and audit:

import json
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("sql_agent")

def run_with_logging(question: str) -> dict:
    result = executor.invoke({"input": question})
    for step in result.get("intermediate_steps", []):
        action, observation = step
        logger.info(json.dumps({
            "tool": action.tool,
            "tool_input": action.tool_input,
            "observation": str(observation)[:500],
        }))
    return result

Metrics: token usage and latency

If you’re routing through a gateway that meters per-token usage, wrap the invoke call:

import time
from langchain_core.callbacks import get_usage_metadata_callback

def run_metered(question: str) -> dict:
    start = time.perf_counter()
    with get_usage_metadata_callback() as cb:
        result = executor.invoke({"input": question})
    latency_ms = (time.perf_counter() - start) * 1000
    usage = cb.usage_metadata
    logger.info(json.dumps({
        "question": question,
        "latency_ms": latency_ms,
        "input_tokens": usage.get("input_tokens"),
        "output_tokens": usage.get("output_tokens"),
    }))
    return result

Step 8: Evaluate with a small test harness

Don’t ship without a regression set. Save 20–30 real questions and expected SQL patterns.

# eval.py
import json
from agent import executor

TEST_CASES = [
    {
        "question": "How many customers do we have?",
        "must_contain": ["COUNT", "customers"],
    },
    {
        "question": "Top 3 customers by revenue",
        "must_contain": ["SUM", "amount_cents", "ORDER BY", "DESC", "LIMIT 3"],
    },
    {
        "question": "Orders from the last 7 days",
        "must_contain": ["INTERVAL '7 days'", "orders.placed_at"],
    },
]

def evaluate():
    passed = 0
    for tc in TEST_CASES:
        result = executor.invoke({"input": tc["question"]})
        sql_calls = [
            step[0].tool_input
            for step in result["intermediate_steps"]
            if step[0].tool == "query_sql_db"
        ]
        sql = " ".join(sql_calls).upper()
        ok = all(token.upper() in sql for token in tc["must_contain"])
        status = "PASS" if ok else "FAIL"
        print(f"[{status}] {tc['question']}")
        if not ok:
            print(f"  Expected tokens: {tc['must_contain']}")
            print(f"  Actual SQL: {sql[:200]}")
        else:
            passed += 1
    print(f"\n{passed}/{len(TEST_CASES)} passed")

if __name__ == "__main__":
    evaluate()

Run it after every schema change or prompt tweak.

Common failure modes and fixes

Symptom	Cause	Fix
Agent hallucinates column names	Skipped `info_sql_db`	Strengthen system prompt: “Never assume column names. Always call info_sql_db first.”
Query returns 0 rows but data exists	Wrong date anchor	Use `CURRENT_DATE` not `NOW()` for date windows; verify timezone
Agent loops on error	`max_iterations` too high or parsing errors not handled	Set `max_iterations=6`, `handle_parsing_errors=True`
`SELECT *` causes wide results	Prompt not explicit	Add “Use explicit column lists, not SELECT *”
Join produces duplicate rows	Missing `DISTINCT` or bad grain	Check primary keys; add `DISTINCT` in prompt rules

What to extend next

Multi-tenant isolation: Add WHERE tenant_id = :current_tenant via a tool wrapper that injects the parameter.
Semantic layer: Replace raw info_sql_db with a curated metrics catalog (dbt metrics, Cube, or a JSON file) so the agent reasons over business concepts, not columns.
Streaming: Use executor.astream_events() for token-by-token UI updates.
Human-in-the-loop: For destructive operations (if you ever allow them), add a confirmation tool that pauses for approval.

The pattern — schema discovery → validated query → execution → natural language answer — scales. The prompt and tool wrappers are your control plane. Invest there.