Agent for text-to-SQL with automatic error correction

In this tutorial, we’ll see how to implement an agent that leverages SQL using transformers.agents.

What’s the advantage over a standard text-to-SQL pipeline?

A standard text-to-sql pipeline is brittle, since the generated SQL query can be incorrect. Even worse, the query could be incorrect, but not raise an error, instead giving some incorrect/useless outputs without raising an alarm.

👉 Instead, an agent system is able to critically inspect outputs and decide if the query needs to be changed or not, thus giving it a huge performance boost.

Let’s build this agent! 💪

Setup SQL tables

from sqlalchemy import (
    create_engine,
    MetaData,
    Table,
    Column,
    String,
    Integer,
    Float,
    insert,
    inspect,
    text,
)

engine = create_engine("sqlite:///:memory:")
metadata_obj = MetaData()

# create city SQL table
table_name = "receipts"
receipts = Table(
    table_name,
    metadata_obj,
    Column("receipt_id", Integer, primary_key=True),
    Column("customer_name", String(16), primary_key=True),
    Column("price", Float),
    Column("tip", Float),
)
metadata_obj.create_all(engine)

rows = [
    {"receipt_id": 1, "customer_name": "Alan Payne", "price": 12.06, "tip": 1.20},
    {"receipt_id": 2, "customer_name": "Alex Mason", "price": 23.86, "tip": 0.24},
    {"receipt_id": 3, "customer_name": "Woodrow Wilson", "price": 53.43, "tip": 5.43},
    {"receipt_id": 4, "customer_name": "Margaret James", "price": 21.11, "tip": 1.00},
]
for row in rows:
    stmt = insert(receipts).values(**row)
    with engine.begin() as connection:
        cursor = connection.execute(stmt)

Let’s check that our system works with a basic query:

>>> with engine.connect() as con:
...     rows = con.execute(text("""SELECT * from receipts"""))
...     for row in rows:
...         print(row)

(1, 'Alan Payne', 12.06, 1.2)
(2, 'Alex Mason', 23.86, 0.24)
(3, 'Woodrow Wilson', 53.43, 5.43)
(4, 'Margaret James', 21.11, 1.0)

Build our agent

Now let’s make our SQL table retrievable by a tool.

The tool’s description attribute will be embedded in the LLM’s prompt by the agent system: it gives the LLM information about how to use the tool. So that is where we want to describe the SQL table.

>>> inspector = inspect(engine)
>>> columns_info = [(col["name"], col["type"]) for col in inspector.get_columns("receipts")]

>>> table_description = "Columns:\n" + "\n".join([f"  - {name}: {col_type}" for name, col_type in columns_info])
>>> print(table_description)

Columns:
  - receipt_id: INTEGER
  - customer_name: VARCHAR(16)
  - price: FLOAT
  - tip: FLOAT

Now the tool will need a few mandatory arguments (read the documentation for more detail):

The description mentioned above
A name by which the agent system can call it.
A dictionary of inputs: describe the tool’s input, each with a type and a description key.
An output_type for compatibility with gradio demos for instance.
A forward method that gets called when the tool is called.

from transformers.agents import Tool


class SQLExecutorTool(Tool):
    name = "sql_engine"
    description = f"""Allows you to perform SQL queries on the table. Returns a string representation of the result.
The table is named 'receipts'. Its description is as follows: \n{table_description}"""
    inputs = {
        "query": {
            "type": "text",
            "description": f"The query to perform. This should be correct SQL.",
        }
    }
    output_type = "text"

    def forward(self, query: str) -> str:
        output = ""
        with engine.connect() as con:
            rows = con.execute(text(query))
            for row in rows:
                output += "\n" + str(row)
        return output

Now let us create an agent that leverages this tool.

We use the ReactCodeAgent, which is transformers.agents’ main agent class: an agent that writes actions in code and can iterate on previous output according to the ReAct framework.

The llm_engine is the LLM that powers the agent system. HfEngine allows you to call LLMs using HF’s Inference API, either via Serverless or Dedicated endpoint, but you could also use any proprietary API: check out this other cookbook to learn how to adapt it.

from transformers.agents import ReactCodeAgent, HfEngine

agent = ReactCodeAgent(
    tools=[SQLExecutorTool()],
    llm_engine=HfEngine("meta-llama/Meta-Llama-3-70B-Instruct"),
)

agent.run("Can you give me the name of the client who got the most expensive receipt?")

Increasing difficulty: Table joins

Now let’s make it more challenging! We want our agent to handle joins across multiple tables.

So let’s make a second table recording the names of waiters for each receipt_id!

table_name = "waiters"
receipts = Table(
    table_name,
    metadata_obj,
    Column("receipt_id", Integer, primary_key=True),
    Column("waiter_name", String(16), primary_key=True),
)
metadata_obj.create_all(engine)

rows = [
    {"receipt_id": 1, "waiter_name": "Corey Johnson"},
    {"receipt_id": 2, "waiter_name": "Michael Watts"},
    {"receipt_id": 3, "waiter_name": "Michael Watts"},
    {"receipt_id": 4, "waiter_name": "Margaret James"},
]
for row in rows:
    stmt = insert(receipts).values(**row)
    with engine.begin() as connection:
        cursor = connection.execute(stmt)

We need to update the SQLExecutorTool with this table’s description to let the LLM properly leverage information from this table.

>>> updated_description = """Allows you to perform SQL queries on the table. Beware that this tool's output is a string representation of the execution output.
... It can use the following tables:"""

>>> inspector = inspect(engine)
>>> for table in ["receipts", "waiters"]:
...     columns_info = [(col["name"], col["type"]) for col in inspector.get_columns(table)]

...     table_description = f"Table '{table}':\n"

...     table_description += "Columns:\n" + "\n".join([f"  - {name}: {col_type}" for name, col_type in columns_info])
...     updated_description += "\n\n" + table_description

>>> print(updated_description)

Allows you to perform SQL queries on the table. Beware that this tool's output is a string representation of the execution output.
It can use the following tables:

Table 'receipts':
Columns:
  - receipt_id: INTEGER
  - customer_name: VARCHAR(16)
  - price: FLOAT
  - tip: FLOAT

Table 'waiters':
Columns:
  - receipt_id: INTEGER
  - waiter_name: VARCHAR(16)

>>> sql_tool = SQLExecutorTool()
>>> sql_tool.description = updated_description

>>> agent = ReactCodeAgent(
...     tools=[sql_tool],
...     llm_engine=HfEngine("meta-llama/Meta-Llama-3-70B-Instruct"),
... )

>>> agent.run("Which waiter got more total money from tips?")

['time', 'queue', 'random', 'unicodedata', 're', 'collections', 'itertools', 'stat', 'math', 'statistics']

It directly works! The setup was surprisingly simple, wasn’t it?

✅ Now you can go build this text-to-SQL system you’ve always dreamt of! ✨

< > Update on GitHub

Open-Source AI Cookbook

Agent for text-to-SQL with automatic error correction

Setup SQL tables

Build our agent

Increasing difficulty: Table joins