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:
- Schema understanding: Reads and interprets the database structure, including tables, columns, relationships, and constraints.
- Query generation: Translates natural language questions into syntactically correct SQL statements.
- Result interpretation: Executes the query, then summarizes the returned data in plain language.
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:
- Multi-agent orchestration: Define distinct roles (planner, coder, executor, reviewer) and let them collaborate through structured conversations.
- Built-in conversation management: AutoGen handles message routing, termination conditions, and conversation history automatically.
- Tool integration: Easily register Python functions as tools that agents can call, such as a function that runs SQL against your database.
- Human-in-the-loop support: Add a UserProxyAgent to approve queries before execution, which is critical for database safety.
- Framework agnostic: Works with OpenAI, Azure OpenAI, Anthropic, and local models through LiteLLM.
Prerequisites and Setup
Before we start coding, ensure you have the following:
- Python 3.10 or higher
- An OpenAI API key (or another supported LLM provider)
- A SQLite database for testing (we'll create one inline)
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:
- Schema Agent: Provides the database schema and answers questions about table structures.
- SQL Coder Agent: Writes SQL queries based on the user's question and the schema.
- SQL Executor (Tool): A Python function registered as a tool that runs the query and returns results.
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:
- The
UserProxyAgentsubmits the natural language question to the team. - The
SQLCoderagent receives the question, analyzes the schema from its system prompt, and generates a SQL query. - The agent calls
execute_sql_querywith the generated SQL. AutoGen intercepts this tool call, executes the function, and feeds the result back to the agent. - Because
reflect_on_tool_useisTrue, the agent produces a natural language summary of the results rather than just returning raw data. - The agent appends
TERMINATEto signal the conversation is complete.
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:
- Use read-only database credentials. If possible, connect with a user that has only SELECT privileges at the database level, not just in your Python code.
- Never expose connection strings in prompts. Keep database configuration in environment variables.
- Implement query timeouts. A runaway query can lock tables. Set a timeout in your executor:
cur.execute(query)wrapped in asignal.alarmor using SQLite'sprogress_handler. - Log every executed query. Maintain an audit trail for compliance and debugging.
Prompt Engineering for SQL
- Include the full schema in the system prompt. The agent cannot query what it does not know exists. Include column types, foreign keys, and constraints.
- Provide example queries. Few-shot examples of common query patterns dramatically improve output quality, especially for complex joins.
- Specify the SQL dialect. SQLite, PostgreSQL, and MySQL have different functions and syntax. State which dialect the agent should use.
- Use low temperature. SQL is unforgiving of creative variation. A temperature of 0.0 to 0.2 produces more reliable queries.
Handling Edge Cases
- Empty results: Train the agent to distinguish between "no data exists" and "the query was wrong." Encourage it to verify assumptions.
- Ambiguous questions: Instruct the agent to state assumptions rather than guessing silently.
- Large result sets: Add a LIMIT clause recommendation in the system prompt to prevent overwhelming the context window.
- Self-correction: The agent should treat error messages from the executor as feedback and retry with a corrected query rather than giving up.
Performance Optimization
- Cache schema lookups. If your schema is large, consider a retrieval step that fetches only relevant table definitions rather than embedding the entire schema in the prompt.
- Use structured outputs. For programmatic integration, configure the model client to return JSON with fields like
sql_query,explanation, andsummary. - Batch related questions. If a user asks multiple questions, let the agent plan all queries upfront and execute them in sequence to reduce round trips.
Extending the Agent
Once you have the basic system working, consider these enhancements:
- Schema introspection tool: Instead of hardcoding the schema, add a tool that runs
PRAGMA table_infoqueries so the agent discovers the schema dynamically. This makes the agent portable across databases. - Visualization tool: Register a tool that generates matplotlib charts from query results. The agent can decide when a chart is more useful than a text summary.
- Multi-database support: Add a tool parameter for database selection, letting the agent query across multiple data sources.
- Conversation memory: Use AutoGen's
ChatAgentwith a memory store so the agent remembers previous questions and can answer follow-ups like "What about last month?"
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.