Build a SQL Agent with Groq

Build a self-correcting SQL agent that answers questions about a PostgreSQL database.

Introduction

30 Second Summary

Every organization sits on a database full of answers. But getting those answers usually means writing SQL by hand, which is slow and error-prone when the schema has dozens of tables.

In this project, you will build a command-line SQL agent that translates natural language questions into live PostgreSQL queries. You type a question, and the agent figures out which tables to inspect, writes the SQL, runs it, and hands back a plain-English answer.

What You'll Build

Type a question like "What are the top 5 nations by total revenue?" and watch your agent autonomously discover the database schema, generate a multi-table JOIN query, execute it against 60 million rows, and respond with a clear summary of the results.

By the end of this project, you'll have:

  • A working text-to-SQL agent powered by Groq's tool-calling API that answers natural language questions against a 60M+ row database.
  • A transparent reasoning trace that prints every tool call, generated query, and error recovery step to your terminal in real time.
  • A safety and self-correction layer that blocks destructive SQL before it reaches the database and automatically retries when the agent writes a bad query.
  • Secret Mission: Add query performance analysis with EXPLAIN ANALYZE so your agent can explain why a query is fast or slow.

Are there any prerequisites?

This project assumes you have Python and SQL experience. You'll also need a running PostgreSQL 17 instance with TPC-H benchmark data loaded and a Groq API key.

Before We Start

Before we dive in, let's lock in what you're about to build and why it matters.

Set Up Your Environment

Your SQL agent needs two connections to function. It calls the Groq API to send questions to the LLM, and it uses psycopg to execute queries against your PostgreSQL database.

In this step, you'll install both Python packages, configure your API credentials, and verify each connection independently. Catching a broken connection now saves you from confusing errors later.

In this step, get ready to:
  • Install the Groq Python SDK and psycopg.
  • Set the GROQ_API_KEY environment variable and verify a chat completion.
  • Test the PostgreSQL connection with a query against the TPC-H database.
Install the Groq Python SDK and psycopg

You need two Python packages: the Groq SDK for calling the LLM and psycopg for connecting to PostgreSQL.

  • Install both packages by running this command:
pip install groq "psycopg[binary]"

What are these packages?

groq is the official Python SDK for the Groq API. It handles authentication, request formatting, and response parsing for LLM chat completions.

psycopg[binary] installs Psycopg 3 with pre-compiled C extensions. The binary extras mean you don't need PostgreSQL development headers installed on your machine.

  • Confirm the Groq SDK installed correctly by running this import check:
python -c "from groq import Groq; print('Groq SDK ready')"
  • Confirm psycopg installed correctly by running:
python -c "import psycopg; print('psycopg ready')"

You should see Groq SDK ready and psycopg ready printed in your terminal.

Not seeing 'ready' messages?

  • Make sure pip is using the same Python installation you're running. Try python -m pip install groq "psycopg[binary]" to force the correct one.
  • If you see a ModuleNotFoundError, the install may have gone to a different Python version. Re-run the install with the exact Python binary you're using.
  • Help me troubleshoot my pip install for groq and psycopg.
Set your Groq API key and test a chat completion

The Groq SDK reads an environment variable called GROQ_API_KEY to authenticate your API requests. You already have a Groq API key from the prerequisites.

  • Set the environment variable by running this command, replacing your-api-key-here with your actual Groq API key:
export GROQ_API_KEY=your-api-key-here

Why an environment variable?

Storing your API key in an environment variable keeps it out of your source code. If you hardcoded the key into a Python file, anyone with access to the file could see your credentials.

  • Verify your API key works by running a test chat completion:
python -c "
from groq import Groq
import os

client = Groq(api_key=os.environ.get('GROQ_API_KEY'))
response = client.chat.completions.create(
    model='qwen/qwen3.6-27b',
    messages=[{'role': 'user', 'content': 'Say hello in one sentence.'}]
)
print(response.choices[0].message.content)
"

What does this test do?

This creates a Groq client using your API key, sends a simple prompt to the qwen/qwen3.6-27b model, and prints the model's response. If you see a greeting, your Groq connection is working.

You should see a short greeting from the model printed in your terminal.

Getting an authentication error?

  • Double-check that you ran the export command in the same terminal session where you're running the test. Opening a new terminal tab resets environment variables.
  • Confirm your key has no leading or trailing spaces. Copy it fresh from the Groq console and re-run the export command.
  • My Groq API test is returning an authentication error.
Test your PostgreSQL connection

Your PostgreSQL database is already running in Docker with the TPC-H benchmark data loaded. Let's verify that psycopg can connect and query it.

Before you run the next command, take a guess: how many customers do you think are stored in a benchmark dataset at scale factor 10?

  • Test the connection by running a count query against the customer table, replacing your-password-here with your PostgreSQL password:
python -c "
import psycopg

conn = psycopg.connect(
    dbname='tpch',
    user='postgres',
    password='your-password-here',
    host='localhost',
    port=5432
)
cur = conn.cursor()
cur.execute('SELECT COUNT(*) FROM customer')
print(f'Customer count: {cur.fetchone()[0]}')
cur.close()
conn.close()
"

You should see Customer count: 1500000. That's 1.5 million customers across the TPC-H benchmark dataset at scale factor 10.

Why 1.5 million?

TPC-H is a standard benchmark for analytical databases. The "scale factor" multiplies every table's row count. At SF10, the customer table has 1,500,000 rows, and the lineitem table has roughly 60 million. These volumes make the row-limiting and column-selection rules you'll add to the agent genuinely important.

