Building a SQL Query Agent with vLLM: Complete Guide

Building a SQL Query Agent with vLLM: Complete Guide

Text-to-SQL agents are one of the most practical applications of large language models in enterprise environments. Instead of forcing business users to learn SQL, these agents translate natural language questions into executable queries, run them against a database, and return human-readable answers. When you pair this capability with vLLM — a high-throughput, memory-efficient inference engine — you get a system that is fast, private, and cost-effective enough to run in production.

This tutorial walks you through building a complete SQL Query Agent from scratch using vLLM as the inference backend. We will cover the architecture, database schema extraction, prompt engineering, query execution with safety guards, and deployment best practices.

What Is a SQL Query Agent?

A SQL Query Agent is an LLM-powered system that accepts natural language input, generates a SQL query, executes it against a target database, and returns the result in a conversational format. Unlike a simple text-to-SQL converter, an agent can iterate: it inspects errors, refines queries, and explains results.

The core loop typically looks like this:

• Receive a user question in natural language
• Retrieve the relevant database schema (tables, columns, types, foreign keys)
• Prompt the LLM to generate a SQL query
• Validate and execute the query in a sandboxed environment
• If errors occur, feed them back to the LLM for correction
• Format the result set into a human-readable answer

Why vLLM?

vLLM is an open-source inference engine optimized for throughput and memory efficiency. It uses PagedAttention to manage the KV cache, allowing it to serve models much faster than naive HuggingFace Transformers. For a SQL agent, this matters for several reasons:

• Low latency: Users expect sub-second responses for interactive analytics.
• High concurrency: Multiple users can query simultaneously without queueing.
• Data privacy: You run the model on your own infrastructure, so sensitive schema and data never leave your network.
• Cost efficiency: No per-token API charges — you pay only for hardware.
• Model flexibility: You can use any open-weights model, including SQL-tuned ones like CodeLlama, DeepSeek-Coder, or Qwen2.5-Coder.

Architecture Overview

Our agent will consist of four components:

• vLLM Server: Serves the LLM through an OpenAI-compatible API.
• Schema Extractor: Introspects the database and builds a compact schema description.
• Agent Orchestrator: Manages the generate-execute-reflect loop.
• Query Executor: Safely runs SQL against the database with read-only protections.

Prerequisites and Installation

You will need a machine with a CUDA-capable GPU (an A10G or L4 is sufficient for 7B–14B models). Install the required Python packages:

pip install vllm sqlalchemy psycopg2-binary openai python-dotenv tabulate

For this tutorial we will use SQLite for simplicity, but the same code works with PostgreSQL, MySQL, or Snowflake by changing the connection string.

Starting the vLLM Server

Launch vLLM as a standalone server. We will use Qwen2.5-Coder-7B-Instruct, a model that performs exceptionally well on SQL generation tasks:

vllm serve Qwen/Qwen2.5-Coder-7B-Instruct \
  --port 8000 \
  --max-model-len 8192 \
  --gpu-memory-utilization 0.9 \
  --dtype auto

Once running, the server exposes an OpenAI-compatible endpoint at http://localhost:8000/v1. You can test it with a simple curl request:

curl http://localhost:8000/v1/chat/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "Qwen/Qwen2.5-Coder-7B-Instruct",
    "messages": [{"role": "user", "content": "Say hello"}]
  }'

Setting Up the Sample Database

Let us create a small e-commerce database to demonstrate the agent. Save this as setup_db.py and run it once:

import sqlite3

