← Back to DevBytes

Building a SQL Query Agent with LangGraph: Complete Guide

Building a SQL Query Agent with LangGraph: Complete Guide

Large language models are powerful at reasoning over natural language, but they cannot natively inspect your database. A SQL query agent bridges that gap by letting an LLM write, validate, and execute SQL against a real database, then reason over the results to answer user questions. LangGraph, built on top of LangChain, gives you a structured way to model this workflow as a stateful graph of nodes and edges. In this guide, you will build a complete SQL query agent from scratch using LangGraph, a SQLite database, and an OpenAI chat model.

What Is a SQL Query Agent?

A SQL query agent is an autonomous system that takes a natural-language question, translates it into SQL, runs that SQL against a database, inspects the results, and either returns an answer or refines its query if something went wrong. Unlike a single-shot text-to-SQL pipeline, an agent can iterate: it can catch errors, inspect the schema, retry with a corrected query, and explain its findings in plain language.

LangGraph models this behavior as a directed graph. Each node performs a discrete action β€” such as listing tables, retrieving a schema, writing a query, executing it, or generating a final answer β€” and the edges define the flow between those actions. Conditional edges let the agent decide dynamically what to do next based on the current state.

Why It Matters

Prerequisites and Setup

Before you start coding, install the required packages and prepare your environment. You will need Python 3.10 or newer, an OpenAI API key, and a few libraries.

pip install langgraph langchain langchain-openai langchain-community \
            sqlalchemy sqlite3-utils

Set your OpenAI API key as an environment variable:

export OPENAI_API_KEY="sk-your-key-here"

For this tutorial, you will use a small SQLite database so you can run everything locally without external dependencies. Create a script that builds a sample database with a few tables representing an online store.

import sqlite3

def create_sample_db(db_path: str = "store.db") -> None:
    conn = sqlite3.connect(db_path)
    cur = conn.cursor()

    cur.executescript("""
    DROP TABLE IF EXISTS customers;
    DROP TABLE IF EXISTS products;
    DROP TABLE IF EXISTS orders;
    DROP TABLE IF EXISTS order_items;

    CREATE TABLE customers (
        customer_id INTEGER PRIMARY KEY,
        name TEXT NOT NULL,
        email TEXT UNIQUE NOT NULL,
        country TEXT NOT NULL,
        signup_date TEXT NOT NULL
    );

    CREATE TABLE products (
        product_id INTEGER PRIMARY KEY,
        name TEXT NOT NULL,
        category TEXT NOT NULL,
        price REAL NOT NULL,
        stock INTEGER NOT NULL
    );

    CREATE TABLE orders (
        order_id INTEGER PRIMARY KEY,
        customer_id INTEGER NOT NULL,
        order_date TEXT NOT NULL,
        status TEXT NOT NULL,
        FOREIGN KEY (customer_id) REFERENCES customers(customer_id)
    );

    CREATE TABLE order_items (
        order_item_id INTEGER PRIMARY KEY,
        order_id INTEGER NOT NULL,
        product_id INTEGER NOT NULL,
        quantity INTEGER NOT NULL,
        unit_price REAL NOT NULL,
        FOREIGN KEY (order_id) REFERENCES orders(order_id),
        FOREIGN KEY (product_id) REFERENCES products(product_id)
    );

    INSERT INTO customers (name, email, country, signup_date) VALUES
        ('Alice Johnson', 'alice@example.com', 'USA', '2023-01-15'),
        ('Bob Smith', 'bob@example.com', 'UK', '2023-03-22'),
        ('Carol Lee', 'carol@example.com', 'Singapore', '2023-06-10'),
        ('David MΓΌller', 'david@example.com', 'Germany', '2023-09-05'),
        ('Eve Garcia', 'eve@example.com', 'Spain', '2024-01-20');

    INSERT INTO products (name, category, price, stock) VALUES
        ('Wireless Mouse', 'Electronics', 29.99, 150),
        ('Mechanical Keyboard', 'Electronics', 89.99, 80),
        ('USB-C Hub', 'Electronics', 49.99, 200),
        ('Notebook Pro', 'Stationery', 12.50, 500),
        ('Desk Lamp', 'Furniture', 39.99, 60);

    INSERT INTO orders (customer_id, order_date, status) VALUES
        (1, '2024-02-01', 'delivered'),
        (2, '2024-02-15', 'delivered'),
        (1, '2024-03-10', 'shipped'),
        (3, '2024-03-20', 'pending'),
        (4, '2024-04-01', 'delivered'),
        (5, '2024-04-12', 'cancelled');

    INSERT INTO order_items (order_id, product_id, quantity, unit_price) VALUES
        (1, 1, 2, 29.99),
        (1, 3, 1, 49.99),
        (2, 2, 1, 89.99),
        (3, 4, 5, 12.50),
        (4, 5, 2, 39.99),
        (5, 1, 3, 29.99),
        (6, 2, 1, 89.99);
    """)

    conn.commit()
    conn.close()
    print(f"Sample database created at {db_path}")

