RAG Training SQL Generator —— 打造智能 SQL 生成训练的新生态

在大数据时代，数据库查询与分析成为企业和开发者的重要需求。如何高效而准确地生成 SQL 查询语句，一直是技术创新的焦点。今天，我们隆重推出基于 RAG 思想的【RAG Training SQL Generator】项目，为您提供一套自动化、智能化的 SQL 生成训练流程。无论是 PostgreSQL、SQLite 还是其他关系型数据库，本项目均能灵活适配，满足不同业务场景的需求。

项目简介

【RAG Training SQL Generator】利用最新的 RAG（Retrieval-Augmented Generation）技术，从大量问答数据中自动构造 SQL 查询。项目主要特性包括：

• 自动构造 SQL 查询：通过训练问答- SQL 对，自动生成精准的 SQL 语句；

• 向量训练模型：基于 ChromaDB 存储数据库表结构、文档信息及问答案例，并借助向量匹配实现智能表选择与分组；

• 多智能体协同：提供两种智能体——一种自动选择相关表并生成 SQL，另一种通过表分组提高生成精度；

• 丰富的数据库支持：支持 PostgreSQL、SQLite 等多种数据库配置（基于 SQLAlchemy）；

• 自动训练数据构建：通过 Excel 文件导入问答数据，并自动构建训练数据，为大规模数据库场景保驾护航；

• 可视化流程图：项目支持生成训练流程图，帮助用户直观了解模型工作流程。

核心技术与实现思路

项目的设计充分借鉴了 vanna 和 chat2db 的优秀经验，同时在大规模数据库场景下进行了技术革新，主要有以下两个技术路线：

1. 智能 SQL 生成智能体

• 表匹配与 SQL 生成：首先利用向量匹配技术，从知识库中选取与用户问题最相似的问答案例。如果案例足够匹配，则根据对应表结构和文档信息生成 SQL；如果匹配不足，系统会依托大模型自主定位问题对应的表结构，并借助文档提示生成正确 SQL。

• 辅助文档：针对表结构信息与查询策略，项目专门提供了辅助文档，确保生成结果具有较高的准确性。

2. 表分组与智能筛选

• 表分类机制：在面对上百张表的大型数据库时，项目采用了表分组策略，先通过大模型判断用户问题所属的业务类别，然后精准匹配对应分类下的表结构信息（DDL）。

• 灵活定制：用户可以选择自动表分组，或上传自己定义的业务分类方式，从而更好地适应不同业务场景。

这种多节点、多信息的协同机制，使得项目在大规模数据库下仍能保持优异的 SQL 生成性能。

使用方法与安装指南

直接安装与快速上手

项目已经发布至 PyPI，直接使用 pip 即可安装：

pip install opentext2sql

安装完成后，你只需简单几行代码即可启动服务并访问前端页面：

from opentext2sql.easy_start import start
start()

启动后访问 localhost:3000 即可查看流程图和交互界面，直观了解训练流程。

配置文件示例

项目支持多种数据库配置，例如 PostgreSQL 与 SQLite。以下分别展示两种数据库配置方式的示例：

PostgreSQL 配置：

config = {
“train_data_directory”: “./train_data/“,
“excel_filename”: “train_data.xlsx”,
“db_config”: {
“dialect”: “postgresql”,
“host”: host,
“port”: port,
“database”: database,
“username”: username,
“password”: password
},
“model_name”: model_name,
“openai_api_base”: openai_api_base,
“openai_api_key”: openai_api_key
}

SQLite 配置：

config = {
“train_data_directory”: “./train_data/“,
“excel_filename”: “train_data.xlsx”,
“db_config”: {
“dialect”: “sqlite”,
“database”: “my_database.db”
},
“model_name”: model_name,
“openai_api_base”: openai_api_base,
“openai_api_key”: openai_api_key
}

构建训练数据

利用配置初始化训练模型后，可通过以下流程自动构建训练数据：

1. 获取表结构信息（必须）
   my_train.build_tabel_schema_train_data_from_conn()
2. 自动分组表（为 Text2SqlAgentAutoSelectAspect 提供支持）
   my_train.build_grouped_tables()
3. 从 Excel 加载问答训练数据
   my_train.build_train_data_from_excel(“train_data.xlsx”)
4. 添加 question-SQL 案例
   my_train.build_question_sql_train_data(“question”, “sql”)