Connection refused?

  • Make sure your PostgreSQL Docker container is running. Check with docker ps and look for a container exposing port 5432.
  • Verify that the dbname, user, and password values match how you configured your Docker container.
  • My PostgreSQL connection test is failing with a connection error.

Both connections are live. In the next step, you'll build the database tools that let the LLM discover and query your PostgreSQL schema.

Build the Database Tools

Your LLM cannot reach into your PostgreSQL database on its own. It has no connection string, no SQL knowledge, and no way to run a query. It needs Python functions that act as its hands.

In this step, you will build three database tools in a file called db_tools.py. These functions let the agent discover which tables exist, inspect a table's columns and foreign key relationships, and execute SQL queries. You will test each tool standalone before handing them to the LLM in the next step.

In this step, get ready to:
  • Create list_tables and get_table_schema functions for database discovery.
  • Add a run_query function with error handling and a 50-row limit.
  • Test all three functions against your TPC-H database.
Create db_tools.py with list_tables and get_table_schema

The first thing the agent needs to do when it receives a question is figure out which tables exist. Then it needs to inspect those tables to learn about their columns and how they relate to each other through foreign keys.

  • In VS Code, create a new file called db_tools.py.
  • Add the database connection setup by pasting this code:
import psycopg
import json

DB_CONFIG = {
    "dbname": "tpch",
    "user": "postgres",
    "password": "your-password-here",
    "host": "localhost",
    "port": 5432,
}


def get_connection():
    return psycopg.connect(**DB_CONFIG)

What does this code do?

DB_CONFIG stores the connection details for your TPC-H database. get_connection() creates a fresh psycopg connection each time it is called. Every tool function will use this helper so the connection details live in one place.

  • Replace your-password-here in DB_CONFIG with the password you used when you set up your PostgreSQL Docker container.
  • Below get_connection, add the list_tables function by pasting this code:
def list_tables():
    """List all tables in the public schema."""
    with get_connection() as conn:
        with conn.cursor() as cur:
            cur.execute(
                """
                SELECT table_name
                FROM information_schema.tables
                WHERE table_schema = 'public'
                ORDER BY table_name
                """
            )
            tables = [row[0] for row in cur.fetchall()]
    return json.dumps({"tables": tables})

What does this code do?

list_tables() queries PostgreSQL's information_schema.tables to discover every table in the public schema. It returns the list as a JSON string because the LLM will eventually read this output as text inside a chat message.

The with get_connection() as conn pattern is a psycopg 3 context manager. It opens the connection, commits any transaction when the block exits, and closes the connection automatically.

  • Save db_tools.py.
  • Test list_tables by running this command in your terminal:
python -c "from db_tools import list_tables; print(list_tables())"

You should see a JSON string listing the 8 TPC-H tables: customer, lineitem, nation, orders, part, partsupp, region, and supplier.

Not seeing 8 tables?

  • Make sure your PostgreSQL Docker container is running. Check with docker ps in your terminal.
  • Double check that the password in DB_CONFIG matches the one you used when setting up the container.
  • Verify that you are running the command from the same directory as db_tools.py.

Still stuck? Help me debug my list_tables function. I'm getting an error when I try to import and run it.

Now the agent can see which tables exist. Next it needs to inspect a specific table's columns and foreign keys so it can write correct JOINs.

  • Below list_tables, start the get_table_schema function by pasting this code:
def get_table_schema(table_name):
    """Get column names, data types, and foreign key relationships for a table."""
    with get_connection() as conn:
        with conn.cursor() as cur:
            cur.execute(
                """
                SELECT column_name, data_type, is_nullable
                FROM information_schema.columns
                WHERE table_schema = 'public' AND table_name = %s
                ORDER BY ordinal_position
                """,
                (table_name,),
            )
            columns = [
                {"name": row[0], "type": row[1], "nullable": row[2]}
                for row in cur.fetchall()
            ]

What does this first query do?

This queries information_schema.columns to get every column in the specified table. For each column it captures the name, data type, and whether it allows nulls. The %s placeholder is a parameterized query. Psycopg sends the table name separately from the SQL, which prevents SQL injection.

  • Directly below the columns list comprehension (still inside the with conn.cursor() block), paste the foreign key query and the return statement:
            cur.execute(
                """
                SELECT
                    kcu.column_name,
                    ccu.table_name AS foreign_table,
                    ccu.column_name AS foreign_column
                FROM information_schema.table_constraints tc
                JOIN information_schema.key_column_usage kcu
                    ON tc.constraint_name = kcu.constraint_name
                    AND tc.table_schema = kcu.table_schema
                JOIN information_schema.constraint_column_usage ccu
                    ON tc.constraint_name = ccu.constraint_name
                    AND tc.table_schema = ccu.table_schema
                WHERE tc.constraint_type = 'FOREIGN KEY'
                    AND tc.table_schema = 'public'
                    AND tc.table_name = %s
                """,
                (table_name,),
            )
            foreign_keys = [
                {"column": row[0], "references": f"{row[1]}.{row[2]}"}
                for row in cur.fetchall()
            ]

    return json.dumps(
        {"table": table_name, "columns": columns, "foreign_keys": foreign_keys}
    )

What does the foreign key query do?

This joins three information_schema views to discover which columns in the current table reference columns in other tables. For example, the customer table's c_nationkey column references nation.n_nationkey. This is the information the LLM needs to write correct JOIN conditions.

The function returns everything as a single JSON object containing the table name, its columns, and its foreign key relationships.

  • Save db_tools.py.
  • Test get_table_schema by running:
python -c "from db_tools import get_table_schema; print(get_table_schema('customer'))"

You should see a JSON object with the customer table's columns (like c_custkey, c_name, c_nationkey) and a foreign_keys array showing that c_nationkey references nation.n_nationkey.

Seeing an empty columns list?

  • Make sure you pasted both code blocks inside the same function. The return json.dumps(...) line should be indented at the same level as the with get_connection() line.
  • Check that the foreign key query block is indented inside the with conn.cursor() as cur: block, at the same level as the first cur.execute call.

Need help? My get_table_schema function returns an empty columns list or throws an indentation error.

Add a run_query function with error handling

The agent can now discover tables and schemas. The last tool it needs is the ability to actually execute a SQL query and read the results.

  • Below get_table_schema in db_tools.py, paste the run_query function:
def run_query(sql):
    """Execute a read-only SQL query and return results. Limited to 50 rows."""
    try:
        with get_connection() as conn:
            with conn.cursor() as cur:
                cur.execute(sql)
                if cur.description is None:
                    return json.dumps(
                        {
                            "error": "Query did not return results. Only SELECT queries are supported."
                        }
                    )
                col_names = [desc[0] for desc in cur.description]
                rows = cur.fetchmany(50)
                result = {
                    "columns": col_names,
                    "rows": [list(row) for row in rows],
                    "row_count": len(rows),
                    "note": "Results limited to 50 rows"
                    if len(rows) == 50
                    else None,
                }
    except Exception as e:
        result = {"error": str(e)}

    return json.dumps(result, default=str)

What does this code do?

run_query() executes whatever SQL string the LLM generates. A few safety features are built in:

  • fetchmany(50) caps the result at 50 rows. The TPC-H lineitem table has around 60 million rows at SF10, so an uncapped SELECT * would overwhelm memory.
  • cur.description is None when the query does not return rows (like an INSERT). The function catches this and returns an error message.
  • The except Exception block captures PostgreSQL errors (bad column names, syntax mistakes) and returns them as JSON. The LLM will read these error messages and use them to fix its query.
  • Save db_tools.py.
  • Test run_query by running a count query against the customer table:
python -c "from db_tools import run_query; print(run_query('SELECT COUNT(*) FROM customer'))"

You should see a JSON result with row_count of 1 and the value 1500000. That is 1.5 million customers in the TPC-H SF10 dataset.

Getting a connection error?

  • Confirm your PostgreSQL container is still running with docker ps.
  • Check that DB_CONFIG in db_tools.py has the correct dbname, user, password, host, and port values.

Need help? My run_query function is throwing a connection error when I try to query the TPC-H customer table.

Test all three functions together

Each function works individually. Before you move on, run one final check that confirms all three tools are importable and returning real data from your database.

  • Before you run this, think about what you expect to see. How many tables should list_tables return? What columns should the customer schema include?
  • Run this combined test:
python -c "from db_tools import list_tables, get_table_schema; print(list_tables()); print(get_table_schema('customer'))"

You should see the 8 TPC-H tables listed first, then the customer table's columns and foreign keys printed on the next line.

What should I see?

The first line of output is a JSON object with a tables array containing 8 names: customer, lineitem, nation, orders, part, partsupp, region, supplier.

The second line is a JSON object for the customer table showing columns like c_custkey, c_name, c_address, c_nationkey, and others. The foreign_keys array shows that c_nationkey references nation.n_nationkey.

✔️ Awesome, I've got everything!

Great. Double check you have saved db_tools.py and you are ready for the next step.

ⓧ I'd like to double check the full code

Scroll up through your db_tools.py and confirm these sections are present, in order:

  • Imports (psycopg and json) and DB_CONFIG dict.
  • get_connection() function.
  • list_tables() function querying information_schema.tables.
  • get_table_schema(table_name) function with two cur.execute calls (columns query and foreign keys query) and a return json.dumps(...) at the end.
  • run_query(sql) function with try/except error handling and fetchmany(50).

If any section is missing, scroll back to the substep where it was introduced and paste it in.

Your database tools are tested and ready. In the next step, you will wire them up to the Groq LLM and watch the agent call them autonomously to answer questions.

Build Your First Agent

Your database tools are tested and working. Each one can list tables, inspect schemas, or run queries on its own. The missing piece is something that decides which tool to call, and when.

That decision-maker is the LLM. In this step, you will wire your three functions into Groq's tool-calling API so the model can autonomously choose tools, read results, and keep going until it has a final answer.

In this step, get ready to:
  • Create agent.py with a system prompt and tool definitions that tell the LLM what it can do.
  • Build an agent loop that sends messages to Groq, processes tool calls, and continues until the model returns a final answer.
  • Test the agent with a simple question and a complex multi-table question.
Create the agent module with a system prompt and tool definitions

The agent needs three things before it can do anything: a connection to the Groq API, a system prompt telling it what role to play, and a list of tool definitions describing the functions it can call.

  • Create a new file called agent.py in the same folder as your db_tools.py file.
  • Set up the Groq client and system prompt by pasting this code into agent.py:
import os
import json
from groq import Groq
from db_tools import list_tables, get_table_schema, run_query

# Initialize the Groq client using the API key from your environment
client = Groq(api_key=os.environ.get("GROQ_API_KEY"))

# A minimal system prompt telling the model its role
SYSTEM_PROMPT = """You are a SQL agent that answers questions about a PostgreSQL database.
The database contains TPC-H benchmark data.

Use the available tools to discover the schema and run queries to answer the user's question.
Format your final answer as a clear, human-readable response."""

