结构化数据

LlamaIndex + 结构化数据指南

现代数据系统大量依赖结构化数据，例如Postgres数据库或Snowflake数据仓库。LlamaIndex提供了许多由大语言模型驱动的高级功能，既能从非结构化数据生成结构化数据，也能通过增强的文本转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表：

# 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引擎，使其可在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 = "Which city has the highest population?"
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("Which city has the highest population?")
print(response)

现在查询检索器查询引擎时，它将检索相关表模式，合成SQL查询并从查询结果生成响应。

结语

目前内容就这些！我们持续改进对结构化数据的支持。如有问题，请通过Discord联系我们。

相关资源：

• Airbyte SQL索引指南