5. 添加辅助文档：表与分组类型信息、生成 SQL 的文档说明
   my_train.build_documentation_table_train_data(“doc”)
   my_train.build_documentation_train_data(“doc”)

Excel 文件示例格式：

创建智能体并生成 SQL

项目内置两个智能体，分别负责表选择和表分组。以下为自动选择表的智能体用法示例：

from opentext2sql.agent.create_agent import Text2SqlAgentAutoSelectTable, Text2SqlAgentAutoSelectAspect
from opentext2sql.chroma.chromadb_vector import ChromaDB_VectorStore

初始化配置

config = { … } # 参考上面的配置示例

初始化训练模型并构建训练数据

my_train = ChromaDB_VectorStore(config=config)
my_train.build_tabel_schema_train_data_from_conn()
my_train.build_grouped_tables()
my_train.build_train_data_from_excel(“train_data.xlsx”)
my_train.build_question_sql_train_data(“question”, “sql”)
my_train.build_documentation_table_train_data(“doc”)
my_train.build_documentation_train_data(“doc”)

创建自动选表智能体

agent = Text2SqlAgentAutoSelectTable(
my_train,
use_exmple_question=True,
save_flow_graph=True
)

使用智能体查询和生成 SQL

query = “查询有哪些歌手”
result = agent.ask(query)
print(“查询结果:”, result)

sql_query = agent.generate_sql(input=query)
print(“生成的 SQL:”, sql_query)

项目优势与未来规划

与市面上其他方案对比

• Vanna 路线：传统方案先将所有匹配信息一次性传递给大模型生成 SQL，可能出现匹配不精准的问题；

• Chat2DB 路线：虽然先提取相关表名，但在大规模数据库中依然存在匹配不足的挑战；

本项目通过灵活的信息利用和分布在不同节点的智能体协作，针对大规模数据库场景进行了优化，大幅提高 SQL 生成的准确性和效率。

未来功能规划

• 前端页面制作：为非专业用户提供易用的交互页面，直接连接数据库进行 SQL 查询与数据分析；

• 更新英文提示词：针对 ChatGPT、Claude 等大模型，进一步优化英文提示词，提升跨语言支持能力；

• 增加检查节点：在查询失败时添加检查节点，自动对关键信息进行校验与调整；

• 基于向量数据训练更多模型：设计更多专注于数据库查询的智能体，实现多场景智能应对；

• 扩展数据库连接测试：支持更多类型的数据库，确保兼容性和稳定性；

• 更多向量数据库支持：未来将扩展对其他向量数据库的支持，应对不断变化的数据存储需求。

快速上手指南与社区支持

无论你是希望提升企业级数据查询能力，还是想在开发过程中体验智能 SQL 生成的便捷，【RAG Training SQL Generator】都是一个不可多得的选择。你可以通过以下步骤快速上手：

1. 安装代码

   pip install opentext2sql

2. 克隆仓库获取最新代码

git clone https://github.com/iooo2333/opentext2sql.git
cd opentext2sql
pip install -r requirements.txt

3. 修改配置文件及准备 Excel 数据
   根据实际需求调整 config 配置，并准备符合格式的 train_data.xlsx 文件。

4. 运行示例代码
   启动程序，查看查询结果和生成的 SQL 语句，感受智能化数据查询的魅力。

如果你在使用过程中遇到任何问题，欢迎通过 GitHub Issues 或邮件（1766980815@qq.com）与我们取得联系，共同推动项目的不断完善和进步。

贡献与许可

【RAG Training SQL Generator】项目基于 MIT 许可证开源发布。我们欢迎所有开发者对本项目提出建议、贡献代码或进行改进，欢迎大家在 GitHub 上提交 Issue 或 Pull Request，共同构建智能 SQL 生成的未来。

总的来说，【RAG Training SQL Generator】项目以其先进的 RAG 思想、灵活的信息利用方式以及针对大规模数据库的优化能力，为 SQL 查询自动化提供了一套高效、精准的解决方案。如果你正在寻找一款能够大幅提升数据分析和 SQL 生成效率的工具，不妨立即尝试这一项目，并与我们共同探讨、改进，共同推动智能数据库应用的发展！

立即加入 GitHub 体验：github.com/iooo2333/opentext2sql.g...

Happy Coding!

本作品采用《CC 协议》，转载必须注明作者和本文链接