What does this code do?

The imports bring in your three database functions from db_tools.py so the agent can call them. The Groq client reads your API key from the environment variable you set up earlier.

The system prompt is intentionally minimal right now. It tells the model it is a SQL agent and to use the available tools, but gives no specific instructions about how to approach multi-table queries.

The LLM needs to know what tools it can call. Groq's tool-calling API uses JSON schemas to describe each function's name, purpose, and parameters. You need one definition for each of your three database functions.

  • Add the tool definitions for list_tables and get_table_schema by pasting this below your SYSTEM_PROMPT:
# Tool definitions describe each function for the LLM
TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "list_tables",
            "description": "List all tables in the database",
            "parameters": {"type": "object", "properties": {}, "required": []},
        },
    },
    {
        "type": "function",
        "function": {
            "name": "get_table_schema",
            "description": "Get column names, data types, and foreign key relationships for a specific table",
            "parameters": {
                "type": "object",
                "properties": {
                    "table_name": {
                        "type": "string",
                        "description": "Name of the table to inspect",
                    }
                },
                "required": ["table_name"],
            },
        },
    },

What are tool definitions?

Each tool definition is a JSON object with a name, a description, and a parameters schema. The LLM reads these to decide which function to call and what arguments to pass.

Notice that list_tables has no parameters (it takes no arguments), while get_table_schema requires a table_name string.

  • Add the run_query tool definition, the closing bracket, and the configuration variables by pasting this directly below the code you just added:
    {
        "type": "function",
        "function": {
            "name": "run_query",
            "description": "Execute a read-only SQL query against the database and return results (max 50 rows)",
            "parameters": {
                "type": "object",
                "properties": {
                    "sql": {
                        "type": "string",
                        "description": "The SELECT query to execute",
                    }
                },
                "required": ["sql"],
            },
        },
    },
]

# Map tool names to the actual Python functions
AVAILABLE_TOOLS = {
    "list_tables": list_tables,
    "get_table_schema": get_table_schema,
    "run_query": run_query,
}

# Safety limit to prevent infinite loops
MAX_ITERATIONS = 10

What does the configuration do?

AVAILABLE_TOOLS is a dictionary that maps tool name strings to the actual Python functions. When the LLM says "call run_query", your code looks up the real function here.

MAX_ITERATIONS caps how many times the agent loop can repeat. Without this, a confused model could loop forever, burning through your API quota.

  • Save agent.py.
  • Verify the module loads without errors by running this command in your terminal:
python -c "import agent; print('Agent module loaded')"

You should see Agent module loaded printed with no errors.

Seeing an import error?

  • If you see ModuleNotFoundError: No module named 'groq', make sure you installed the Groq SDK in Step 1. Run pip install groq to fix it.
  • If you see ModuleNotFoundError: No module named 'db_tools', make sure agent.py and db_tools.py are in the same folder. Your terminal must also be in that folder.
  • If you see a SyntaxError, double check that the TOOLS list brackets and braces all match up. The closing ] must appear after all three tool definitions.

Help me debug my agent.py import error

Build the agent loop

The agent loop is the core pattern that makes this an agent instead of a one-shot API call. It sends the conversation to Groq, checks if the model wants to call any tools, executes those tools, feeds the results back, and repeats until the model produces a plain-text answer.

  • Add the run_agent function definition by pasting this below MAX_ITERATIONS in agent.py:
def run_agent(user_question):
    """Run the SQL agent with a user question."""
    print(f"\nQuestion: {user_question}")
    print("-" * 60)

    # Start the conversation with the system prompt and user question
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": user_question},
    ]

    for i in range(MAX_ITERATIONS):
        # Send the full conversation to Groq with available tools
        response = client.chat.completions.create(
            model="qwen/qwen3.6-27b",
            messages=messages,
            tools=TOOLS,
            tool_choice="auto",
        )

        message = response.choices[0].message
        messages.append(message)

What does this code do?

The messages list is the full conversation history. It starts with the system prompt and the user's question. Every API call sends the entire conversation so the model has full context.

The tool_choice="auto" parameter lets the model decide whether to call a tool or respond with text. Each iteration, the model either asks to use a tool or gives a final answer.

The second half of the loop handles the model's response. If the model requested tool calls, your code executes each one and feeds the results back. If no tool calls were requested, the model has produced its final answer.

  • Add the tool-call processing and final-answer handling by pasting this directly below the code you just added, inside the for loop (indented by 8 spaces):
        if message.tool_calls:
            for tool_call in message.tool_calls:
                func_name = tool_call.function.name
                func_args = json.loads(tool_call.function.arguments)

                # Print each tool call so you can see the agent's reasoning
                print(f"  Tool: {func_name}({func_args})")

                if func_name in AVAILABLE_TOOLS:
                    result = AVAILABLE_TOOLS[func_name](**func_args)
                else:
                    result = json.dumps({"error": f"Unknown tool: {func_name}"})

                print(f"  Result: {result[:200]}...")

                # Feed the tool result back into the conversation
                messages.append(
                    {
                        "role": "tool",
                        "tool_call_id": tool_call.id,
                        "name": func_name,
                        "content": result,
                    }
                )
        else:
            print(f"\nAnswer: {message.content}")
            return message.content

    print("\nAgent reached maximum iterations without a final answer.")
    return None

How does the agent loop work?

  • When message.tool_calls is present, the model wants to call a function. Your code looks up the function name in AVAILABLE_TOOLS, runs it with the arguments the model chose, and appends the result as a tool message.
  • When there are no tool calls, the model has finished reasoning and produced its final text answer. The loop exits with return.
  • The print statements let you watch the agent's reasoning in real time. You will see every tool it calls and a preview of each result.

Why is this called an agent loop?

A regular LLM call is one-shot: you send a prompt, you get a response. An agent loop keeps the conversation going across multiple rounds. The model calls a tool, gets the result, thinks about what to do next, maybe calls another tool, and keeps going until the question is answered.

  • Add the interactive CLI loop at the bottom of agent.py so you can type questions from the terminal:
if __name__ == "__main__":
    print("SQL Agent ready. Ask questions about the TPC-H database.")
    print("Type 'quit' to exit.\n")

    # Keep asking questions until the user types quit
    while True:
        question = input("Ask a question: ")
        if question.lower() in ("quit", "exit", "q"):
            break
        run_agent(question)

What does this code do?

The if __name__ == "__main__" block only runs when you execute the file directly with python agent.py. It starts an interactive loop where you can keep asking the agent questions. Type quit to exit.

  • Save agent.py.

✔️ Awesome, I've got everything!

Great. Double check you have saved agent.py before moving on.

ⓧ I'd like to double check the full code

Scroll up and confirm you have each of these sections in agent.py, in order:

  • Imports (os, json, Groq, and three functions from db_tools) and the client initialization.
  • SYSTEM_PROMPT with the minimal agent role description.
  • TOOLS list with all three tool definitions (list_tables, get_table_schema, run_query).
  • AVAILABLE_TOOLS dictionary and MAX_ITERATIONS = 10.
  • run_agent function with the message setup, API call loop, tool-call processing, and the final-answer else branch.
  • if __name__ == "__main__" block with the interactive while True loop.
Test with a simple single-table question

The agent is ready for its first real test. A simple question that only needs one table is the best way to confirm the agent loop, tool calling, and database connection all work together.

Before you run this, think about what you expect the agent to do. How many tool calls do you think it will make to answer a basic counting question?

  • Start the agent by running:
python agent.py
  • When you see the Ask a question: prompt, type:
How many customers are there?

Watch the terminal output. You should see the agent call list_tables to discover the available tables, then call run_query with a COUNT query on the customer table. The final answer should report 1,500,000 customers.

Why 1,500,000?

Your database uses TPC-H benchmark data at scale factor 10. The TPC-H specification sets the customer count at 150,000 per scale factor. At SF10, that is exactly 1,500,000 rows.

Agent not responding or hitting an error?

  • If you see an authentication error, make sure your GROQ_API_KEY environment variable is still set in this terminal session. Run echo $GROQ_API_KEY to check.
  • If the agent loops without producing an answer, check that your AVAILABLE_TOOLS dictionary keys match the name values in your TOOLS definitions. They must match exactly.
  • If you see a database connection error, make sure your PostgreSQL Docker container is still running.

Help me debug my SQL agent

Test with a complex multi-table question

The agent handled a single-table question with no trouble. Now try something harder that requires joining multiple tables together.

Before you ask the next question, predict: will the agent produce the correct answer when the question involves joins across several tables?

  • At the Ask a question: prompt, type:
What are the top 5 nations by total revenue?

Watch the agent's tool calls closely. You will likely see it generate SQL that guesses at join conditions or uses incorrect column relationships. The query may error out, return wrong results, or miss tables entirely.

Why does the agent struggle with this question?

Answering "top 5 nations by total revenue" requires joining at least 4 tables: nation, customer, orders, and lineitem. The model needs to know the foreign key relationships to write the correct JOIN conditions.

Your current system prompt says nothing about inspecting foreign keys before writing SQL. The model has the get_table_schema tool available, but the prompt does not instruct it to use schemas to discover join conditions. Without that guidance, the model guesses at column names and relationships.

This failure is expected. The agent has the right tools but no structured process for using them. In the next step, you will rewrite the system prompt to guide the agent through a step-by-step schema discovery process that fixes exactly this problem.

  • Type quit to exit the agent.

Your agent loop is working. It can autonomously call tools, read results, and produce answers for simple questions. The gap is that it does not know how to discover and use foreign key relationships for complex joins. Next up, you will make the agent schema-aware so it handles multi-table queries correctly.

Make the Agent Schema-Aware

In the last step, your agent answered a simple customer count correctly. But when you asked "What are the top 5 nations by total revenue?", it wrote broken SQL because it had no idea how the tables connect to each other.

The root cause is the system prompt. It tells the model to "use the available tools" but never explains how. Without explicit instructions to inspect foreign key relationships before writing SQL, the model guesses at join conditions and gets them wrong. A structured process in the prompt fixes this.

In this step, get ready to:
  • Rewrite the system prompt with a step-by-step schema-discovery process.
  • Test the agent on the same multi-table question that previously failed.
Rewrite the system prompt with a structured process

The current SYSTEM_PROMPT is four lines of vague instruction. The model needs a concrete process: discover tables, inspect schemas, use foreign keys for joins, then run the query.

  • In agent.py, find the SYSTEM_PROMPT variable near the top of the file. It currently looks like this:
SYSTEM_PROMPT = """You are a SQL agent that answers questions about a PostgreSQL database.
The database contains TPC-H benchmark data.

Use the available tools to discover the schema and run queries to answer the user's question.
Format your final answer as a clear, human-readable response."""
  • Replace the entire SYSTEM_PROMPT assignment with this updated version:
SYSTEM_PROMPT = """You are a SQL agent that answers questions about a PostgreSQL database.
The database contains TPC-H benchmark data: an order management system with parts, suppliers, customers, orders, and line items across nations and regions.

Follow this process for every question:
1. Call list_tables to see available tables.
2. Call get_table_schema for each table relevant to the question. Pay close attention to foreign key relationships to determine correct JOIN conditions.
3. Write a SQL query using only the columns and relationships you discovered. Never guess column names.
4. Call run_query to execute it.

Rules:
- Only generate SELECT queries. Never use INSERT, UPDATE, DELETE, DROP, or any DDL/DML.
- Always LIMIT results to avoid returning too many rows. Use LIMIT 20 unless the user asks for more.
- Never SELECT * from large tables. Always select only the columns needed.
- When joining tables, always use the foreign key relationships from the schema.
- Format your final answer as a clear, human-readable response based on the query results."""

What changed in the prompt?

The old prompt said "use the available tools" and left the model to figure out the rest. The new prompt gives the model a numbered process to follow for every question.

  • Steps 1-2 force the model to call list_tables and get_table_schema before writing any SQL. This is how the agent discovers foreign key relationships instead of guessing at joins.
  • Step 3 tells the model to only use columns and relationships it actually discovered. No guessing.
  • The Rules section adds guardrails like row limits and column selection. On a 60M+ row database, a SELECT * from the lineitem table would be catastrophic.

Why teach the agent a process instead of hardcoding the schema?

You could paste the entire database schema into the prompt and skip the tool calls. But that approach breaks the moment the database changes. By teaching the agent to discover the schema at query time, it stays accurate even if tables or columns are added later.

  • Save agent.py.

✔️ Awesome, I've got everything!

Great. Double check you've saved agent.py with the updated SYSTEM_PROMPT.

ⓧ I'd like to double check the full code

The only change in this step is the SYSTEM_PROMPT variable. Scroll up and confirm your SYSTEM_PROMPT matches the updated version above. The rest of agent.py (the TOOLS list, AVAILABLE_TOOLS dict, run_agent function, and __main__ block) should be unchanged from the previous step.

Seeing a syntax error when you save?

  • Make sure the triple quotes """ are present at both the start and end of the prompt string. A missing closing """ causes Python to swallow the rest of the file into the string.
  • Check that you replaced the old SYSTEM_PROMPT entirely. If parts of the old prompt remain, you may have duplicate variable assignments.

Help me fix a syntax error in my SYSTEM_PROMPT string in agent.py.

Test the schema-aware agent on a multi-table question

Your agent loop already prints each tool call and its result as the agent runs. This trace shows you exactly which tables the agent inspects, what foreign keys it discovers, and what SQL it writes.

  • Run the agent by running this command:
python agent.py

Before you type a question, think about what you expect to see differently this time. The prompt now tells the model to call get_table_schema for relevant tables before writing SQL. How many schema-inspection tool calls do you expect for a question that spans nations, orders, and line items?

  • At the Ask a question: prompt, type:
What are the top 5 nations by total revenue?

Watch the tool call trace scroll by in your terminal. You should see the agent follow the structured process from the prompt.

What should I see?

The agent should make several tool calls in sequence, following the process you defined in the prompt.

  • First, a list_tables call to discover the 8 TPC-H tables.
  • Then, multiple get_table_schema calls for tables relevant to the question (like nation, orders, lineitem, and customer). Each result includes the foreign key relationships the model needs for correct joins.
  • Finally, a run_query call with a multi-table JOIN query that uses the FK relationships it discovered.

The final answer lists 5 nations ranked by total revenue. The exact SQL the model generates may vary between runs, but it should produce correct results because it built the joins from actual foreign key data.

Why does the model's SQL vary between runs?

LLMs are non-deterministic. The model might join through customer on one run and through supplier on another, depending on how it interprets "revenue." Both paths can produce valid results because the structured prompt ensures the model always checks the actual foreign keys before writing the join.

  • Type quit to exit the agent.

Agent still writing bad SQL?

  • Check the tool call trace. If you don't see get_table_schema calls before the run_query call, the model is skipping schema inspection. Make sure the numbered process in your SYSTEM_PROMPT matches the version above exactly.
  • If the agent returns an error from run_query and then gives up, that is expected at this stage. The agent does not yet have error recovery. You will add that in the next step.
  • LLM responses are non-deterministic. Try running the same question again. A different run may produce a correct query.

Help me debug why my SQL agent is not inspecting table schemas before writing queries.

Your agent now discovers the database schema on the fly and writes accurate multi-table joins. Next up, you will add safety guardrails that block destructive queries and error recovery that lets the agent fix its own mistakes.

Add Safety and Error Recovery

Your agent now handles complex multi-table queries by discovering foreign key relationships before writing SQL. The analytical power is there.

But there's a dangerous gap. Nothing stops the agent from generating a DELETE or DROP statement if someone phrases their question the right way. And when the model writes bad SQL, the agent crashes instead of recovering. In this step, you'll add a validation gate that blocks destructive queries and an error recovery loop that lets the agent self-correct.

In this step, get ready to:
  • Add a query validation function that blocks destructive SQL statements.
  • Wire the validation tool into the agent and update the system prompt.
  • Add error recovery so the agent retries failed queries up to 3 times.
Add a query validation function

The agent's system prompt tells it to only write SELECT queries. But a prompt instruction is not enforcement. You need a Python function that inspects the actual SQL string and rejects anything destructive before it reaches the database.

  • In db_tools.py, find the run_query function.
  • Add the following validate_query function directly above run_query:
def validate_query(sql):
    """Check a SQL query for forbidden DML statements before execution."""
    forbidden = ["INSERT", "UPDATE", "DELETE", "DROP", "ALTER", "TRUNCATE", "CREATE"]
    sql_upper = sql.strip().upper()
    for keyword in forbidden:
        if sql_upper.startswith(keyword):
            return json.dumps(
                {
                    "valid": False,
                    "reason": f"{keyword} statements are not allowed. Only SELECT queries are permitted.",
                }
            )
    return json.dumps({"valid": True, "sql": sql})

What does this code do?

The forbidden list contains seven SQL keywords that modify or destroy data. The function uppercases the incoming SQL and checks whether it starts with any of them.

If a match is found, it returns a JSON object with valid: false and a reason. If the SQL is safe, it returns valid: true with the original query passed through.

  • Save db_tools.py.
  • Test the function directly by running this command in your terminal:
python -c "from db_tools import validate_query; print(validate_query('SELECT * FROM customer')); print(validate_query('DELETE FROM customer'))"

The first call should return {"valid": true, "sql": "SELECT * FROM customer"}. The second should return {"valid": false, "reason": "DELETE statements are not allowed. Only SELECT queries are permitted."}.

Seeing an ImportError?

  • Make sure you added validate_query above run_query and not inside another function.
  • Check that the function name is spelled exactly validate_query with an underscore.
  • Ask for help: Why can't Python import my validate_query function from db_tools?

✔️ Awesome, I've got everything!

Great. Make sure you've saved db_tools.py before moving on.

ⓧ I'd like to double check the full code

Your db_tools.py should now have these functions in this order. Scroll up and confirm each one is present:

  • The imports (psycopg and json) and DB_CONFIG at the top.
  • get_connection() helper.
  • list_tables() function.
  • get_table_schema(table_name) function.
  • validate_query(sql) function (the one you just added).
  • run_query(sql) function.
Wire the validation tool into the agent

The validate_query function exists, but the agent doesn't know about it yet. You need to import it, give the LLM a tool definition for it, and update the system prompt so the agent always validates before executing.

  • In agent.py, find the import line at the top of the file:
from db_tools import list_tables, get_table_schema, run_query
  • Add validate_query to the import so the line looks like this:
from db_tools import list_tables, get_table_schema, validate_query, run_query
  • Next, scroll down to the TOOLS list. Find the get_table_schema tool definition and add the following validate_query definition directly after it (before the run_query definition):
    {
        "type": "function",
        "function": {
            "name": "validate_query",
            "description": "Check a SQL query for forbidden statements before execution",
            "parameters": {
                "type": "object",
                "properties": {
                    "sql": {
                        "type": "string",
                        "description": "The SQL query to validate",
                    }
                },
                "required": ["sql"],
            },
        },
    },

What is a tool definition?

Each entry in the TOOLS list is a JSON schema that tells the LLM what the function does, what parameters it accepts, and which are required. The model reads these descriptions to decide when and how to call each tool.

  • Now scroll down to the AVAILABLE_TOOLS dictionary. Find this code:
AVAILABLE_TOOLS = {
    "list_tables": list_tables,
    "get_table_schema": get_table_schema,
    "run_query": run_query,
}
  • Add validate_query to the dictionary so it looks like this:
AVAILABLE_TOOLS = {
    "list_tables": list_tables,
    "get_table_schema": get_table_schema,
    "validate_query": validate_query,
    "run_query": run_query,
}
  • Save agent.py and verify the wiring by running this quick check:
python -c "from agent import TOOLS, AVAILABLE_TOOLS; print(len(TOOLS), len(AVAILABLE_TOOLS))"

You should see 4 4. Previously you had 3 tools. The new count confirms validate_query is registered.

Still seeing 3?

  • Make sure you placed the validate_query tool definition inside the TOOLS list (between the square brackets), not outside it.
  • Check that the import line includes validate_query from db_tools.
  • Ask for help: My TOOLS list still shows 3 tools after I added validate_query. What did I miss?

The agent now knows the tool exists, but the system prompt doesn't tell it to use validation yet. Time to fix the prompt.

  • In agent.py, scroll up to the SYSTEM_PROMPT and find this line in the process steps:
4. Call run_query to execute it.
  • Replace that line with these two new steps:
4. Call validate_query to check your SQL before executing it.
5. If validation passes, call run_query to execute it.

Why validate before executing?

The old step 4 went straight from writing SQL to executing it. Now the agent must pass its SQL through validate_query first. If the validation fails, the model sees the rejection reason in the tool result and can tell the user why the query was blocked.

Before you test this, what do you think the agent will do when you ask it to delete all customers? Will it refuse immediately, or will it try to build a DELETE query first?

  • Save agent.py and run the agent:
python agent.py
  • At the prompt, type:
Delete all customers

Watch the tool trace in your terminal. You should see the agent call validate_query and receive the valid: false response. The agent then tells you the query was blocked instead of running it.

  • Type quit to exit the agent.

Agent didn't call validate_query?

  • Check that step 4 in your SYSTEM_PROMPT says Call validate_query and not Call run_query.
  • Confirm the validate_query tool definition is inside the TOOLS list, not outside it.
  • Ask for help: My agent skipped validate_query and ran the SQL directly. How do I fix the prompt?
Add error recovery to the system prompt

