LlamaIndex + 结构化数据指南

许多现代数据系统依赖结构化数据，比如 Postgres 数据库或 Snowflake 数据仓库。LlamaIndex 提供了大量由 LLM 驱动的进阶能力：既能从非结构化数据中生成结构化数据，也能借助增强的 text-to-SQL 能力来分析这些结构化数据。

**注意：**任何 Text-to-SQL 应用都应意识到，执行任意 SQL 查询存在安全风险。建议按需采取防护措施，例如使用受限角色、只读数据库、沙箱隔离等。

本指南会逐一讲解这些能力。具体涵盖以下主题：

• 准备：搭好我们的示例 SQL 表。
• 构建 Table Index：如何从 SQL 数据库走到 Table Schema Index。
• 使用自然语言 SQL 查询：如何用自然语言查询我们的 SQL 数据库。

我们会走一遍一个包含城市/人口/国家信息的玩具示例表。本教程的 notebook 在这里。

设置

首先用 SQLAlchemy 建一个简单的 sqlite 数据库：

from sqlalchemy import (
    create_engine,
    MetaData,
    Table,
    Column,
    String,
    Integer,
    select,
    column,
)

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

然后创建一张玩具 city_stats 表：

# create city SQL table
table_name = "city_stats"
city_stats_table = Table(
    table_name,
    metadata_obj,
    Column("city_name", String(16), primary_key=True),
    Column("population", Integer),
    Column("country", String(16), nullable=False),
)
metadata_obj.create_all(engine)

接下来插入一些数据点！

如果你想试试从非结构化数据中推断出结构化数据点来填充这张表，可以往下看相应章节；否则你也可以直接向这张表写入数据：

from sqlalchemy import insert

rows = [
    {"city_name": "Toronto", "population": 2731571, "country": "Canada"},
    {"city_name": "Tokyo", "population": 13929286, "country": "Japan"},
    {"city_name": "Berlin", "population": 600000, "country": "Germany"},
]
for row in rows:
    stmt = insert(city_stats_table).values(**row)
    with engine.begin() as connection:
        cursor = connection.execute(stmt)

最后，用我们的 SQLDatabase 包装器把 SQLAlchemy engine 包起来；这样这个 db 就能在 LlamaIndex 里使用了：

from llama_index.core import SQLDatabase

sql_database = SQLDatabase(engine, include_tables=["city_stats"])

自然语言 SQL

SQL 数据库建好之后，就可以使用 NLSQLTableQueryEngine 来构建自然语言查询，并由它合成出 SQL 查询。

注意我们必须指定这个查询引擎要使用哪些表。若不指定，查询引擎会把全部 schema 上下文拉进来，可能撑爆 LLM 的上下文窗口。

from llama_index.core.query_engine import NLSQLTableQueryEngine

query_engine = NLSQLTableQueryEngine(
    sql_database=sql_database,
    tables=["city_stats"],
)
query_str = "Which city has the highest population?"
response = query_engine.query(query_str)

凡是你能够事先确定要查询哪些表的场合，或者「全部表 schema 加上提示词其余部分」的总体积能装进上下文窗口的场合，都可以用这个查询引擎。

构建我们的 Table Index

如果我们事先不知道要用哪张表，且表 schema 的总体积会超出上下文窗口，那就应该把表 schema 存进一个索引，以便查询时取回正确的那份 schema。

做法是使用 SQLTableNodeMapping 对象：它接收一个 SQLDatabase，并为传入 ObjectIndex 构造函数的每个 SQLTableSchema 对象生成一个 Node。

from llama_index.core.objects import (
    SQLTableNodeMapping,
    ObjectIndex,
    SQLTableSchema,
)

table_node_mapping = SQLTableNodeMapping(sql_database)
table_schema_objs = [
    (SQLTableSchema(table_name="city_stats")),
    ...,
]  # one SQLTableSchema for each table

obj_index = ObjectIndex.from_objects(
    table_schema_objs,
    table_node_mapping,
    VectorStoreIndex,
)

这里你可以看到我们定义了 table_node_mapping，以及一个表名为 “city_stats” 的 SQLTableSchema。把它们连同我们打算使用的 VectorStoreIndex 类定义一起传给 ObjectIndex 构造函数。这样我们就得到一个 VectorStoreIndex，其中每个 Node 都包含表 schema 及其他上下文信息。你还可以按需添加更多上下文信息。

# manually set extra context text
city_stats_text = (
    "This table gives information regarding the population and country of a given city.\n"
    "The user will query with codewords, where 'foo' corresponds to population and 'bar'"
    "corresponds to city."
)

table_node_mapping = SQLTableNodeMapping(sql_database)
table_schema_objs = [
    (SQLTableSchema(table_name="city_stats", context_str=city_stats_text))
]

使用自然语言 SQL 查询

表 schema 索引 obj_index 定义好之后，就可以通过传入我们的 SQLDatabase、以及从对象索引构建的检索器，来创建 SQLTableRetrieverQueryEngine。

from llama_index.core.indices.struct_store import SQLTableRetrieverQueryEngine

query_engine = SQLTableRetrieverQueryEngine(
    sql_database, obj_index.as_retriever(similarity_top_k=1)
)
response = query_engine.query("Which city has the highest population?")
print(response)

现在我们向这个检索器查询引擎提问时，它会先取回相关的表 schema，然后合成出一条 SQL 查询，并根据查询结果给出回答。

小结

目前就是这些！我们一直在寻找改进结构化数据支持的方式。有任何问题，欢迎到我们的 Discord 里告诉我们。