手撸 Text2SQL 应用

端午假期在家无事, 闲着也是闲着, 自己动手写一个 Text2SQL 生成器, 通过 Langchain 实现起来也很容易

只要我们有了 DB Schema, 加上大模型, 写 SQL 不再麻烦, 直接用自然语言描述你的要求, 直接就能生成所需的 SQL 以及相关的 MyBatis 语句.

1. 程序整体架构

这个程序是一个基于RAG技术的Text-to-SQL生成系统，主要包含以下组件：

Text2SQLGenerator
├── 环境配置 (SSL/HTTP设置)
├── 核心组件
│   ├── LLM (ChatOpenAI)
│   ├── 向量数据库 (Chroma)
│   ├── 嵌入模型 (HuggingFaceEmbeddings)
│   └── SQLite对话历史存储
└── 业务流程
    ├── 数据库模式提取
    ├── 向量索引构建
    ├── 检索增强生成
    └── 对话历史管理

2. RAG技术详解

RAG的实现体现在以下关键流程中：

a) 检索阶段 (Retrieval)

def _get_vector_store(self, mysql_url: str):
    # 从MySQL数据库提取模式(Schema)
    schema = self._get_mysql_schema(mysql_url)  

    # 将模式文本分割为文档
    docs = self.text_splitter.create_documents([schema])

    # 构建向量存储
    return Chroma.from_documents(
        docs, 
        self.embedding,
        collection_name="schema",
        persist_directory=self.PERSIST_DIRECTORY
    )

关键技术点： - 使用HuggingFace的bge-small-en-v1.5模型生成嵌入向量 - 将数据库模式(DDL语句)转换为向量表示 - 通过ChromaDB建立可持久化的向量索引

b) 增强阶段 (Augmentation)

# 获取相关数据库模式片段
docs = retriever.get_relevant_documents(request.query)
context = "\n---\n".join(doc.page_content for doc in docs)

# 构建包含以下内容的提示词：
prompt = """
Previous conversation: {history}  
Database Schema: {context}
Question: {question}
"""

增强逻辑： 1. 根据用户 query 检索最相关的数据库模式片段 2. 将检索结果与对话历史一起注入提示词 3. 提供明确的JSON响应格式要求

c) 生成阶段 (Generation)

result_msg = self.llm.invoke(prompt)
result_json = self._extract_json(result_msg.content)

生成特点： - 使用ChatOpenAI作为生成模型 - 强制要求JSON格式输出 - 包含SQL和MyBatis两种输出形式

3. 关键技术组件

a) 分层检索架构

用户问题 → 向量相似度检索 → 数据库模式片段 → 提示词构建 → LLM生成

b) 动态上下文管理

• 短期记忆：通过SQLite保存对话历史
• 长期记忆：向量化的数据库模式知识
• 会话隔离：基于session_id的对话历史检索

c) 异常处理机制

def _extract_json(content: str):
    # 尝试多种JSON提取方式：
    # 1. 直接解析
    # 2. 提取```json代码块
    # 3. 提取{...}模式
    # 确保非结构化响应也能被处理

4. 工作流程示例

对于查询："how many applications are there in the system?"

1. 检索阶段：
2. 从向量库中找到application表的相关模式

3. 例如：CREATE TABLE application (id INT, name VARCHAR(255), status VARCHAR(50))

4. 提示词构建： ```sql Database Schema: CREATE TABLE application ( id INT, name VARCHAR(255), status VARCHAR(50) )

Question: how many applications are there in the system? ```

1. 生成输出： json { "sql": "SELECT COUNT(*) FROM application", "mybatis": "<select id='countApplications' resultType='int'>...</select>" }

5. 小结

1. 知识更新便捷：
2. 只需重新导入数据库模式即可更新知识库

3. 无需重新训练模型

4. 多轮对话支持：

5. 通过session_id维持会话上下文

6. 历史问答自动注入后续提示

7. 混合式知识管理：

8. 结构化知识：数据库模式
9. 非结构化知识：LLM的通用知识

10. 会话知识：对话历史

11. 生产就绪特性：

12. 持久化存储向量索引和对话记录
13. 完善的错误处理和日志记录

这样两百多行代码也就实现用 RAG 将传统数据库知识与大语言模型能力相结合，轻松地将自然语言转换为了 SQL。

# python text2sql_demo_2.py "how many applications in the system"
Results:
SQL: SELECT COUNT(DISTINCT id) FROM application;
MyBatis: <select id='countApplications' resultMap='stringResultMap'>
  SELECT COUNT(DISTINCT id) AS count
  FROM application
</select>

完整代码参见 https://github.com/walterfan/lazy-rabbit-agent/blob/master/example/text2sql_demo_2.py

本作品采用知识共享署名-非商业性使用-禁止演绎 4.0 国际许可协议进行许可。