Build a LangChain SQL Agent for Text-to-SQL Queries (2026)

pip install langchain langchain-openai langchain-community \
    sqlalchemy python-dotenv fastapi uvicorn

# Create a demo database for testing
import sqlite3

def create_demo_database(db_path: str = "demo.db"):
    conn = sqlite3.connect(db_path)
    cursor = conn.cursor()

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

        CREATE TABLE customers (
            id INTEGER PRIMARY KEY,
            name TEXT NOT NULL,
            email TEXT UNIQUE,
            city TEXT,
            created_at DATE
        );

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

        CREATE TABLE orders (
            id INTEGER PRIMARY KEY,
            customer_id INTEGER REFERENCES customers(id),
            product_id INTEGER REFERENCES products(id),
            quantity INTEGER,
            total_amount REAL,
            order_date DATE,
            status TEXT
        );

        INSERT INTO customers VALUES
            (1, 'Alice Johnson', 'alice@example.com', 'New York', '2025-01-15'),
            (2, 'Bob Smith', 'bob@example.com', 'Los Angeles', '2025-02-20'),
            (3, 'Carol White', 'carol@example.com', 'Chicago', '2025-03-10'),
            (4, 'David Brown', 'david@example.com', 'New York', '2025-04-05');

        INSERT INTO products VALUES
            (1, 'Python Course', 'Education', 149.99, 999),
            (2, 'AI Toolkit Pro', 'Software', 299.99, 500),
            (3, 'Data Science Bundle', 'Education', 199.99, 750),
            (4, 'Cloud Computing Guide', 'Education', 89.99, 1000);

        INSERT INTO orders VALUES
            (1, 1, 1, 2, 299.98, '2025-05-01', 'completed'),
            (2, 1, 2, 1, 299.99, '2025-05-15', 'completed'),
            (3, 2, 3, 1, 199.99, '2025-05-20', 'pending'),
            (4, 3, 1, 1, 149.99, '2025-06-01', 'completed'),
            (5, 4, 4, 3, 269.97, '2025-06-10', 'shipped'),
            (6, 2, 2, 2, 599.98, '2025-06-15', 'completed');
    """)

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

create_demo_database()

import os
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_community.utilities import SQLDatabase
from langchain_community.agent_toolkits import SQLDatabaseToolkit, create_sql_agent

load_dotenv()

# Connect to the database
db = SQLDatabase.from_uri(
    "sqlite:///demo.db",
    include_tables=["customers", "products", "orders"],  # allowlist
    sample_rows_in_table_info=3,   # show sample rows in schema info
)

print("Database tables:", db.get_usable_table_names())
print("\nSchema info:")
print(db.get_table_info())

# The agent sees something like this for table info:
"""
CREATE TABLE customers (
    id INTEGER PRIMARY KEY,
    name TEXT NOT NULL,
    email TEXT UNIQUE,
    city TEXT,
    created_at DATE
)

/*
3 rows from customers table:
id  name            email               city        created_at
1   Alice Johnson   alice@example.com   New York    2025-01-15
2   Bob Smith       bob@example.com     Los Angeles 2025-02-20
3   Carol White     carol@example.com   Chicago     2025-03-10
*/
"""

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

# Create the toolkit (contains query, check, list tables tools)
toolkit = SQLDatabaseToolkit(db=db, llm=llm)

# Create the agent
agent_executor = create_sql_agent(
    llm=llm,
    toolkit=toolkit,
    agent_type="openai-tools",
    verbose=True,
    max_iterations=10,
    handle_parsing_errors=True,
)

# Test it
result = agent_executor.invoke({
    "input": "How many customers are from New York?"
})
print(result["output"])
# Expected: "There are 2 customers from New York: Alice Johnson and David Brown."

import re
from langchain_core.callbacks import BaseCallbackHandler

class SQLSecurityGuard(BaseCallbackHandler):
    """Intercepts SQL queries and blocks dangerous operations."""

    BLOCKED_KEYWORDS = {
        "DROP", "DELETE", "TRUNCATE", "ALTER", "CREATE",
        "INSERT", "UPDATE", "GRANT", "REVOKE", "EXEC"
    }

    def on_agent_action(self, action, **kwargs):
        if action.tool in ["sql_db_query", "sql_db_query_checker"]:
            query = action.tool_input
            if isinstance(query, dict):
                query = query.get("query", "")

            query_upper = query.upper()
            for keyword in self.BLOCKED_KEYWORDS:
                if re.search(r'\b' + keyword + r'\b', query_upper):
                    raise ValueError(
                        f"SECURITY BLOCK: Query contains prohibited keyword '{keyword}'. "
                        f"Only SELECT queries are allowed.\nQuery: {query}"
                    )

security_guard = SQLSecurityGuard()

# Recreate agent with security callback
agent_executor = create_sql_agent(
    llm=llm,
    toolkit=toolkit,
    agent_type="openai-tools",
    verbose=True,
    max_iterations=10,
    callbacks=[security_guard],
)

# This should work fine
safe_result = agent_executor.invoke({
    "input": "What is the total revenue from completed orders?"
})
print(safe_result["output"])

# This should be blocked
try:
    dangerous_result = agent_executor.invoke({
        "input": "Delete all orders from customer 1"
    })
except ValueError as e:
    print(f"BLOCKED: {e}")

from langchain_core.prompts import (
    ChatPromptTemplate,
    FewShotChatMessagePromptTemplate,
    MessagesPlaceholder,
)

# Define examples specific to your schema
examples = [
    {
        "input": "How many orders did each customer make?",
        "query": """SELECT c.name, COUNT(o.id) as order_count 