The validation gate handles malicious queries. But the agent still has a weakness: when the model writes syntactically wrong SQL, the error comes back from run_query and the agent doesn't know it should try again. You need to tell it to read the error message and retry.

  • In agent.py, find step 5 in your SYSTEM_PROMPT:
5. If validation passes, call run_query to execute it.
  • Add the following as step 6 on the next line, directly after step 5:
6. If you get an error, read the error message carefully, fix the query, and try again (up to 3 attempts).

How does error recovery work here?

When run_query catches a PostgreSQL exception, it returns the error message as JSON in the tool result. The LLM sees this message in the conversation, understands what went wrong (a missing column, a bad join, a syntax error), and generates a corrected query on its next turn.

The "up to 3 attempts" cap prevents infinite loops if the model keeps generating broken SQL. With MAX_ITERATIONS = 10 already set in your agent loop, the agent has room to retry while still having a hard stop.

Now for the real test. What do you think will happen when you ask a complex question that the model might not get right on the first SQL attempt?

  • Save agent.py and run the agent:
python agent.py
  • At the prompt, type:
Show me parts with a supply cost over 900 from European suppliers

Watch the tool trace in your terminal. You should see the agent call list_tables, then get_table_schema for the relevant tables, then validate_query to check its SQL, and finally run_query to execute it. If the model writes a bad query on its first attempt, you'll see the error come back and the agent try a corrected version.

  • Type quit to exit the agent.

Agent hit maximum iterations without answering?

  • This can happen if the model keeps writing bad SQL and runs out of retries. Try the question again. LLM responses vary between runs.
  • Check that step 6 in your prompt includes the phrase "up to 3 attempts" so the model knows to limit its retries.
  • Ask for help: My SQL agent keeps hitting the maximum iteration limit when I ask complex questions. How can I debug this?

✔️ Awesome, I've got everything!

Great. Make sure you've saved agent.py before moving on.

ⓧ I'd like to double check the full code

Your agent.py should now have all of these pieces. Scroll up and confirm each one is present in order:

  • Import line with list_tables, get_table_schema, validate_query, and run_query.
  • SYSTEM_PROMPT with 6 numbered process steps (including validate_query at step 4 and error retry at step 6).
  • TOOLS list with 4 tool definitions: list_tables, get_table_schema, validate_query, run_query.
  • AVAILABLE_TOOLS dictionary with all 4 functions mapped.
  • The run_agent function and the if __name__ == "__main__" interactive loop (unchanged from before).

Secret mission

Add Query Performance Analysis

Your agent answers questions and blocks dangerous queries. But can it tell you why a query runs fast or slow? In this mission, you will add an EXPLAIN ANALYZE tool that lets the agent break down query execution plans in plain English.

Clean Up Your Resources

Clean Up Your Resources

This project runs entirely locally, and the Groq free tier has no charges. Decide whether to keep your resources running, pause them to come back later, or delete them entirely.

Resources you used:

  • db_tools.py and agent.py project files.
  • groq, psycopg, and psycopg-binary Python packages.
  • GROQ_API_KEY environment variable (session-only, clears when you close the terminal).
  • Pre-existing PostgreSQL Docker container (not created by this project, but can be paused).

Keep everything running

No action needed. Choose this if you are still actively building on your SQL agent or want to keep experimenting with questions.

  • Your agent.py and db_tools.py files stay on your machine wherever you saved them.
  • The groq and psycopg packages remain installed in your Python environment.
  • The Groq free tier has no ongoing charges. Your PostgreSQL Docker container uses memory only while running.

Pause - I'll come back to this later

Stop the PostgreSQL Docker container to free up memory. Your project files and installed packages stay in place so you can pick up where you left off.

  • Find your PostgreSQL container name by running:
docker ps
  • Stop the container by running:
docker stop [[CONTAINER_NAME="your PostgreSQL container name"]]
  • When you are ready to resume, start the container again by running:
docker start [[CONTAINER_NAME="your PostgreSQL container name"]]
  • Re-export your GROQ_API_KEY if you have opened a new terminal session since your last run.

Delete - I don't want to use this again

Remove the project files and uninstall the Python packages you added during this project.

  • Delete the project files by running:
rm db_tools.py agent.py
  • Uninstall the Python packages by running:
pip uninstall groq psycopg psycopg-binary
  • Type y and press Enter when prompted to confirm each package removal.
  • Your GROQ_API_KEY environment variable clears automatically when you close the terminal.

What about the PostgreSQL Docker container?

Your PostgreSQL container with the TPC-H data existed before this project. If you still need it for other work, leave it running. If you want to remove it entirely, run docker rm -f followed by your container name.

Nice Work!

Nice Work!

You just built a fully functional SQL agent from scratch. You can type a plain-English question, and your agent autonomously discovers the database schema, writes SQL, validates it, executes it against a 60M+ row PostgreSQL database, and responds with a clear answer.

You've learned how to:

  • Built database introspection tools that dynamically discover table schemas and foreign key relationships from PostgreSQL's information_schema, giving the LLM the context it needs to write correct multi-table queries.
  • Implemented an agent loop that processes multi-step tool calls using Groq's tool-calling API with the qwen/qwen3.6-27b model, letting the LLM autonomously reason about which tables to inspect, what SQL to write, and how to present results.
  • Added a validation layer that blocks destructive SQL before it reaches the database, and an error recovery loop that lets the agent read PostgreSQL error messages and self-correct bad queries.
  • Secret Mission: Added EXPLAIN ANALYZE query performance analysis so the agent can explain why a query is fast or slow.

Ready to quiz yourself?