Buckets:

hf-doc-build/doc-dev / cookbook /pr_362 /en /agent_text_to_sql.md
|
download
raw
6.6 kB
# Agent for text-to-SQL with automatic error correction
_Authored by: [Aymeric Roucher](https://huggingface.co/m-ric)_
In this tutorial, we'll see how to implement an agent that leverages SQL using `smolagents`.
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
```python
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)
```
```python
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:
```python
>>> 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.
Our `sql_engine` tool needs the following: (read [the documentation](https://huggingface.co/docs/transformers/en/agents#create-a-new-tool) for more detail)
- A docstring with an `Args:` part. This docstring will be parsed to become the tool's `description` attribute, which will be used as the instruction manual for the LLM powering the agent, so it's important to provide it!
- Type hints for inputs and output.
```python
from smolagents import tool
@tool
def sql_engine(query: str) -> str:
"""
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:
Columns:
- receipt_id: INTEGER
- customer_name: VARCHAR(16)
- price: FLOAT
- tip: FLOAT
Args:
query: The query to perform. This should be correct SQL.
"""
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 `CodeAgent`, 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. `InferenceClientModel` allows you to call LLMs using Hugging Face's Inference API, either via Serverless or Dedicated endpoint, but you could also use any proprietary API: check out [this other cookbook](agent_change_llm) to learn how to adapt it.
```python
from smolagents import CodeAgent, InferenceClientModel
agent = CodeAgent(
tools=[sql_engine],
model=InferenceClientModel("meta-llama/Meta-Llama-3-8B-Instruct"),
)
```
```python
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`!
```python
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.
```python
>>> 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)
Since this request is a bit harder than the previous one, we'll switch the llm engine to use the more powerful [Qwen/Qwen2.5-72B-Instruct](https://huggingface.co/Qwen/Qwen2.5-72B-Instruct)!
```python
sql_engine.description = updated_description
agent = CodeAgent(
tools=[sql_engine],
model=InferenceClientModel("Qwen/Qwen2.5-72B-Instruct"),
)
agent.run("Which waiter got more total money from tips?")
```
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! ✨

Xet Storage Details

Size:
6.6 kB
·
Xet hash:
cd4d444a58c7861d1f7eb8ccc7d60637d7156b89ec09dd68ebc387b009293523

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.