FROM customers c 
LEFT JOIN orders o ON c.id = o.customer_id 
GROUP BY c.id, c.name 
ORDER BY order_count DESC;"""
    },
    {
        "input": "What is the total revenue from each product category?",
        "query": """SELECT p.category, SUM(o.total_amount) as total_revenue
FROM orders o
JOIN products p ON o.product_id = p.id
WHERE o.status = 'completed'
GROUP BY p.category
ORDER BY total_revenue DESC;"""
    },
    {
        "input": "Which customers have spent more than $500 total?",
        "query": """SELECT c.name, c.email, SUM(o.total_amount) as total_spent
FROM customers c
JOIN orders o ON c.id = o.customer_id
WHERE o.status = 'completed'
GROUP BY c.id, c.name, c.email
HAVING SUM(o.total_amount) > 500
ORDER BY total_spent DESC;"""
    },
    {
        "input": "What are the top 3 best-selling products by quantity?",
        "query": """SELECT p.name, p.category, SUM(o.quantity) as total_sold
FROM products p
JOIN orders o ON p.id = o.product_id
GROUP BY p.id, p.name, p.category
ORDER BY total_sold DESC
LIMIT 3;"""
    },
    {
        "input": "Show me all pending orders with customer names",
        "query": """SELECT c.name, c.email, p.name as product, 
       o.quantity, o.total_amount, o.order_date
FROM orders o
JOIN customers c ON o.customer_id = c.id
JOIN products p ON o.product_id = p.id
WHERE o.status = 'pending'
ORDER BY o.order_date DESC;"""
    },
]

# Create few-shot prompt
example_prompt = ChatPromptTemplate.from_messages([
    ("human", "{input}"),
    ("ai", "{query}"),
])

few_shot_prompt = FewShotChatMessagePromptTemplate(
    example_prompt=example_prompt,
    examples=examples,
)

# Full system prompt with examples
full_prompt = ChatPromptTemplate.from_messages([
    ("system", f"""You are a SQL expert. Given a question, write a syntactically correct SQLite query.
    
Database schema:
{db.get_table_info()}

Rules:
- Only use tables: customers, products, orders
- Always use table aliases for readability
- Use proper JOINs when needed (never implicit joins)
- Add LIMIT 100 to queries that might return many rows
- Return only the SQL query, nothing else"""),
    few_shot_prompt,
    ("human", "{input}"),
])

# Test the prompt
from langchain_core.output_parsers import StrOutputParser

sql_generator = full_prompt | llm | StrOutputParser()

test_question = "Which city has the highest average order value?"
generated_sql = sql_generator.invoke({"input": test_question})
print(f"Question: {test_question}")
print(f"Generated SQL:\n{generated_sql}")

from langchain_community.agent_toolkits import create_sql_agent
from langchain_community.agent_toolkits.sql.prompt import SQL_PREFIX

custom_system_prompt = f"""You are a SQL expert for a demo e-commerce database.

Database tables: customers, products, orders
- customers: customer id, name, email, city, created_at
- products: product id, name, category, price, stock_quantity  
- orders: order id, customer_id (FK), product_id (FK), quantity, total_amount, order_date, status (completed/pending/shipped)

Key relationships:
- orders.customer_id → customers.id
- orders.product_id → products.id

When writing SQL:
1. Always use explicit JOINs, never implicit
2. Use aliases: c=customers, p=products, o=orders
3. For aggregations, always include GROUP BY correctly
4. Add LIMIT 100 for large result sets
5. Never run DELETE, UPDATE, DROP, or INSERT

Here are examples of correct SQL for this schema:
- "How many orders per customer?" → SELECT c.name, COUNT(o.id) FROM customers c LEFT JOIN orders o ON c.id=o.customer_id GROUP BY c.id, c.name
- "Total revenue by product?" → SELECT p.name, SUM(o.total_amount) FROM products p JOIN orders o ON p.id=o.product_id WHERE o.status='completed' GROUP BY p.id, p.name
"""

agent_executor = create_sql_agent(
    llm=llm,
    toolkit=toolkit,
    agent_type="openai-tools",
    verbose=True,
    max_iterations=15,
    prefix=custom_system_prompt,
)

# Test with progressively complex queries
test_questions = [
    "How much total revenue has been generated from completed orders?",
    "Which customer has placed the most orders?",
    "What is the average order value per city?",
    "Show me products that have never been ordered",
]

for question in test_questions:
    print(f"\n{'='*60}")
    print(f"Question: {question}")
    result = agent_executor.invoke({"input": question})
    print(f"Answer: {result['output']}")

# Pattern 1: Multi-level aggregation
result = agent_executor.invoke({
    "input": "For each customer, show their name, number of orders, total spent, and average order value, sorted by total spent descending"
})
print(result["output"])

# Pattern 2: Date-based filtering
result = agent_executor.invoke({
    "input": "How many orders were placed in the last 30 days? Count by status."
})
print(result["output"])

# Pattern 3: Finding records with no related records
result = agent_executor.invoke({
    "input": "Which products have never been ordered?"
})
print(result["output"])

# Pattern 4: Ranking within groups
result = agent_executor.invoke({
    "input": "For each city, show the customer who has spent the most"
})
print(result["output"])

# sql_api.py
from fastapi import FastAPI, HTTPException, Depends
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from pydantic import BaseModel
import logging
from datetime import datetime

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI(title="Text-to-SQL API", version="1.0")
security = HTTPBearer()

API_TOKEN = os.getenv("API_TOKEN", "demo-token-change-in-production")

def verify_token(credentials: HTTPAuthorizationCredentials = Depends(security)):
    if credentials.credentials != API_TOKEN:
        raise HTTPException(status_code=401, detail="Invalid token")
    return credentials.credentials

class QueryRequest(BaseModel):
    question: str
    max_results: int = 50

class QueryResponse(BaseModel):
    question: str
    answer: str
    sql_query: str | None
    execution_time_ms: float

@app.post("/query", response_model=QueryResponse)
async def query_database(
    request: QueryRequest,
    token: str = Depends(verify_token)
):
    if not request.question.strip():
        raise HTTPException(status_code=400, detail="Question cannot be empty")

    if len(request.question) > 500:
        raise HTTPException(status_code=400, detail="Question too long (max 500 chars)")

    start_time = datetime.now()
    logger.info(f"Query: {request.question}")

    try:
        result = agent_executor.invoke({"input": request.question})
        answer = result["output"]

        # Extract SQL from intermediate steps if available
        sql_used = None
        if "intermediate_steps" in result:
            for step in result["intermediate_steps"]:
                if hasattr(step[0], "tool") and "query" in step[0].tool:
                    sql_used = str(step[0].tool_input)
                    break

        execution_ms = (datetime.now() - start_time).total_seconds() * 1000

        logger.info(f"Answer generated in {execution_ms:.0f}ms")

        return QueryResponse(
            question=request.question,
            answer=answer,
            sql_query=sql_used,
            execution_time_ms=execution_ms,
        )

    except ValueError as e:
        # Security guard blocked a dangerous query
        logger.warning(f"Security block: {e}")
        raise HTTPException(status_code=403, detail=str(e))
    except Exception as e:
        logger.error(f"Query failed: {e}")
        raise HTTPException(status_code=500, detail="Failed to process query")

@app.get("/schema")
async def get_schema(token: str = Depends(verify_token)):
    return {
        "tables": db.get_usable_table_names(),
        "description": "E-commerce demo database with customers, products, and orders"
    }

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8001)