if __name__ == "__main__":
    create_sample_db()

Run this script once to generate store.db. You now have a realistic schema with customers, products, orders, and line items β€” enough to ask interesting analytical questions.

Architecture Overview

Your agent will follow a simple but powerful loop. The graph has the following nodes:

The key design decision is the conditional edge after execute_query. If the query succeeds, the graph flows to generate_answer. If it fails, the graph loops back to write_query so the model can correct its mistake. To prevent infinite loops, you will track the number of attempts in the state and cap it at a reasonable limit.

Defining the Agent State

LangGraph uses a typed state object that every node reads from and writes to. Define it using Python's TypedDict so the structure is explicit and type-checked.

from typing import TypedDict, Annotated, List, Optional
from langgraph.graph.message import add_messages

class AgentState(TypedDict):
    question: str
    tables: List[str]
    schema: str
    query: str
    query_result: str
    query_error: Optional[str]
    attempts: int
    answer: str

Each field serves a specific purpose. The question is the original user input. tables and schema hold metadata about the database. query stores the most recent SQL string. query_result and query_error capture the outcome of execution. attempts is a counter that prevents infinite retry loops, and answer holds the final natural-language response.

Building the Database Tools

Before defining graph nodes, create utility functions that interact with SQLite. Keeping these separate from the graph logic makes the code easier to test and reuse.

import sqlite3
from typing import List

DB_PATH = "store.db"

def get_table_names() -> List[str]:
    conn = sqlite3.connect(DB_PATH)
    cur = conn.cursor()
    cur.execute(
        "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%';"
    )
    tables = [row[0] for row in cur.fetchall()]
    conn.close()
    return tables

def get_table_schema(table_name: str) -> str:
    conn = sqlite3.connect(DB_PATH)
    cur = conn.cursor()
    cur.execute(f"PRAGMA table_info({table_name});")
    columns = cur.fetchall()
    conn.close()
    lines = []
    for col in columns:
        col_id, name, col_type, not_null, default, pk = col
        constraints = []
        if pk:
            constraints.append("PRIMARY KEY")
        if not_null:
            constraints.append("NOT NULL")
        if default is not None:
            constraints.append(f"DEFAULT {default}")
        constraint_str = " ".join(constraints)
        lines.append(f"  {name} {col_type} {constraint_str}".strip())
    return f"CREATE TABLE {table_name} (\n" + ",\n".join(lines) + "\n);"

def get_full_schema(table_names: List[str]) -> str:
    return "\n\n".join(get_table_schema(t) for t in table_names)

def execute_sql(query: str) -> tuple:
    """Returns (results_string, error_string). One of them is always None."""
    conn = sqlite3.connect(DB_PATH)
    try:
        cur = conn.cursor()
        cur.execute(query)
        rows = cur.fetchall()
        col_names = [desc[0] for desc in cur.description] if cur.description else []
        if not rows:
            return ("Query executed successfully. No rows returned.", None)
        formatted = ", ".join(col_names) + "\n"
        formatted += "\n".join(
            ", ".join(str(val) for val in row) for row in rows
        )
        return (formatted, None)
    except Exception as e:
        return (None, str(e))
    finally:
        conn.close()

