跳转到内容

结构化数据

许多现代数据系统依赖于结构化数据,例如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数据库,就可以使用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")),
...,
] # one SQLTableSchema for each table
obj_index = ObjectIndex.from_objects(
table_schema_objs,
table_node_mapping,
VectorStoreIndex,
)

在这里您可以看到我们定义了表节点映射,以及一个名为“city_stats”的SQLTableSchema。我们将这些传入ObjectIndex构造函数,同时传入我们想要使用的VectorStoreIndex类定义。这将为我们提供一个VectorStoreIndex,其中每个节点都包含表模式和其他上下文信息。您还可以添加任何您希望的其他上下文信息。

# 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))
]

一旦我们定义了表模式索引 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中告知我们。