← Back to DevBytes

Building a SQL Query Agent with AutoGen: Complete Guide

Introduction to Building a SQL Query Agent with AutoGen

Modern applications increasingly rely on natural language interfaces to databases. Instead of forcing users to learn SQL syntax, developers can build intelligent agents that translate plain English questions into valid SQL queries, execute them safely, and return human-readable insights. Microsoft's AutoGen framework makes this remarkably achievable by orchestrating multiple conversational agents that collaborate to solve complex tasks.

In this complete guide, you'll learn how to build a production-ready SQL Query Agent using AutoGen. We'll cover the architecture, walk through a full implementation, and discuss best practices for security, reliability, and performance.

What Is a SQL Query Agent?

A SQL Query Agent is an AI-powered assistant that bridges the gap between non-technical users and relational databases. It performs three core functions:

With AutoGen, we can decompose these responsibilities across specialized agents that converse with each other. One agent writes SQL, another executes it, and a third explains the results. This multi-agent approach produces more accurate and robust behavior than a single monolithic prompt.

Why Use AutoGen for SQL Agents?

AutoGen offers several advantages over building a SQL agent from scratch:

Prerequisites and Setup

Before we start coding, ensure you have the following:

Install the required packages:

pip install "autogen-agentchat==0.4.7" "autogen-ext[openai]" python-dotenv sqlalchemy

Create a .env file to store your API key securely:

OPENAI_API_KEY=sk-your-key-here

Building the Sample Database

For this tutorial, we'll create a small e-commerce database with customers, orders, and products. This gives the agent enough complexity to demonstrate joins, aggregations, and filtering.

import sqlite3

def create_sample_db(db_path: str = "ecommerce.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;

    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 DEFAULT 0
    );

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

    cur.executemany(
        "INSERT INTO customers VALUES (?, ?, ?, ?, ?)",
        [
            (1, "Alice Johnson", "alice@example.com", "USA", "2024-01-15"),
            (2, "Bob Smith", "bob@example.com", "UK", "2024-02-20"),
            (3, "Carol Lee", "carol@example.com", "USA", "2024-03-10"),
            (4, "David Kim", "david@example.com", "Canada", "2024-04-05"),
        ],
    )

    cur.executemany(
        "INSERT INTO products VALUES (?, ?, ?, ?, ?)",
        [
            (1, "Laptop", "Electronics", 1200.00, 15),
            (2, "Mouse", "Electronics", 25.00, 100),
            (3, "Keyboard", "Electronics", 75.00, 50),
            (4, "Notebook", "Stationery", 5.00, 200),
            (5, "Desk Lamp", "Furniture", 45.00, 30),
        ],
    )

    cur.executemany(
        "INSERT INTO orders VALUES (?, ?, ?, ?, ?)",
        [
            (1, 1, 1, 1, "2024-05-01"),
            (2, 1, 2, 2, "2024-05-03"),
            (3, 2, 3, 1, "2024-05-10"),
            (4, 3, 1, 1, "2024-05-15"),
            (5, 3, 4, 5, "2024-05-18"),
            (6, 4, 5, 2, "2024-05-20"),
            (7, 2, 2, 3, "2024-05-25"),
        ],
    )

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

if __name__ == "__main__":
    create_sample_db()

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

Designing the Agent Architecture

Our SQL Query Agent system will consist of three cooperating agents:

We'll also use a UserProxyAgent to represent the human user and manage the conversation flow.

Implementing the SQL Execution Tool

The executor is the most security-sensitive component. It must validate queries, prevent destructive operations, and return results in a format the agents can understand.

import sqlite3
import re
from typing import Annotated

DB_PATH = "ecommerce.db"

# Keywords that indicate potentially destructive operations
FORBIDDEN_KEYWORDS = re.compile(
    r"\b(DROP|DELETE|INSERT|UPDATE|ALTER|TRUNCATE|CREATE|GRANT|REVOKE)\b",
    re.IGNORECASE,
)