These functions are intentionally simple. get_table_names lists tables, get_full_schema builds a textual schema description, and execute_sql runs a query and returns either a formatted result string or an error message. In a production system, you would add query timeouts, row limits, and read-only enforcement here.

Initializing the Language Model

Create a single LLM instance that all nodes will share. Using langchain-openai gives you a consistent interface and built-in retry logic.

from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model="gpt-4o", temperature=0)

Setting temperature=0 is important for SQL generation. You want deterministic, reproducible output rather than creative variation. For more complex agents, you might use a higher temperature for the final answer generation node, but for this guide, a single low-temperature model works well.

Defining Graph Nodes

Each node is a function that takes the current state and returns a dictionary of updates. LangGraph merges these updates into the state automatically.

Node 1: List Tables

def list_tables_node(state: AgentState) -> dict:
    tables = get_table_names()
    print(f"[list_tables] Found tables: {tables}")
    return {"tables": tables}

This node simply queries the database for table names and stores them in state. It runs once at the beginning of every query.

Node 2: Get Schema

def get_schema_node(state: AgentState) -> dict:
    schema = get_full_schema(state["tables"])
    print(f"[get_schema] Retrieved schema for {len(state['tables'])} tables")
    return {"schema": schema}

With the table names in hand, this node builds a complete schema string. The LLM will use this schema to understand column names, types, and constraints when writing SQL.

Node 3: Write Query

def write_query_node(state: AgentState) -> dict:
    schema = state["schema"]
    question = state["question"]
    previous_error = state.get("query_error")
    previous_query = state.get("query")

    system_prompt = """You are an expert SQLite developer. Given a database schema and a
natural-language question, write a single SQL query that answers the question.

Rules:
- Use only SQLite-compatible syntax.
- Return ONLY the SQL query. Do not include explanations, markdown, or backticks.
- If correcting a previous error, fix the specific issue that caused the failure.
- Limit results to 50 rows unless the question asks for an aggregate or count.
"""

    user_content = f"Database schema:\n\n{schema}\n\nQuestion: {question}"

    if previous_error and previous_query:
        user_content += (
            f"\n\nPrevious query that failed:\n{previous_query}"
            f"\n\nError:\n{previous_error}"
            f"\n\nPlease write a corrected query."
        )

    messages = [
        ("system", system_prompt),
        ("human", user_content),
    ]

    response = llm.invoke(messages)
    query = response.content.strip()

    # Strip markdown code fences if the model added them
    if query.startswith(""):
        lines = query.split("\n")
        query = "\n".join(lines[1:-1] if lines[-1].strip() == "" else lines[1:])

    attempts = state.get("attempts", 0) + 1
    print(f"[write_query] Attempt {attempts}: {query}")
    return {"query": query, "attempts": attempts, "query_error": None}

This is the core reasoning node. The system prompt is explicit about output format to reduce parsing issues. If the node is being called as a retry, it includes the previous query and error so the model can fix the specific problem. The code also strips markdown code fences, which models sometimes add despite instructions not to.

Node 4: Execute Query

def execute_query_node(state: AgentState) -> dict:
    query = state["query"]
    print(f"[execute_query] Running: {query}")
    result, error = execute_sql(query)
    if error:
        print(f"[execute_query] Error: {error}")
        return {"query_error": error, "query_result": None}
    print(f"[execute_query] Success: {result[:200]}...")
    return {"query_result": result, "query_error": None}

This node runs the SQL and stores either the result or the error. The conditional edge that follows will use query_error to decide whether to retry or proceed.

Node 5: Generate Answer

