Skip to content

结构化数据#

LlamaIndex + 结构化数据指南#

许多现代数据系统依赖于结构化数据,比如 Postgres 数据库或 Snowflake 数据仓库。LlamaIndex 提供了许多由 LLM 驱动的高级功能,既可以从非结构化数据中创建结构化数据,也可以通过增强的文本到 SQL 功能来分析这些结构化数据。

注意: 任何文本到 SQL 的应用程序都应该意识到执行任意 SQL 查询可能存在安全风险。建议采取必要的预防措施,比如使用受限角色、只读数据库、沙盒等。

本指南将帮助您了解每一种这些功能。具体来说,我们将涵盖以下主题:

  • 设置:定义我们的示例 SQL 表。
  • 构建我们的表索引:如何从 SQL 数据库转换为表模式索引。
  • 使用自然语言 SQL 查询:如何使用自然语言查询我们的 SQL 数据库。

我们将以包含城市/人口/国家信息的玩具示例表为例进行讲解。本教程的笔记本可在此处找到。

设置#

首先,我们使用 SQLAlchemy 设置一个简单的 sqlite 数据库:

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

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

然后我们创建一个玩具 city_stats 表:

# 创建城市 SQL 表
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 引擎;这样可以让数据库在 LlamaIndex 中使用:

from llama_index.core import SQLDatabase

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

自然语言 SQL#

一旦我们构建了我们的 SQL 数据库,我们就可以使用 NLSQLTableQueryEngine 来构建自然语言查询,这些查询将被合成为 SQL 查询。

请注意,我们需要在此查询引擎中指定要使用的表。如果不这样做,查询引擎将提取所有模式上下文,这可能会超出 LLM 的上下文窗口。

from llama_index.core.query_engine import NLSQLTableQueryEngine

query_engine = NLSQLTableQueryEngine(
    sql_database=sql_database,
    tables=["city_stats"],
)
query_str = "哪个城市的人口最多?"
response = query_engine.query(query_str)

在任何可以预先指定要查询的表或者总表模式大小加上提示的上下文窗口的情况下,应该使用这个查询引擎。

构建我们的表索引#

如果我们事先不知道要使用哪个表,而且表模式的总大小超出了您的上下文窗口大小,那么我们应该将表模式存储在索引中,以便在查询时可以检索到正确的模式。

我们可以使用 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")),
    ...,
]  # 每个表一个 SQLTableSchema

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

在这里,我们定义了我们的 table_node_mapping,以及一个名为 "city_stats" 的单个 SQLTableSchema。我们将这些传递给 ObjectIndex 构造函数,以及我们想要使用的 VectorStoreIndex 类定义。这将为我们提供一个 VectorStoreIndex,其中每个 Node 包含表模式和其他上下文信息。您还可以添加任何其他您想要的上下文信息。

# 手动设置额外的上下文文本
city_stats_text = (
    "这张表格提供了关于特定城市的人口和国家信息。\n"
    "用户将使用代码词进行查询,其中'foo'对应人口,'bar'对应城市。"
)

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

使用自然语言 SQL 查询#

一旦我们定义了表模式索引 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("哪个城市的人口最多?")
print(response)

现在,当我们查询检索器查询引擎时,它将检索相关的表模式,并从该查询的结果中合成一个 SQL 查询和一个响应。

总结思路#

目前就是这样了!我们一直在寻找改进结构化数据支持的方法。如果您有任何问题,请在我们的 Discord 中告诉我们。

相关资源: