Building a SQL Query Agent with LlamaIndex: Complete Guide

Modern applications increasingly rely on natural language interfaces to unlock data locked away in relational databases. Instead of forcing users to learn SQL or rely on a handful of pre-built dashboards, a SQL Query Agent lets them ask plain-English questions and receive accurate, data-backed answers. LlamaIndex, one of the leading frameworks for building LLM-powered applications, provides robust tooling for exactly this use case. In this guide, we'll walk through everything you need to build, test, and harden a SQL Query Agent using LlamaIndex.

What Is a SQL Query Agent?

A SQL Query Agent is an LLM-driven system that translates natural language questions into SQL queries, executes those queries against a database, and returns the results in a human-readable form. Unlike a simple text-to-SQL model that just emits a query string, an agent operates in a loop: it reasons about the question, decides which tables and columns are relevant, writes SQL, validates it, runs it, and—if something fails—corrects itself and tries again.

LlamaIndex provides this capability through its SQLDatabase abstraction and the NLSQLTableQueryEngine and agent-based toolkits. The agent layer adds planning, tool selection, and multi-step reasoning on top of the core query engine.

Why It Matters

Democratizes data access: Non-technical stakeholders can query complex schemas without writing SQL.
Reduces dashboard sprawl: Instead of building dozens of static reports, one agent can answer ad-hoc questions.
Self-correcting behavior: Agents can catch syntax errors, retry failed queries, and explain their reasoning.
Composable with other tools: A SQL agent can be combined with vector search, web lookup, or calculation tools to answer richer questions.
Auditability: Because the agent emits SQL, you can log, review, and optimize every query it runs.

Prerequisites and Setup

Before we start coding, make sure you have the following:

Python 3.9 or higher
An OpenAI API key (or another supported LLM provider)
A relational database — for this tutorial we'll use SQLite for simplicity, but the same approach works with PostgreSQL, MySQL, Snowflake, and others via SQLAlchemy

Install the required packages:

pip install llama-index llama-index-llms-openai llama-index-embeddings-openai sqlalchemy

Set your API key as an environment variable:

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

Step 1: Preparing a Sample Database

Let's create a small SQLite database with a couple of related tables so we have something realistic to query. We'll model a simple e-commerce schema with customers, orders, and products.

import sqlite3

conn = sqlite3.connect("shop.db")
cursor = conn.cursor()

cursor.executescript("""
CREATE TABLE IF NOT EXISTS customers (
    customer_id INTEGER PRIMARY KEY,
    name TEXT NOT NULL,
    email TEXT,
    country TEXT,
    signup_date TEXT
);

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

CREATE TABLE IF NOT EXISTS orders (
    order_id INTEGER PRIMARY KEY,
    customer_id INTEGER,
    product_id INTEGER,
    quantity INTEGER,
    order_date TEXT,
    total REAL,
    FOREIGN KEY (customer_id) REFERENCES customers(customer_id),
    FOREIGN KEY (product_id) REFERENCES products(product_id)
);
""")

# Insert sample data
cursor.executemany(
    "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"),
        ("Carla Mendes", "carla@example.com", "Brazil", "2023-06-10"),
        ("David Lee", "david@example.com", "USA", "2023-08-05"),
    ],
)

cursor.executemany(
    "INSERT INTO products (name, category, price) VALUES (?, ?, ?)",
    [
        ("Wireless Mouse", "Electronics", 25.99),
        ("Mechanical Keyboard", "Electronics", 89.99),
        ("Coffee Mug", "Kitchen", 12.50),
        ("Notebook", "Stationery", 4.99),
    ],
)

cursor.executemany(
    "INSERT INTO orders (customer_id, product_id, quantity, order_date, total) VALUES (?, ?, ?, ?, ?)",
    [
        (1, 1, 2, "2023-09-01", 51.98),
        (1, 3, 5, "2023-09-03", 62.50),
        (2, 2, 1, "2023-09-10", 89.99),
        (3, 4, 10, "2023-09-12", 49.90),
        (4, 1, 1, "2023-09-15", 25.99),
        (4, 2, 2, "2023-09-20", 179.98),
    ],
)

conn.commit()
conn.close()
print("Database created successfully.")

Run this script once to generate shop.db. This file will be the data source for our agent.