def generate_answer_node(state: AgentState) -> dict:
    question = state["question"]
    query = state["query"]
    result = state["query_result"]

    system_prompt = """You are a helpful data analyst. Given a user's question, the SQL
query that was executed, and the query results, provide a clear, concise answer in
natural language. If the results are empty or unhelpful, say so honestly. Reference
specific numbers from the results when relevant."""

    user_content = (
        f"Question: {question}\n\n"
        f"SQL Query: {query}\n\n"
        f"Query Results:\n{result}"
    )

    messages = [
        ("system", system_prompt),
        ("human", user_content),
    ]

    response = llm.invoke(messages)
    answer = response.content.strip()
    print(f"[generate_answer] {answer}")
    return {"answer": answer}

The final node synthesizes a human-readable answer. It has access to the original question, the executed query, and the raw results, so it can explain what the data means rather than just repeating numbers.

Wiring Up the Graph

Now that all nodes are defined, assemble them into a LangGraph StateGraph. The graph starts at list_tables, flows linearly through schema retrieval and query writing, then branches based on execution success.

from langgraph.graph import StateGraph, START, END

MAX_ATTEMPTS = 3

def should_retry_or_finish(state: AgentState) -> str:
    """Conditional edge after execute_query."""
    if state.get("query_error"):
        if state.get("attempts", 0) >= MAX_ATTEMPTS:
            print(f"[router] Max attempts ({MAX_ATTEMPTS}) reached. Giving up.")
            return "generate_answer"
        print("[router] Query failed. Retrying with corrected query.")
        return "write_query"
    print("[router] Query succeeded. Generating answer.")
    return "generate_answer"

# Build the graph
workflow = StateGraph(AgentState)

# Add nodes
workflow.add_node("list_tables", list_tables_node)
workflow.add_node("get_schema", get_schema_node)
workflow.add_node("write_query", write_query_node)
workflow.add_node("execute_query", execute_query_node)
workflow.add_node("generate_answer", generate_answer_node)

# Add edges
workflow.add_edge(START, "list_tables")
workflow.add_edge("list_tables", "get_schema")
workflow.add_edge("get_schema", "write_query")
workflow.add_edge("write_query", "execute_query")

# Conditional edge: retry or finish
workflow.add_conditional_edges(
    "execute_query",
    should_retry_or_finish,
    {
        "write_query": "write_query",
        "generate_answer": "generate_answer",
    },
)

workflow.add_edge("generate_answer", END)

# Compile
app = workflow.compile()

The should_retry_or_finish function is the decision point. If there is an error and attempts remain, it routes back to write_query. If the query succeeded, or if the maximum number of attempts is exhausted, it routes to generate_answer. This creates the self-correcting loop that makes the agent robust.

Running the Agent

To invoke the agent, pass an initial state with the user's question. LangGraph will execute the nodes in order and return the final state.

def ask(question: str) -> str:
    initial_state = {
        "question": question,
        "tables": [],
        "schema": "",
        "query": "",
        "query_result": "",
        "query_error": None,
        "attempts": 0,
        "answer": "",
    }
    final_state = app.invoke(initial_state)
    return final_state["answer"]

if __name__ == "__main__":
    questions = [
        "Which customer has placed the most orders?",
        "What is the total revenue from delivered orders?",
        "Which product category has the highest average price?",
        "List all customers who have never received a delivered order.",
    ]

    for q in questions:
        print(f"\n{'='*60}")
        print(f"Question: {q}")
        print(f"{'='*60}")
        answer = ask(q)
        print(f"\nAnswer: {answer}")

When you run this, you will see the agent's step-by-step progress in the console. For a question like "What is the total revenue from delivered orders?", the output looks something like this:

[list_tables] Found tables: ['customers', 'products', 'orders', 'order_items']
[get_schema] Retrieved schema for 4 tables
[write_query] Attempt 1: SELECT ROUND(SUM(oi.quantity * oi.unit_price), 2) AS total_revenue
FROM order_items oi
JOIN orders o ON oi.order_id = o.order_id
WHERE o.status = 'delivered';
[execute_query] Running: SELECT ROUND(SUM(oi.quantity * oi.unit_price), 2) ...
[execute_query] Success: total_revenue
329.95...
[router] Query succeeded. Generating answer.
[generate_answer] The total revenue from delivered orders is $329.95.

Answer: The total revenue from delivered orders is $329.95.

If the model makes a mistake β€” for example, referencing a column that does not exist β€” you will see the retry loop in action:

[write_query] Attempt 1: SELECT customer_name, COUNT(*) FROM customers ...
[execute_query] Error: no such column: customer_name
[router] Query failed. Retrying with corrected query.
[write_query] Attempt 2: SELECT c.name, COUNT(*) AS order_count FROM customers c ...
[execute_query] Success: ...
[router] Query succeeded. Generating answer.

Adding a Human-in-the-Loop Checkpoint

For production systems, you may want a human to approve queries before execution, especially for write operations. LangGraph supports this through checkpointing and interrupts. Here is how to add an approval step.

from langgraph.checkpoint.memory import MemorySaver

def write_query_with_interrupt_node(state: AgentState) -> dict:
    # Same logic as write_query_node, but the graph will pause
    # after this node returns because of the interrupt configuration.
    return write_query_node(state)

workflow_with_checkpoint = StateGraph(AgentState)

workflow_with_checkpoint.add_node("list_tables", list_tables_node)
workflow_with_checkpoint.add_node("get_schema", get_schema_node)
workflow_with_checkpoint.add_node("write_query", write_query_with_interrupt_node)
workflow_with_checkpoint.add_node("execute_query", execute_query_node)
workflow_with_checkpoint.add_node("generate_answer", generate_answer_node)

workflow_with_checkpoint.add_edge(START, "list_tables")
workflow_with_checkpoint.add_edge("list_tables", "get_schema")
workflow_with_checkpoint.add_edge("get_schema", "write_query")
workflow_with_checkpoint.add_edge("write_query", "execute_query")
workflow_with_checkpoint.add_conditional_edges(
    "execute_query",
    should_retry_or_finish,
    {
        "write_query": "write_query",
        "generate_answer": "generate_answer",
    },
)
workflow_with_checkpoint.add_edge("generate_answer", END)

checkpointer = MemorySaver()
app_with_checkpoint = workflow_with_checkpoint.compile(
    checkpointer=checkpointer,
    interrupt_before=["execute_query"],
)

# Usage with manual approval
config = {"configurable": {"thread_id": "thread-1"}}
initial_state = {
    "question": "How many products are out of stock?",
    "tables": [], "schema": "", "query": "",
    "query_result": "", "query_error": None,
    "attempts": 0, "answer": "",
}

# Runs up to the interrupt (before execute_query)
partial_state = app_with_checkpoint.invoke(initial_state, config=config)

print(f"Proposed query: {partial_state['query']}")
print("Approve? (yes/no)")
approval = input().strip().lower()

if approval == "yes":
    final_state = app_with_checkpoint.invoke(None, config=config)
    print(f"Answer: {final_state['answer']}")
else:
    print("Query rejected by user.")

The interrupt_before parameter tells LangGraph to pause execution right before the execute_query node. The thread_id in the config identifies the conversation so the checkpointer can resume it. Calling invoke(None, config=config) resumes from the checkpoint, executing the remaining nodes.

Best Practices

Enforce Read-Only Access

In production, never let an agent run arbitrary SQL against your primary database. Use a read-only replica or a database user with only SELECT permissions. You can also validate queries programmatically before execution:

import re

def is_read_only(query: str) -> bool:
    dangerous_keywords = [
        "INSERT", "UPDATE", "DELETE", "DROP", "ALTER",
        "CREATE", "TRUNCATE", "REPLACE", "ATTACH", "DETACH",
    ]
    # Remove string literals to avoid false positives
    cleaned = re.sub(r"'[^']*'", "", query, flags=re.IGNORECASE)
    cleaned = re.sub(r'"[^"]*"', "", cleaned, flags=re.IGNORECASE)
    for keyword in dangerous_keywords:
        if re.search(rf"\b{keyword}\b", cleaned, flags=re.IGNORECASE):
            return False
    return True

