如何基于LlamaIndex与SQLAlchemy为查询引擎设置字段元数据

核心思路 #

将定义好的元数据字典转换为JSON字符串，关联到snapshots表的data字段注释中，再通过SQLTableRetrieverQueryEngine的构建流程，让智能体在查询时读取这些元数据，理解字段的业务含义，进而优化查询逻辑。

步骤详解 #

1. 定义元数据字典
   先明确需要提供给智能体的字段业务说明，示例如下：

metadata_dict = {
    "page_call_phone_clicks_logged_in_unique": "Number of people who logged into Facebook and clicked the Call Now button. Ex: 2"
}

2. 创建绑定元数据的snapshots表
   将元数据字典序列化为JSON字符串，作为snapshots表data字段的注释，代码实现：

from sqlalchemy import create_engine, MetaData, Table, Column, JSON
import json

# 初始化数据库连接与元数据对象
engine = create_engine(pg_uri)
metadata_obj = MetaData()

# 转换元数据字典为JSON字符串
metadata_json = json.dumps(metadata_dict)

# 定义snapshots表，将元数据绑定到data字段的comment属性
table_name = "snapshots"
snapshots_table = Table(
    table_name,
    metadata_obj,
    Column("data", JSON, comment=metadata_json),
)

# 反射数据库现有表结构
metadata_obj.reflect(engine)

3. 构建支持元数据的SQL查询引擎
   整合数据库、元数据索引与大模型，让智能体可利用元数据优化查询：

from llama_index import SQLDatabase, SQLTableNodeMapping, SQLTableSchema
from llama_index.indices.struct_store import ObjectIndex
from llama_index.indices.vector_store import VectorStoreIndex
from llama_index.llms import OpenAI
from llama_index.llms.predictor import LLMPredictor
from llama_index import ServiceContext
from llama_index.query_engine import SQLTableRetrieverQueryEngine

# 初始化SQL数据库对象
sql_database = SQLDatabase(engine)

# 建立表与节点的映射关系
nodes = SQLTableNodeMapping(sql_database)

# 收集所有表的schema信息
schema = []
for table_name in metadata_obj.tables.keys():
    schema.append(SQLTableSchema(table_name=table_name))

# 将表schema存入向量索引，用于查询时动态检索关联表
index = ObjectIndex.from_objects(schema, nodes, VectorStoreIndex)

# 配置大模型与服务上下文
llm = OpenAI(temperature=0.9, model_name="gpt-3.5-turbo")
llm_predictor = LLMPredictor(llm=llm)
service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor)

# 构建带元数据支持的查询引擎
query_engine = SQLTableRetrieverQueryEngine(
    sql_database,
    index.as_retriever(similarity_top_k=1),
    service_context=service_context,
)

关键说明 #

• 元数据通过字段的comment属性存储，SQLTableRetrieverQueryEngine在检索表schema时会读取这些注释，帮助大模型理解字段的业务意义。
• 向量索引的作用是在查询时快速定位到相关表，结合元数据生成更贴合需求的SQL语句。

内容的提问来源于stack exchange，提问作者teddybear