Step 2: Connecting LlamaIndex to the Database

LlamaIndex wraps databases through the SQLDatabase class, which uses SQLAlchemy under the hood. This abstraction reads the schema metadata and exposes it to the LLM so it knows which tables and columns exist.

from llama_index.core import SQLDatabase
from sqlalchemy import create_engine

engine = create_engine("sqlite:///shop.db")
sql_database = SQLDatabase(engine, include_tables=["customers", "products", "orders"])

# Quick sanity check: print the schema
print(sql_database.get_schema())

The get_schema() method returns a textual description of the tables, columns, and foreign key relationships. This is exactly what the LLM will see when deciding how to construct a query.

Step 3: Building a Basic NLSQL Query Engine

Before jumping to a full agent, let's build the simpler NLSQLTableQueryEngine. This engine takes a natural language question, generates SQL, executes it, and synthesizes a natural language answer from the results.

from llama_index.core.query_engine import NLSQLTableQueryEngine
from llama_index.llms.openai import OpenAI

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

query_engine = NLSQLTableQueryEngine(
    sql_database=sql_database,
    tables=["customers", "products", "orders"],
    llm=llm,
)

response = query_engine.query("Which customer spent the most money in total?")
print(response.response)

# Inspect the generated SQL
print("\nGenerated SQL:")
print(response.metadata["sql_query"])

When you run this, the engine will produce something like:

David Lee spent the most money in total, with $205.97 in purchases.

Generated SQL:
SELECT customers.name, SUM(orders.total) AS total_spent
FROM customers
JOIN orders ON customers.customer_id = orders.customer_id
GROUP BY customers.customer_id
ORDER BY total_spent DESC
LIMIT 1

This works well for straightforward questions, but it has limitations. If the generated SQL fails, the engine doesn't retry. If the question requires multiple steps—like first computing an aggregate and then filtering on it—the single-shot engine may struggle. That's where an agent comes in.

Step 4: Upgrading to a SQL Query Agent

An agent can call tools iteratively, observe results, and adapt. LlamaIndex provides SQLAutoVectorQueryEngine and, more importantly, the QueryEngineTool wrapper that lets us plug our SQL engine into a general-purpose agent. Let's build a proper agent using the FunctionAgent (or ReActAgent) with SQL as one of its tools.

from llama_index.core.agent import FunctionAgent
from llama_index.core.tools import QueryEngineTool, ToolMetadata

# Wrap the SQL query engine as a tool
sql_tool = QueryEngineTool.from_defaults(
    query_engine=query_engine,
    name="sql_query",
    description=(
        "Useful for translating natural language questions into SQL queries "
        "over the e-commerce database. The database contains tables: "
        "customers, products, and orders. Provide a clear natural language question."
    ),
)

# Build the agent
agent = FunctionAgent(
    tools=[sql_tool],
    llm=llm,
    system_prompt=(
        "You are a helpful data analyst assistant. "
        "When asked a question about the database, use the sql_query tool. "
        "Always explain your reasoning briefly before giving the final answer. "
        "If a query fails, try to fix the SQL and retry."
    ),
)

Now let's ask it a more complex question:

response = agent.chat(
    "What is the average order total for each product category, "
    "and which category has the highest average?"
)
print(response.response)

The agent will reason about the question, call the sql_query tool, receive the result, and synthesize a final answer. Because it's an agent, you can also ask follow-up questions in the same conversation:

response = agent.chat("Now list the top 2 customers from that category's country.")
print(response.response)

Step 5: Adding Schema Context and Custom Prompts

Out of the box, the LLM only sees the raw schema. In real-world databases, column names are often cryptic and business logic is implicit. You can dramatically improve accuracy by providing additional context.

from llama_index.core import PromptTemplate

# Provide extra context about the schema
context_query_engine = NLSQLTableQueryEngine(
    sql_database=sql_database,
    tables=["customers", "products", "orders"],
    llm=llm,
    context_str_prefix=(
        "Database context: This is an e-commerce database. "
        "The 'total' column in orders represents the final amount paid in USD. "
        "Signup dates are in ISO format (YYYY-MM-DD). "
        "Categories are case-sensitive strings."
    ),
)

context_tool = QueryEngineTool.from_defaults(
    query_engine=context_query_engine,
    name="sql_query",
    description="Query the e-commerce database with natural language.",
)

context_agent = FunctionAgent(
    tools=[context_tool],
    llm=llm,
    system_prompt="You are a precise data analyst. Always cite the numbers you use.",
)

response = context_agent.chat("How many customers signed up in 2023?")
print(response.response)

Step 6: Handling Errors and Retries

One of the biggest advantages of an agent over a single-shot engine is graceful error handling. If the LLM generates invalid SQL, the agent can observe the error message and try again. LlamaIndex surfaces execution errors back to the agent automatically when using the query engine tool, but you can also build a custom tool for finer control.

from llama_index.core.tools import FunctionTool
from sqlalchemy import text

def run_sql(query: str) -> str:
    """Execute a raw SQL query and return results as a string."""
    try:
        with engine.connect() as connection:
            result = connection.execute(text(query))
            rows = result.fetchmany(50)
            if not rows:
                return "Query executed successfully but returned no rows."
            columns = list(result.keys())
            formatted = [", ".join(f"{c}={v}" for c, v in zip(columns, row)) for row in rows]
            return "\n".join(formatted)
    except Exception as e:
        return f"SQL ERROR: {e}. Please fix the query and try again."

raw_sql_tool = FunctionTool.from_defaults(
    fn=run_sql,
    name="run_sql",
    description="Execute a SQL query directly. Use this if sql_query fails.",
)

robust_agent = FunctionAgent(
    tools=[sql_tool, raw_sql_tool],
    llm=llm,
    system_prompt=(
        "You are a SQL expert agent. First try the sql_query tool. "
        "If it fails or returns unexpected results, use run_sql to write "
        "and debug SQL manually. Always verify your results before answering."
    ),
)

response = robust_agent.chat("Show me all orders placed in September 2023.")
print(response.response)

Step 7: Adding Row-Level Security and Read-Only Guards

When deploying a SQL agent in production, safety is paramount. You should never let an LLM execute arbitrary write operations. There are two layers of defense you should always implement.

First, restrict the database user to read-only permissions. For SQLite, you can use a read-only connection string. For PostgreSQL or MySQL, create a dedicated role with only SELECT privileges.

Second, validate queries programmatically before execution:

import re

BLOCKED_KEYWORDS = ["insert", "update", "delete", "drop", "alter", "truncate", "create"]

def is_safe_query(sql: str) -> bool:
    lowered = sql.lower()
    return not any(kw in lowered for kw in BLOCKED_KEYWORDS)

def safe_run_sql(query: str) -> str:
    if not is_safe_query(query):
        return "ERROR: Only SELECT queries are allowed."
    return run_sql(query)

safe_sql_tool = FunctionTool.from_defaults(
    fn=safe_run_sql,
    name="safe_run_sql",
    description="Execute a read-only SELECT query. Write operations are blocked.",
)

Best Practices

1. Keep Schemas Focused

Don't expose every table in a large database to the agent. Use the include_tables parameter to limit the schema to what's relevant. A smaller schema reduces token usage and improves accuracy.

2. Provide Rich Context

Column names like cst_amt_pd are meaningless to an LLM. Always include a context string that explains abbreviations, units, and business rules. Consider maintaining a YAML or JSON file that maps tables and columns to human-readable descriptions.

3. Use the Right Model

Text-to-SQL is a reasoning-heavy task. Smaller or older models often produce syntactically invalid SQL. GPT-4o, Claude 3.5 Sonnet, or similar frontier models are recommended. Set temperature=0 for deterministic, reproducible outputs.

4. Limit Result Sets

LLMs have context limits. If a query returns thousands of rows, the response will be truncated or expensive. Add LIMIT clauses in your prompts, or enforce them in your custom execution tool as shown in the fetchmany(50) example above.

5. Log Every Query

Always log the generated SQL alongside the user's question. This creates an audit trail, helps you debug failures, and surfaces common questions that might warrant a dedicated dashboard or view.

import logging

logging.basicConfig(filename="sql_agent.log", level=logging.INFO)

def logged_run_sql(query: str) -> str:
    logging.info(f"SQL executed: {query}")
    return safe_run_sql(query)

6. Test with a Golden Dataset