def execute_sql_safe(query: str) -> tuple:
    if not is_read_only(query):
        return (None, "Blocked: only SELECT queries are allowed.")
    return execute_sql(query)

Limit Result Rows

Large result sets can blow up your context window and cost. Add a LIMIT clause enforcement or truncate results in the execution function:

def execute_sql_limited(query: str, max_rows: int = 50) -> tuple:
    conn = sqlite3.connect(DB_PATH)
    try:
        cur = conn.cursor()
        cur.execute(query)
        rows = cur.fetchmany(max_rows + 1)
        col_names = [desc[0] for desc in cur.description] if cur.description else []
        truncated = len(rows) > max_rows
        rows = rows[:max_rows]
        if not rows:
            return ("No rows returned.", None)
        formatted = " | ".join(col_names) + "\n"
        formatted += "\n".join(" | ".join(str(v) for v in row) for row in rows)
        if truncated:
            formatted += f"\n... (results truncated to {max_rows} rows)"
        return (formatted, None)
    except Exception as e:
        return (None, str(e))
    finally:
        conn.close()

Provide Few-Shot Examples

LLMs write better SQL when they see examples of correct queries for your specific schema. Add a few-shot prompt section to the write_query node:

FEW_SHOT_EXAMPLES = """
Example 1:
Q: How many customers are from the USA?
SQL: SELECT COUNT(*) FROM customers WHERE country = 'USA';

Example 2:
Q: What is the average order value?
SQL: SELECT ROUND(AVG(total), 2) FROM (SELECT o.order_id, SUM(oi.quantity * oi.unit_price) AS total FROM orders o JOIN order_items oi ON o.order_id = oi.order_id GROUP BY o.order_id);
"""

# In write_query_node, append to user_content:
user_content += f"\n\nExamples:\n{FEW_SHOT_EXAMPLES}"

Log Every Query

For auditing and debugging, log every query the agent executes along with its outcome. This creates a trail you can review later:

import json
from datetime import datetime

def log_query(question: str, query: str, result: str, error: str) -> None:
    entry = {
        "timestamp": datetime.now().isoformat(),
        "question": question,
        "query": query,
        "success": error is None,
        "error": error,
        "result_preview": result[:500] if result else None,
    }
    with open("query_log.jsonl", "a") as f:
        f.write(json.dumps(entry) + "\n")

Choose the Right Model

SQL generation is a task where model quality matters significantly. GPT-4o and Claude 3.5 Sonnet tend to produce more accurate SQL than smaller models. If cost is a concern, use a smaller model for the final answer generation (which is easier) and a stronger model for query writing (which is harder). You can instantiate two separate LLM objects and use them in different nodes.

Handle Ambiguous Questions

Not every user question maps cleanly to a single query. Consider adding a clarification node that asks the user for more details when the question is ambiguous. You can implement this as a conditional edge after get_schema that uses the LLM to assess whether the question is answerable with the available schema.

Conclusion

Building a SQL query agent with LangGraph gives you a robust, transparent, and self-correcting system for natural-language database interaction. By modeling the workflow as a graph with conditional edges, you get iterative error correction for free β€” the agent can read its own mistakes and try again. The separation of concerns between tools, nodes, and routing logic makes the system easy to extend: you can add few-shot examples, human approval gates, read-only enforcement, and result limits without restructuring the core graph. Start with the simple loop presented here, then layer in the production safeguards as your use case demands. With careful prompt engineering, sensible limits, and the right model, a LangGraph SQL agent can reliably turn natural-language questions into accurate, data-driven answers.

β€” Ad β€”

Google AdSense will appear here after approval

← Back to all articles