def execute_sql_query(
    query: Annotated[str, "A valid SQL SELECT query to execute"],
) -> str:
    """Execute a read-only SQL query against the e-commerce database.

    Only SELECT statements are allowed. Any attempt to modify data
    or schema will be rejected.
    """
    query_stripped = query.strip().rstrip(";")

    # Ensure the query starts with SELECT
    if not query_stripped.upper().startswith("SELECT"):
        return "ERROR: Only SELECT queries are permitted."

    # Block destructive keywords anywhere in the query
    if FORBIDDEN_KEYWORDS.search(query_stripped):
        return "ERROR: Destructive SQL operations are not allowed."

    try:
        conn = sqlite3.connect(DB_PATH)
        conn.row_factory = sqlite3.Row
        cur = conn.cursor()
        cur.execute(query_stripped)
        rows = cur.fetchall()

        if not rows:
            return "Query executed successfully but returned no rows."

        columns = [desc[0] for desc in cur.description]
        result_lines = [" | ".join(columns)]
        result_lines.append("-" * len(result_lines[0]))
        for row in rows:
            result_lines.append(" | ".join(str(row[col]) for col in columns))

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

    except sqlite3.Error as e:
        return f"SQL Error: {e}"
    except Exception as e:
        return f"Unexpected error: {e}"

Key safety measures in this tool include a regex check for forbidden keywords, a requirement that queries begin with SELECT, and comprehensive error handling that returns messages the agent can use to self-correct.

Building the Agent System

Now we assemble the agents using AutoGen's AssistantAgent and UserProxyAgent classes. The SQL Coder agent will have access to our execute_sql_query tool.

import os
import asyncio
from dotenv import load_dotenv

from autogen_agentchat.agents import AssistantAgent, UserProxyAgent
from autogen_agentchat.teams import RoundRobinGroupChat
from autogen_agentchat.conditions import TextMentionTermination, MaxMessageTermination
from autogen_ext.models.openai import OpenAIChatCompletionClient

load_dotenv()

SCHEMA_DESCRIPTION = """
Database: ecommerce.db (SQLite)

Tables:

1. customers
   - customer_id INTEGER PRIMARY KEY
   - name TEXT NOT NULL
   - email TEXT UNIQUE NOT NULL
   - country TEXT NOT NULL
   - signup_date TEXT (ISO format YYYY-MM-DD)

2. products
   - product_id INTEGER PRIMARY KEY
   - name TEXT NOT NULL
   - category TEXT NOT NULL
   - price REAL NOT NULL
   - stock INTEGER NOT NULL DEFAULT 0

3. orders
   - order_id INTEGER PRIMARY KEY
   - customer_id INTEGER (FK -> customers.customer_id)
   - product_id INTEGER (FK -> products.product_id)
   - quantity INTEGER NOT NULL
   - order_date TEXT (ISO format YYYY-MM-DD)

Relationships:
- orders.customer_id references customers.customer_id
- orders.product_id references products.product_id
"""

SQL_CODER_SYSTEM_PROMPT = f"""You are an expert SQL analyst. Your job is to answer
the user's question by writing and executing SQL queries against an e-commerce database.

Here is the database schema:

{SCHEMA_DESCRIPTION}

Rules:
1. Write only SELECT queries. Never attempt INSERT, UPDATE, DELETE, or DROP.
2. Use the execute_sql_query tool to run your query.
3. If the tool returns an error, analyze it and rewrite the query.
4. After receiving results, provide a clear, concise natural language summary.
5. If the question is ambiguous, make a reasonable assumption and state it.
6. Always end your final answer with the word TERMINATE on its own line.
"""


async def main() -> None:
    model_client = OpenAIChatCompletionClient(
        model="gpt-4o-mini",
        api_key=os.getenv("OPENAI_API_KEY"),
        temperature=0.1,  # Low temperature for deterministic SQL generation
    )

    # The SQL Coder agent has access to the execution tool
    sql_coder = AssistantAgent(
        name="SQLCoder",
        model_client=model_client,
        system_message=SQL_CODER_SYSTEM_PROMPT,
        tools=[execute_sql_query],
        reflect_on_tool_use=True,  # Agent writes a summary after tool results
    )

    # The user proxy represents the human asking questions
    user_proxy = UserProxyAgent(
        name="User",
        description="A human user asking questions about the e-commerce database.",
    )

    # Termination: stop when agent says TERMINATE or after 10 messages
    termination = TextMentionTermination("TERMINATE") | MaxMessageTermination(10)

    # Round-robin team alternates between User and SQLCoder
    team = RoundRobinGroupChat(
        participants=[user_proxy, sql_coder],
        termination_condition=termination,
    )

    # Run a sample query
    task = "Which product category has generated the most revenue, and who are the top 3 customers by total spending?"

    print("=" * 70)
    print(f"User Question: {task}")
    print("=" * 70)

    result = await team.run(task=task)

    print("\n" + "=" * 70)
    print("Conversation Complete")
    print("=" * 70)
    for msg in result.messages:
        print(f"\n[{msg.source}]:")
        print(msg.content if isinstance(msg.content, str) else str(msg.content))

    await model_client.close()