Maintain a set of question-SQL pairs that you run as regression tests whenever you change prompts, models, or schema. This catches regressions before they reach production.

golden_tests = [
    ("How many customers are from the USA?", "SELECT COUNT(*) FROM customers WHERE country = 'USA'"),
    ("What is the most expensive product?", "SELECT name FROM products ORDER BY price DESC LIMIT 1"),
]

for question, expected_sql in golden_tests:
    result = query_engine.query(question)
    generated = result.metadata.get("sql_query", "")
    status = "PASS" if expected_sql.strip().lower() in generated.lower() else "CHECK"
    print(f"{status}: {question}")

7. Consider Few-Shot Examples

For complex schemas, include a few example question-SQL pairs in your context. This dramatically improves accuracy on edge cases.

few_shot_context = """
Example questions and their SQL:
Q: How many orders did Alice place?
A: SELECT COUNT(*) FROM orders JOIN customers ON orders.customer_id = customers.customer_id WHERE customers.name = 'Alice Johnson';
Q: What is the total revenue from Electronics?
A: SELECT SUM(orders.total) FROM orders JOIN products ON orders.product_id = products.product_id WHERE products.category = 'Electronics';
"""

Putting It All Together

Here's a consolidated, production-ready script that combines everything we've covered:

import logging
from sqlalchemy import create_engine, text
from llama_index.core import SQLDatabase
from llama_index.core.query_engine import NLSQLTableQueryEngine
from llama_index.core.agent import FunctionAgent
from llama_index.core.tools import QueryEngineTool, FunctionTool
from llama_index.llms.openai import OpenAI

logging.basicConfig(filename="sql_agent.log", level=logging.INFO)

# 1. Connect to the database
engine = create_engine("sqlite:///shop.db")
sql_database = SQLDatabase(engine, include_tables=["customers", "products", "orders"])

# 2. Configure the LLM
llm = OpenAI(model="gpt-4o", temperature=0)

# 3. Build the NLSQL query engine with context
query_engine = NLSQLTableQueryEngine(
    sql_database=sql_database,
    tables=["customers", "products", "orders"],
    llm=llm,
    context_str_prefix=(
        "E-commerce database. 'total' in orders is USD. "
        "Dates are ISO format (YYYY-MM-DD)."
    ),
)

# 4. Wrap as a tool
sql_tool = QueryEngineTool.from_defaults(
    query_engine=query_engine,
    name="sql_query",
    description="Answer natural language questions about the e-commerce database.",
)

# 5. Safe raw SQL tool with logging
BLOCKED = ["insert", "update", "delete", "drop", "alter", "truncate", "create"]

def safe_run_sql(query: str) -> str:
    lowered = query.lower()
    if any(kw in lowered for kw in BLOCKED):
        return "ERROR: Only SELECT queries are allowed."
    logging.info(f"SQL: {query}")
    try:
        with engine.connect() as conn:
            result = conn.execute(text(query))
            rows = result.fetchmany(50)
            if not rows:
                return "No rows returned."
            cols = list(result.keys())
            return "\n".join(", ".join(f"{c}={v}" for c, v in zip(cols, r)) for r in rows)
    except Exception as e:
        return f"SQL ERROR: {e}"

raw_tool = FunctionTool.from_defaults(
    fn=safe_run_sql,
    name="run_sql",
    description="Execute a read-only SQL query directly for debugging.",
)

# 6. Build the agent
agent = FunctionAgent(
    tools=[sql_tool, raw_tool],
    llm=llm,
    system_prompt=(
        "You are a precise data analyst. Use sql_query first. "
        "If it fails, use run_sql to debug. Always verify results before answering."
    ),
)

# 7. Run a query
response = agent.chat("Which product category generated the most revenue?")
print(response.response)

Conclusion

Building a SQL Query Agent with LlamaIndex bridges the gap between powerful language models and the structured data that drives business decisions. By combining the NLSQLTableQueryEngine with an agent loop, you get a system that not only translates questions into SQL but also reasons about results, recovers from errors, and explains its answers. The key to a reliable production deployment lies in the details: focused schemas, rich context, read-only guards, query logging, and a golden test suite. With these pieces in place, you can ship a natural language data interface that is both genuinely useful and safe enough for real users.