Building a SQL Query Agent with Pydantic AI: Complete Guide

Building a SQL Query Agent with Pydantic AI: Complete Guide

Large language models are excellent at understanding natural language, but turning that understanding into reliable, executable SQL has historically been a fragile exercise. Pydantic AI changes the equation by combining structured outputs, dependency injection, and a typed agent framework that keeps your database interactions predictable. In this guide, you will build a production-ready SQL Query Agent that accepts plain English questions, translates them into safe SQL, executes them against a database, and returns both the rows and a natural-language summary.

What Is a SQL Query Agent?

A SQL Query Agent is an LLM-powered system that acts as an intermediary between a non-technical user and a relational database. Instead of writing SQL manually, the user asks questions such as "Which customers placed the most orders last month?" The agent inspects the schema, generates a query, runs it, and explains the results. With Pydantic AI, this workflow is encoded as a typed agent with explicit tools, dependencies, and a validated result model, so every step is auditable and testable.

Why It Matters

• Structured outputs: Pydantic AI guarantees the model returns data matching a schema, eliminating brittle string parsing.
• Dependency injection: Database connections and configuration are passed in as typed dependencies, making agents easy to test and swap between environments.
• Tool use: The agent can call functions like list_tables or run_sql in a controlled loop, mirroring how a human analyst works.
• Safety: Read-only connections, query validation, and result limits can be enforced at the framework level rather than relying on prompt discipline alone.
• Observability: Because the agent is typed and instrumented, you can log every tool call, token usage, and validation failure.

Prerequisites and Project Setup

Start by creating a virtual environment and installing the required packages. This tutorial uses SQLite for portability, but the same pattern applies to PostgreSQL, MySQL, or any SQLAlchemy-supported backend.

python -m venv .venv
source .venv/bin/activate
pip install pydantic-ai sqlalchemy aiosqlite python-dotenv

Create a .env file with your model provider key. Pydantic AI supports OpenAI, Anthropic, Gemini, Groq, and others. This guide uses OpenAI, but swapping providers only changes one line.

OPENAI_API_KEY=sk-your-key-here

Defining the Database and Sample Schema

Before building the agent, create a small e-commerce schema with customers, orders, and order items. This gives the agent something realistic to query.

import sqlite3
from pathlib import Path

DB_PATH = "shop.db"

def seed_database() -> None:
    if Path(DB_PATH).exists():
        Path(DB_PATH).unlink()
    conn = sqlite3.connect(DB_PATH)
    cur = conn.cursor()
    cur.executescript("""
    CREATE TABLE customers (
        id INTEGER PRIMARY KEY,
        name TEXT NOT NULL,
        email TEXT UNIQUE NOT NULL,
        country TEXT NOT NULL
    );
    CREATE TABLE orders (
        id INTEGER PRIMARY KEY,
        customer_id INTEGER NOT NULL REFERENCES customers(id),
        placed_at TEXT NOT NULL,
        status TEXT NOT NULL
    );
    CREATE TABLE order_items (
        id INTEGER PRIMARY KEY,
        order_id INTEGER NOT NULL REFERENCES orders(id),
        product_name TEXT NOT NULL,
        quantity INTEGER NOT NULL,
        unit_price REAL NOT NULL
    );
    """)
    cur.executemany(
        "INSERT INTO customers (name, email, country) VALUES (?, ?, ?)",
        [
            ("Alice Lee", "alice@example.com", "Singapore"),
            ("Bob Singh", "bob@example.com", "India"),
            ("Carla Nunez", "carla@example.com", "Spain"),
            ("David Kim", "david@example.com", "South Korea"),
        ],
    )
    cur.executemany(
        "INSERT INTO orders (customer_id, placed_at, status) VALUES (?, ?, ?)",
        [
            (1, "2024-11-02", "delivered"),
            (1, "2024-12-15", "delivered"),
            (2, "2024-12-20", "shipped"),
            (3, "2024-12-28", "pending"),
            (4, "2025-01-05", "delivered"),
        ],
    )
    cur.executemany(
        "INSERT INTO order_items (order_id, product_name, quantity, unit_price) VALUES (?, ?, ?, ?)",
        [
            (1, "USB-C Cable", 3, 9.99),
            (1, "Mouse Pad", 1, 12.50),
            (2, "Mechanical Keyboard", 1, 89.00),
            (3, "Webcam", 2, 45.00),
            (4, "Laptop Stand", 1, 35.00),
            (5, "Notebook", 5, 4.50),
        ],
    )
    conn.commit()
    conn.close()

if __name__ == "__main__":
    seed_database()
    print("Database seeded.")

Run this script once to create shop.db. The agent will introspect this schema at runtime.

Designing the Result Model

One of the core strengths of Pydantic AI is that the final agent output is validated against a Pydantic model. Define a result type that captures the SQL, the rows, and a human-readable explanation.

from typing import Any
from pydantic import BaseModel, Field

class QueryResult(BaseModel):
    question: str = Field(description="The original natural-language question.")
    sql: str = Field(description="The SQL query that was executed.")
    rows: list[dict[str, Any]] = Field(
        default_factory=list,
        description="Rows returned by the query, as dictionaries.",
    )
    row_count: int = Field(description="Number of rows returned.")
    explanation: str = Field(
        description="A concise natural-language summary of the findings."
    )

Because this model is typed, the LLM is forced to produce output that conforms to it. If the model omits a field or returns the wrong type, Pydantic AI retries automatically.

Creating the Dependency Container

Pydantic AI uses dataclasses or Pydantic models as dependency containers. This keeps your agent decoupled from global state and makes it trivial to inject a mock during tests.

from dataclasses import dataclass
import aiosqlite

@dataclass
class DbDeps:
    conn: aiosqlite.Connection
    max_rows: int = 50

async def make_deps() -> DbDeps:
    conn = await aiosqlite.connect(DB_PATH)
    conn.row_factory = aiosqlite.Row
    return DbDeps(conn=conn, max_rows=50)

Building the Agent and Tools

Now assemble the agent. The agent receives a system prompt that explains its role, a set of tools it can call, and the result model it must satisfy. Two tools are exposed: one to inspect the schema and one to execute read-only SQL.

from pydantic_ai import Agent, RunContext
import re

READ_ONLY_PATTERN = re.compile(r"^\s*(select|with|explain)\b", re.IGNORECASE)

agent = Agent(
    "openai:gpt-4o-mini",
    deps_type=DbDeps,
    result_type=QueryResult,
    system_prompt=(
        "You are a SQL analyst agent. You help users answer questions about an "
        "e-commerce database with tables: customers, orders, order_items. "
        "Always inspect the schema first using list_schema, then write a single "
        "read-only SQL query using run_sql. Never invent columns. Limit results "
        "to a reasonable number of rows. After running the query, summarize the "
        "findings in plain English."
    ),
)

@agent.tool
async def list_schema(ctx: RunContext[DbDeps]) -> str:
    """Return the CREATE statements for every table in the database."""
    cursor = await ctx.deps.conn.execute(
        "SELECT sql FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'"
    )
    rows = await cursor.fetchall()
    return "\n\n".join(row[0] for row in rows)

@agent.tool
async def run_sql(ctx: RunContext[DbDeps], query: str) -> dict[str, Any]:
    """Execute a read-only SQL query and return up to max_rows rows."""
    if not READ_ONLY_PATTERN.match(query):
        return {"error": "Only SELECT, WITH, or EXPLAIN queries are allowed."}
    try:
        cursor = await ctx.deps.conn.execute(query)
        rows = await cursor.fetchmany(ctx.deps.max_rows)
        columns = [desc[0] for desc in cursor.description]
        data = [dict(zip(columns, row)) for row in rows]
        return {"columns": columns, "rows": data, "row_count": len(data)}
    except Exception as exc:
        return {"error": str(exc)}

