Text-to-SQL Agent with Pydantic AI

from pydantic_ai import Agent, RunContext
from langchain_community.utilities import SQLDatabase
from langchain_community.agent_toolkits import SQLDatabaseToolkit

agent = Agent(
    "openai:gpt-4o",
    system_prompt="You are a SQL analyst. Convert natural language questions into SQL queries. Only use SELECT statements. Limit results to 10 rows.",
)

@agent.tool
async def list_tables(ctx: RunContext) -> str:
    """List all tables in the database."""
    return "\n".join(inspector.get_table_names())

@agent.tool
async def run_query(ctx: RunContext, sql: str) -> str:
    """Execute a read-only SQL query and return results."""
    if not sql.strip().upper().startswith("SELECT"):
        return "Error: Only SELECT queries are allowed."
    result = db.execute(text(sql))
    rows = result.fetchmany(10)
    return "\n".join(str(row) for row in rows)

result = await agent.run("Show me top 5 customers by total spending this quarter")
print(result.output)