Text-to-SQL

在本教程中，我们将了解如何使用 smolagents 实现一个利用 SQL 的 agent。

让我们从一个关键问题开始：为什么不保持简单，使用标准的 text-to-SQL 流程呢？

标准的 text-to-sql 流程是脆弱的，因为生成的 SQL 查询可能不正确。更糟糕的是，查询可能不正确，但不会引发错误，而是给出一些不正确/无用的输出，而不会发出警报。

👉 相反，agent 系统能够批判性地检查输出，并决定是否需要更改查询，从而大大提高性能。

让我们构建这个 agent！ 💪

运行下面这行代码来安装所需的依赖

!pip install smolagents python-dotenv sqlalchemy --upgrade -q

要调用 HF Inference API，您需要一个有效的 token 作为您的环境变量 HF_TOKEN。我们使用 python-dotenv 来加载它。

from dotenv import load_dotenv
load_dotenv()

然后，我们设置 SQL 环境

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

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

def insert_rows_into_table(rows, table, engine=engine):
    for row in rows:
        stmt = insert(table).values(**row)
        with engine.begin() as connection:
            connection.execute(stmt)

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},
]
insert_rows_into_table(rows, receipts)

构建我们的 agent

现在让我们使我们的 SQL 表可以通过工具检索。

工具的 description 属性将由 agent 系统嵌入到 LLM 的 prompt 中：它为 LLM 提供了关于如何使用工具的信息。这是我们想要描述 SQL 表的地方。

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

现在让我们构建我们的工具。它需要以下内容：（阅读工具文档以了解更多详情）

一个带有 Args: 部分列出参数的 docstring。
输入和输出的类型提示。

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

现在让我们创建一个利用此工具的 agent。

我们使用 CodeAgent，它是 smolagents 的主要 agent 类：一个以代码形式编写操作的 agent，可以根据 ReAct 框架迭代先前的输出。

模型是驱动 agent 系统的 LLM。HfApiModel 允许您使用 HF 的 Inference API 调用 LLM，可以通过 Serverless 或 Dedicated endpoint，但您也可以使用任何专有 API。

from smolagents import CodeAgent, HfApiModel

agent = CodeAgent(
    tools=[sql_engine],
    model=HfApiModel(model_id="meta-llama/Meta-Llama-3.1-8B-Instruct"),
)
agent.run("Can you give me the name of the client who got the most expensive receipt?")

第二级：表连接

现在让我们让它更具挑战性！我们希望我们的 agent 处理跨多个表的连接。

因此，让我们创建第二个表来记录每个 receipt_id 的服务员姓名！

table_name = "waiters"
waiters = 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"},
]
insert_rows_into_table(rows, waiters)

由于我们更改了表，我们使用此表的描述更新 SQLExecutorTool，以便 LLM 能够正确利用来自此表的信息。

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)

由于此请求比之前的请求更难一些，我们将切换 LLM 引擎以使用更强大的 Qwen/Qwen2.5-Coder-32B-Instruct！

sql_engine.description = updated_description

agent = CodeAgent(
    tools=[sql_engine],
    model=HfApiModel(model_id="Qwen/Qwen2.5-Coder-32B-Instruct"),
)

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

它可以直接工作！设置出奇的简单，不是吗？

这个例子完成了！我们已经接触了以下概念

构建新工具。
更新工具的描述。
切换到更强大的 LLM 有助于 agent 推理。

✅ 现在您可以去构建您一直梦想的 text-to-SQL 系统了！ ✨

< > 在 GitHub 上更新