if __name__ == "__main__":
    asyncio.run(main())

How the Conversation Flows

When you run this script, the following sequence occurs:

Adding a Query Reviewer Agent

For more complex scenarios, adding a second agent that reviews SQL before execution dramatically improves accuracy. The reviewer checks for correctness, efficiency, and safety.

REVIEWER_SYSTEM_PROMPT = f"""You are a senior database engineer who reviews SQL queries
written by a junior analyst before they are executed.

Database schema:

{SCHEMA_DESCRIPTION}

Your responsibilities:
1. Verify the query references valid tables and columns.
2. Check that JOINs use correct foreign key relationships.
3. Ensure aggregations include proper GROUP BY clauses.
4. Flag any syntax errors or logic mistakes.
5. If the query is correct, respond with "APPROVED" followed by a brief explanation.
6. If the query needs changes, provide the corrected query and explain what was wrong.
7. Always end with TERMINATE on its own line.
"""


async def build_reviewed_team():
    model_client = OpenAIChatCompletionClient(
        model="gpt-4o-mini",
        api_key=os.getenv("OPENAI_API_KEY"),
        temperature=0.1,
    )

    sql_coder = AssistantAgent(
        name="SQLCoder",
        model_client=model_client,
        system_message=SQL_CODER_SYSTEM_PROMPT,
        tools=[execute_sql_query],
        reflect_on_tool_use=True,
    )

    reviewer = AssistantAgent(
        name="Reviewer",
        model_client=model_client,
        system_message=REVIEWER_SYSTEM_PROMPT,
    )

    user_proxy = UserProxyAgent(name="User")

    termination = TextMentionTermination("TERMINATE") | MaxMessageTermination(15)

    team = RoundRobinGroupChat(
        participants=[user_proxy, sql_coder, reviewer],
        termination_condition=termination,
    )

    return team, model_client

In this three-agent setup, the SQLCoder drafts a query, the Reviewer validates or corrects it, and the SQLCoder executes the approved version. This mirrors how real database teams operate and catches errors that a single agent might miss.

Best Practices

Security First

Database access by AI agents carries real risk. Follow these principles:

Prompt Engineering for SQL

Handling Edge Cases

Performance Optimization

Extending the Agent

Once you have the basic system working, consider these enhancements:

Here's a quick example of a schema introspection tool:

def get_table_schema(
    table_name: Annotated[str, "The name of the table to inspect"],
) -> str:
    """Return column information for a specific table."""
    try:
        conn = sqlite3.connect(DB_PATH)
        cur = conn.cursor()
        cur.execute(f"PRAGMA table_info({table_name})")
        columns = cur.fetchall()
        conn.close()

        if not columns:
            return f"Table '{table_name}' does not exist."

        lines = [f"Table: {table_name}"]
        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}")
            lines.append(f"  - {name} {col_type} {' '.join(constraints)}".strip())

        return "\n".join(lines)
    except sqlite3.Error as e:
        return f"Error: {e}"

Conclusion

Building a SQL Query Agent with AutoGen combines the power of large language models with the structure of multi-agent orchestration to create a robust natural language database interface. By separating concerns into specialized agents, enforcing strict read-only safety boundaries in the execution tool, and carefully crafting system prompts with complete schema context, you can deliver an agent that handles real-world questions reliably. Start with the single-agent implementation, add a reviewer as your query complexity grows, and continuously refine your prompts based on the questions your users actually ask. With the foundation in this guide, you're well equipped to extend the system with dynamic schema discovery, visualization, and multi-database support as your needs evolve.

— Ad —

Google AdSense will appear here after approval

← Back to all articles