DEV Community

Nebula
Nebula

Posted on

How to Build a Text-to-SQL Agent with Python in 10 Minutes

#ai #tutorial #python #beginners

You want to ask your database questions in plain English. Most tutorials make this harder than it needs to be — spinning up PostgreSQL, installing heavy ORMs, writing 200 lines of boilerplate.

Here's a text-to-SQL agent in under 40 lines of Python. It uses PydanticAI for the agent logic and SQLite so you don't need any database server.

The Code 

import sqlite3
import asyncio
from pydantic_ai import Agent, RunContext, ModelRetry
from pydantic import BaseModel
from dataclasses import dataclass

# 1. Set up a sample SQLite database
conn = sqlite3.connect(":memory:")
conn.execute(\"\"\"CREATE TABLE employees (
    id INTEGER PRIMARY KEY,
    name TEXT,
    department TEXT,
    salary INTEGER,
    hire_date TEXT
)\"\"\")
conn.executemany(
    "INSERT INTO employees (name, department, salary, hire_date) VALUES (?, ?, ?, ?)",
    [
        ("Alice", "Engineering", 120000, "2023-01-15"),
        ("Bob", "Marketing", 85000, "2023-06-01"),
        ("Carol", "Engineering", 135000, "2022-03-20"),
        ("Dave", "Sales", 90000, "2024-01-10"),
        ("Eve", "Engineering", 110000, "2024-07-01"),
    ],
)
conn.commit()


# 2. Define dependencies and output schema
@dataclass
class Deps:
    conn: sqlite3.Connection


class QueryResult(BaseModel):
    sql: str
    explanation: str
    rows: list[dict]


# 3. Create the agent
agent = Agent(
    "openai:gpt-4o-mini",
    deps_type=Deps,
    system_prompt=(
        "You are a SQL assistant. Given the 'employees' table with columns: "
        "id, name, department, salary, hire_date — generate a SELECT query "
        "for the user's question. Return ONLY the SQL query, nothing else."
    ),
)


# 4. Add a tool that runs the generated SQL
@agent.tool
async def run_query(ctx: RunContext[Deps], sql_query: str) -> str:
    \"\"\"Execute a SQL query against the employees database and return results.\"\"\"
    if not sql_query.strip().upper().startswith("SELECT"):
        raise ModelRetry("Only SELECT queries are allowed. Try again.")
    try:
        cursor = ctx.deps.conn.execute(sql_query)
        columns = [desc[0] for desc in cursor.description]
        rows = [dict(zip(columns, row)) for row in cursor.fetchall()]
        return f"Query: {sql_query}\nResults ({len(rows)} rows):\n{rows}"
    except sqlite3.Error as e:
        raise ModelRetry(f"SQL error: {e}. Fix the query and try again.")


# 5. Run it
async def main():
    questions = [
        "Who earns the most?",
        "How many engineers do we have?",
        "Show me everyone hired in 2024",
    ]
    deps = Deps(conn=conn)
    for q in questions:
        print(f"\n> {q}")
        result = await agent.run(q, deps=deps)
        print(result.output)


asyncio.run(main())
Enter fullscreen mode Exit fullscreen mode

What's Happening

Lines 1-26: The database. We create an in-memory SQLite database with an employees table and five rows. No server install, no Docker — sqlite3 ships with Python.

Lines 29-35: Dependencies and schema. The Deps dataclass passes the database connection into the agent's tools. QueryResult defines what we expect back, though here we let the agent respond freely since the tool formats the output.

Lines 38-44: The agent. One line creates the agent with a model, dependency type, and system prompt. The system prompt tells the LLM exactly which table and columns exist. This is the key to accurate SQL generation — always give the agent your schema.

Lines 48-59: The tool. The @agent.tool decorator registers run_query as something the agent can call. When the user asks "Who earns the most?", the LLM generates SQL, then calls this tool to execute it. Two safety features are built in:

  • Only SELECT queries run (no drops or deletes)
  • SQL errors trigger ModelRetry, which tells the LLM to fix its query and try again

Lines 62-71: Run it. We loop through three questions. The agent generates SQL, runs it, and returns results.

Expected Output 

> Who earns the most?
Carol in Engineering earns the most at $135,000.

> How many engineers do we have?
There are 3 engineers in the Engineering department.

> Show me everyone hired in 2024
Dave (Sales, hired 2024-01-10) and Eve (Engineering, hired 2024-07-01).
Enter fullscreen mode Exit fullscreen mode

Make It Your Own

Swap the in-memory database for a real one by changing the connection string:

conn = sqlite3.connect("your_database.db")
Enter fullscreen mode Exit fullscreen mode

For PostgreSQL, swap sqlite3 for asyncpg and update the tool's execute call. The agent logic stays identical.

The pattern here — agent + tool + retry — works for any data source. Swap the SQL tool for a REST API call, a CSV reader, or a vector search, and you have a different agent with the same clean architecture.

If you're building AI agents that need to coordinate multiple tools — database queries, API calls, and more — platforms like Nebula can help you orchestrate these workflows without managing infrastructure.

This is part of the AI Agent Quick Tips series. Previously: How to Build Your First MCP Server in 10 Minutes.

Top comments (0)