def create_database():
    conn = sqlite3.connect("ecommerce.db")
    cur = conn.cursor()

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

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

    CREATE TABLE IF NOT EXISTS orders (
        order_id INTEGER PRIMARY KEY,
        customer_id INTEGER,
        order_date TEXT,
        total_amount REAL,
        status TEXT,
        FOREIGN KEY (customer_id) REFERENCES customers(customer_id)
    );

    CREATE TABLE IF NOT EXISTS order_items (
        item_id INTEGER PRIMARY KEY,
        order_id INTEGER,
        product_id INTEGER,
        quantity INTEGER,
        unit_price REAL,
        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', '2024-01-15'),
        ('Bob Smith', 'bob@example.com', 'UK', '2024-02-20'),
        ('Carol Lee', 'carol@example.com', 'USA', '2024-03-10'),
        ('David Kim', 'david@example.com', 'Korea', '2024-04-05');

    INSERT INTO products (name, category, price, stock) VALUES
        ('Laptop', 'Electronics', 1200.00, 50),
        ('Mouse', 'Electronics', 25.00, 200),
        ('Keyboard', 'Electronics', 75.00, 150),
        ('Notebook', 'Stationery', 5.00, 500),
        ('Pen', 'Stationery', 2.00, 1000);

    INSERT INTO orders (customer_id, order_date, total_amount, status) VALUES
        (1, '2024-05-01', 1225.00, 'completed'),
        (2, '2024-05-03', 100.00, 'completed'),
        (1, '2024-05-10', 80.00, 'pending'),
        (3, '2024-05-15', 1275.00, 'completed'),
        (4, '2024-05-20', 7.00, 'cancelled');

    INSERT INTO order_items (order_id, product_id, quantity, unit_price) VALUES
        (1, 1, 1, 1200.00),
        (1, 2, 1, 25.00),
        (2, 3, 1, 75.00),
        (2, 2, 1, 25.00),
        (3, 3, 1, 75.00),
        (4, 1, 1, 1200.00),
        (4, 2, 3, 25.00),
        (5, 4, 1, 5.00),
        (5, 5, 1, 2.00);
    """)

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

if __name__ == "__main__":
    create_database()

Building the Schema Extractor

The agent needs to understand the database structure. Rather than dumping the entire schema, we extract a compact, LLM-friendly representation including table names, columns, types, and foreign key relationships.

import sqlite3
from typing import Dict, List

class SchemaExtractor:
    def __init__(self, db_path: str):
        self.db_path = db_path

    def get_schema(self) -> str:
        conn = sqlite3.connect(self.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()]

        schema_parts: List[str] = []
        for table in tables:
            cur.execute(f"PRAGMA table_info({table})")
            columns = cur.fetchall()

            col_lines = []
            for col in columns:
                col_id, col_name, col_type, not_null, default, pk = col
                constraints = []
                if pk:
                    constraints.append("PRIMARY KEY")
                if not_null:
                    constraints.append("NOT NULL")
                constraint_str = f" {', '.join(constraints)}" if constraints else ""
                col_lines.append(f"  {col_name} {col_type}{constraint_str}")

            cur.execute(f"PRAGMA foreign_key_list({table})")
            fks = cur.fetchall()
            for fk in fks:
                ref_table = fk[2]
                from_col = fk[3]
                to_col = fk[4]
                col_lines.append(f"  FOREIGN KEY ({from_col}) REFERENCES {ref_table}({to_col})")

            schema_parts.append(f"CREATE TABLE {table} (\n" + ",\n".join(col_lines) + "\n);")

        conn.close()
        return "\n\n".join(schema_parts)

    def get_sample_rows(self, table: str, limit: int = 3) -> str:
        conn = sqlite3.connect(self.db_path)
        cur = conn.cursor()
        cur.execute(f"SELECT * FROM {table} LIMIT {limit}")
        rows = cur.fetchall()
        col_names = [desc[0] for desc in cur.description]
        conn.close()

        lines = [", ".join(col_names)]
        for row in rows:
            lines.append(", ".join(str(v) for v in row))
        return "\n".join(lines)

Providing a few sample rows alongside the schema dramatically improves the LLM's ability to generate correct queries, because it sees real data patterns and value formats.

Building the Query Executor

Safety is critical. The executor must enforce read-only access and prevent destructive operations. Even in development, you should never let an LLM run arbitrary SQL without guards.

import sqlite3
import re
from typing import Tuple, Optional

class QueryExecutor:
    FORBIDDEN_KEYWORDS = [
        "DROP", "DELETE", "INSERT", "UPDATE", "ALTER",
        "CREATE", "TRUNCATE", "ATTACH", "DETACH", "PRAGMA"
    ]

    def __init__(self, db_path: str):
        self.db_path = db_path

    def validate_query(self, query: str) -> Tuple[bool, Optional[str]]:
        cleaned = re.sub(r"--.*$", "", query, flags=re.MULTILINE)
        cleaned = re.sub(r"/\*.*?\*/", "", cleaned, flags=re.DOTALL)
        upper = cleaned.upper().strip()

        if not upper.startswith("SELECT") and not upper.startswith("WITH"):
            return False, "Only SELECT or WITH queries are allowed."

        for keyword in self.FORBIDDEN_KEYWORDS:
            pattern = r'\b' + keyword + r'\b'
            if re.search(pattern, upper):
                return False, f"Forbidden keyword detected: {keyword}"

        if ";" in cleaned.rstrip().rstrip(";"):
            return False, "Multiple statements are not allowed."

        return True, None

    def execute(self, query: str, limit: int = 100) -> Tuple[bool, any, Optional[str]]:
        is_valid, error = self.validate_query(query)
        if not is_valid:
            return False, None, error

        try:
            conn = sqlite3.connect(self.db_path)
            conn.row_factory = sqlite3.Row
            cur = conn.cursor()

            if not query.upper().strip().startswith("WITH"):
                if "LIMIT" not in query.upper():
                    query = query.rstrip(";") + f" LIMIT {limit};"

            cur.execute(query)
            rows = cur.fetchall()
            columns = [desc[0] for desc in cur.description] if cur.description else []
            result = [dict(row) for row in rows]

            conn.close()
            return True, {"columns": columns, "rows": result}, None

        except Exception as e:
            return False, None, str(e)

Building the Agent Orchestrator

The orchestrator ties everything together. It constructs the prompt, calls vLLM, executes the query, and handles error correction. This is the heart of the system.

from openai import OpenAI
from tabulate import tabulate
import json

class SQLAgent:
    SYSTEM_PROMPT = """You are an expert SQL assistant. Given a database schema and a natural language question, you must:

1. Write a single, correct SQL query that answers the question.
2. Use only standard SQLite-compatible SQL.
3. Return ONLY the SQL query inside a sql code block. Do not include explanations before or after the code block.
4. If the question is ambiguous, make reasonable assumptions and proceed.
5. Always use table aliases for readability.
6. Never use SELECT * — always specify columns explicitly.

Database Schema:
{schema}

Sample data from key tables:
{samples}
"""

    CORRECTION_PROMPT = """The previous query failed with this error:

{error}

Original query:
{query}

Please fix the query and return only the corrected SQL inside a sql code block."""

    def __init__(self, db_path: str, vllm_url: str = "http://localhost:8000/v1", model: str = None):
        self.client = OpenAI(base_url=vllm_url, api_key="dummy")
        self.model = model or "Qwen/Qwen2.5-Coder-7B-Instruct"
        self.schema_extractor = SchemaExtractor(db_path)
        self.executor = QueryExecutor(db_path)
        self.max_retries = 3

    def _build_system_prompt(self) -> str:
        schema = self.schema_extractor.get_schema()
        samples = ""
        for table in ["customers", "products", "orders"]:
            samples += f"\n-- {table}:\n"
            samples += self.schema_extractor.get_sample_rows(table)
            samples += "\n"
        return self.SYSTEM_PROMPT.format(schema=schema, samples=samples)

    def _extract_sql(self, response: str) -> str:
        if "sql" in response:
            start = response.index("sql") + 6
            end = response.index("", start)
            return response[start:end].strip()
        elif "" in response:
            start = response.index("") + 3
            end = response.index("", start)
            return response[start:end].strip()
        return response.strip()

    def _generate_query(self, question: str, correction_context: str = None) -> str:
        messages = [
            {"role": "system", "content": self._build_system_prompt()},
            {"role": "user", "content": question}
        ]

        if correction_context:
            messages.append({"role": "assistant", "content": "sql\n" + correction_context["query"] + "\n"})
            messages.append({"role": "user", "content": self.CORRECTION_PROMPT.format(
                error=correction_context["error"],
                query=correction_context["query"]
            )})

        response = self.client.chat.completions.create(
            model=self.model,
            messages=messages,
            temperature=0.1,
            max_tokens=1024
        )

        return self._extract_sql(response.choices[0].message.content)

    def _format_results(self, result: dict, question: str) -> str:
        if not result["rows"]:
            return "The query returned no results."

        table_data = result["rows"]
        formatted_table = tabulate(
            table_data,
            headers="keys",
            tablefmt="grid",
            maxcolwidths=30
        )

        summary_prompt = f"""Based on these SQL query results, provide a concise natural language answer to the original question.

Original question: {question}

Results:
{formatted_table}

Answer the question directly and concisely. Mention specific numbers when relevant."""

        response = self.client.chat.completions.create(
            model=self.model,
            messages=[{"role": "user", "content": summary_prompt}],
            temperature=0.3,
            max_tokens=512
        )

        answer = response.choices[0].message.content
        return f"{answer}\n\n--- Query Results ---\n{formatted_table}"

    def ask(self, question: str) -> str:
        print(f"\nQuestion: {question}")

        correction_context = None
        for attempt in range(self.max_retries):
            query = self._generate_query(question, correction_context)
            print(f"Attempt {attempt + 1} - Generated query:\n{query}\n")

            success, result, error = self.executor.execute(query)

            if success:
                return self._format_results(result, question)
            else:
                print(f"Error: {error}")
                correction_context = {"query": query, "error": error}

        return f"Failed to generate a valid query after {self.max_retries} attempts. Last error: {error}"

Running the Agent

Now let us put it all together in a main script:

from sql_agent import SQLAgent

def main():
    agent = SQLAgent(
        db_path="ecommerce.db",
        vllm_url="http://localhost:8000/v1",
        model="Qwen/Qwen2.5-Coder-7B-Instruct"
    )

    questions = [
        "What are the top 3 customers by total order amount?",
        "Which product category has the highest total revenue?",
        "How many orders were cancelled and what was their total value?",
        "List all customers from the USA who have placed at least one completed order."
    ]

    for question in questions:
        answer = agent.ask(question)
        print(f"Answer: {answer}")
        print("=" * 80)

if __name__ == "__main__":
    main()

When you run this, you should see output like:

Question: What are the top 3 customers by total order amount?
Attempt 1 - Generated query:
SELECT c.name, SUM(o.total_amount) AS total_spent
FROM customers c
JOIN orders o ON c.customer_id = o.customer_id
WHERE o.status = 'completed'
GROUP BY c.customer_id, c.name
ORDER BY total_spent DESC
LIMIT 3;

Answer: The top 3 customers by total completed order amount are:
1. Alice Johnson with $1,225.00
2. Carol Lee with $1,275.00
3. Bob Smith with $100.00

Carol Lee is the highest spender with a total of $1,275.00 in completed orders.

Adding Few-Shot Examples

For complex schemas, few-shot examples significantly improve accuracy. You can embed example question-query pairs in the system prompt:

FEW_SHOT_EXAMPLES = """
Example Questions and Queries:

Q: How many products are in each category?
sql
SELECT category, COUNT(*) AS product_count
FROM products
GROUP BY category
ORDER BY product_count DESC;
Q: What is the average order value for completed orders?
sql
SELECT AVG(total_amount) AS avg_order_value
FROM orders
WHERE status = 'completed';
Q: Which customers have not placed any orders?
sql
SELECT c.name, c.email
FROM customers c
LEFT JOIN orders o ON c.customer_id = o.customer_id
WHERE o.order_id IS NULL;
"""

Append this to the system prompt after the schema. The LLM will learn your preferred query style and common patterns specific to your database.

Best Practices

1. Always Enforce Read-Only Access

Never trust LLM-generated SQL blindly. Use a dedicated read-only database user in production. In PostgreSQL, for example:

CREATE USER agent_reader WITH PASSWORD 'secure_password';
GRANT CONNECT ON DATABASE ecommerce TO agent_reader;
GRANT USAGE ON SCHEMA public TO agent_reader;
GRANT SELECT ON ALL TABLES IN SCHEMA public TO agent_reader;
ALTER DEFAULT PRIVILEGES IN SCHEMA public GRANT SELECT ON TABLES TO agent_reader;

2. Limit Result Sets

Always inject a LIMIT clause if one is not present. Returning millions of rows will consume excessive context window tokens and slow down the response. A reasonable default is 100 rows for display and 1000 for aggregation checks.

3. Use Low Temperature for Query Generation

Set temperature between 0.0 and 0.2 for SQL generation. You want deterministic, precise output. Use a slightly higher temperature (0.3–0.5) only for the natural language summary step where creativity is acceptable.

4. Cache the Schema

Extracting the schema on every request is wasteful. Cache it with a TTL or invalidate it only when the schema changes:

import time

class CachedSchemaExtractor(SchemaExtractor):
    def __init__(self, db_path: str, ttl_seconds: int = 3600):
        super().__init__(db_path)
        self.ttl = ttl_seconds
        self._cache = None
        self._cache_time = 0

    def get_schema(self) -> str:
        if self._cache is None or (time.time() - self._cache_time) > self.ttl:
            self._cache = super().get_schema()
            self._cache_time = time.time()
        return self._cache

5. Log Everything for Debugging

Log every generated query, execution result, and error. This helps you identify systematic failure patterns and build better few-shot examples over time:

import logging

logging.basicConfig(
    filename="sql_agent.log",
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s"
)

# Inside the ask method:
logging.info(f"Question: {question}")
logging.info(f"Query: {query}")
logging.info(f"Success: {success}, Error: {error}")

6. Choose the Right Model

Model selection has a huge impact on accuracy. For SQL tasks, these models perform well with vLLM:

• Qwen2.5-Coder-7B-Instruct: Excellent SQL generation, fits on a single 16GB GPU.
• DeepSeek-Coder-V2-Lite-Instruct: Strong reasoning, good for complex joins.
• CodeLlama-13B-Instruct: Solid baseline, widely tested.
• Llama-3.1-8B-Instruct: General-purpose but capable on SQL.

For production, benchmark several models on your specific schema using a held-out test set of question-query pairs.

7. Handle Ambiguous Questions Gracefully

Not every user question maps cleanly to a query. Add a clarification step where the agent can ask follow-up questions before generating SQL. You can implement this by adding a special "CLARIFY" output token that the orchestrator detects and routes back to the user.

Extending the Agent

Once the basic agent works, consider these enhancements:

• Schema linking: Use embeddings to retrieve only the relevant tables for each question, reducing prompt size on large databases.
• Query validation: Run EXPLAIN before execution to catch performance issues like full table scans.
• Multi-turn conversations: Maintain conversation history so users can ask follow-up questions like "What about only completed orders?"
• Visualization: Generate chart specifications (Chart.js, Plotly) alongside text summaries for numeric results.
• RAG integration: Store past successful question-query pairs in a vector database and retrieve them as dynamic few-shot examples.

Conclusion

Building a SQL Query Agent with vLLM gives you a powerful, private, and cost-effective way to let non-technical users interact with your databases. By combining vLLM's high-throughput inference with careful prompt engineering, robust query validation, and an iterative error-correction loop, you can create a production-grade system that handles real-world questions reliably. The key to success lies in giving the model rich context — schema, sample data, and few-shot examples — while maintaining strict safety guardrails around query execution. Start with the simple architecture described here, benchmark different models on your schema, and iteratively improve accuracy by logging failures and adding targeted examples. With vLLM handling inference efficiently, your agent can scale to serve entire organizations without the ongoing costs of external API calls.