Notice the guardrail: run_sql rejects any statement that does not begin with SELECT, WITH, or EXPLAIN. This is defense in depth on top of using a read-only database user in production.

Running the Agent

With the agent defined, wire it into an async entry point. The agent loop handles tool calls, validation retries, and final result construction automatically.

import asyncio

async def ask(question: str) -> QueryResult:
    deps = await make_deps()
    try:
        result = await agent.run(question, deps=deps)
        return result.data
    finally:
        await deps.conn.close()

async def main() -> None:
    questions = [
        "Which customer spent the most money in total?",
        "How many orders are still pending?",
        "List all products ordered by customers outside India.",
    ]
    for q in questions:
        print(f"\nQ: {q}")
        answer = await ask(q)
        print(f"SQL: {answer.sql}")
        print(f"Rows: {answer.row_count}")
        print(f"Answer: {answer.explanation}")

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

When you run this script, the agent will call list_schema, generate a query, call run_sql, and then produce a QueryResult with a natural-language summary. Each step is logged by Pydantic AI and can be inspected via the result.all_messages() method.

Streaming and Conversation History

For interactive applications, you may want to stream the explanation token by token and maintain context across turns. Pydantic AI supports both through agent.run_stream and message history.

async def chat_session() -> None:
    deps = await make_deps()
    history = None
    try:
        while True:
            user_input = input("You: ").strip()
            if user_input.lower() in {"exit", "quit"}:
                break
            async with agent.run_stream(
                user_input, deps=deps, message_history=history
            ) as stream:
                print("Agent: ", end="", flush=True)
                async for chunk in stream.stream_text(delta=True):
                    print(chunk, end="", flush=True)
                print()
                history = stream.all_messages()
    finally:
        await deps.conn.close()

Best Practices

• Use read-only credentials: Never rely solely on regex validation. Connect with a database user that has only SELECT privileges.
• Cap result sets: Enforce max_rows in the tool and add a LIMIT clause in the system prompt to prevent runaway queries.
• Provide schema context: Include column descriptions and foreign-key relationships in the system prompt or a dedicated describe_table tool to reduce hallucinated columns.
• Log every query: Persist the generated SQL, parameters, and row counts for audit and debugging. Pydantic AI's message log makes this straightforward.
• Test with deterministic mocks: Inject a fake DbDeps in unit tests so you can assert agent behavior without hitting a real database or LLM.
• Handle empty results gracefully: Instruct the model to explain when a query returns zero rows rather than fabricating data.
• Rotate providers carefully: Because the result model is provider-agnostic, you can switch from openai:gpt-4o-mini to anthropic:claude-3-5-sonnet by changing one string, but re-run your eval suite to confirm tool-calling reliability.

Extending the Agent

Once the core loop works, you can extend it with additional tools: a get_sample_rows tool that returns a few example rows per table, a validate_sql tool that runs EXPLAIN before execution, or a charting tool that converts results into a base64-encoded PNG. Each new tool is just a decorated async function with typed parameters, so the agent's capability surface grows without restructuring existing code.

You can also layer a second agent on top. For example, a "planner" agent could decompose a complex question into sub-questions, each handled by the SQL agent, and a "synthesizer" agent could merge the answers. Pydantic AI's typed dependencies make chaining agents safe and composable.

Conclusion

Building a SQL Query Agent with Pydantic AI gives you the best of both worlds: the flexibility of natural-language interaction and the rigor of typed, validated, testable code. By defining a clear result model, injecting database dependencies, exposing constrained tools, and enforcing read-only guardrails, you create a system that is safe enough for real users and maintainable enough for long-term development. Start with the schema introspection and single-query pattern shown here, then iterate by adding domain-specific tools, richer schema metadata, and evaluation harnesses to measure accuracy as your data and user questions evolve.