Skip to content

Text2SQL:结构化数据检索

在 RAG 系统中,除了非结构化文本(PDF、Word),我们经常还需要查询结构化数据(SQL 数据库)。 Text2SQL(Text-to-SQL)技术旨在将用户的自然语言问题转化为 SQL 查询语句。

1. 核心原理

Text2SQL 的本质是一个翻译任务: Input (Question + Schema) -> LLM -> Output (SQL)

关键步骤:

  1. Schema Linking: 将用户问题中的实体映射到数据库的表名和列名。
  2. SQL Generation: 利用 LLM 生成 SQL。
  3. Execution: 在数据库执行 SQL。
  4. Response Synthesis: 将查询结果转化为自然语言回答。

2. Golang 实现 Text2SQL

我们可以使用 langchaingo 的 Chain 或者是原生的 text/template 来实现。

2.1 定义数据库 Schema

为了让 LLM 写出正确的 SQL,我们需要把数据库结构(DDL)喂给它。

go
const dbSchema = `
CREATE TABLE users (
    id INT PRIMARY KEY,
    name VARCHAR(50),
    age INT,
    city VARCHAR(50)
);

CREATE TABLE orders (
    id INT PRIMARY KEY,
    user_id INT,
    amount DECIMAL(10, 2),
    created_at DATETIME
);
`

2.2 构建 Prompt

go
const text2sqlTemplate = `
你是一个 SQL 专家。请根据以下数据库 Schema,将用户的自然语言问题转化为可执行的 SQL 查询语句。
只返回 SQL,不要包含任何解释或 Markdown 格式。

【Schema】:
{{ .Schema }}

【用户问题】: {{ .Query }}

【SQL】:
`

2.3 完整代码示例

go
package main

import (
	"bytes"
	"context"
	"database/sql"
	"fmt"
	"log"
	"strings"
	"text/template"

	_ "github.com/mattn/go-sqlite3" // 引入 SQLite 驱动
	"github.com/tmc/langchaingo/llms"
	"github.com/tmc/langchaingo/llms/openai"
)

func main() {
	// 1. 初始化 LLM
	ctx := context.Background()
	llm, err := openai.New()
	if err != nil {
		log.Fatal(err)
	}

	// 2. 准备数据 (模拟)
	db, _ := sql.Open("sqlite3", ":memory:")
	defer db.Close()
	initDB(db)

	// 3. 用户提问
	query := "查询来自 'Beijing' 的用户总共有多少个?"

	// 4. 生成 SQL
	generatedSQL, err := generateSQL(ctx, llm, query)
	if err != nil {
		log.Fatal(err)
	}
	fmt.Printf("生成的 SQL: %s\n", generatedSQL)

	// 5. 执行 SQL (注意:生产环境必须做只读权限控制!)
	rows, err := db.Query(generatedSQL)
	if err != nil {
		log.Fatal("SQL 执行失败:", err)
	}
	defer rows.Close()

	// 6. 打印结果
	for rows.Next() {
		var count int
		rows.Scan(&count)
		fmt.Printf("查询结果: %d\n", count)
	}
}

func generateSQL(ctx context.Context, llm llms.Model, query string) (string, error) {
	tmpl, _ := template.New("sql").Parse(text2sqlTemplate)
	var buf bytes.Buffer
	tmpl.Execute(&buf, map[string]string{
		"Schema": dbSchema,
		"Query":  query,
	})

	// 调用 LLM
	resp, err := llm.Call(ctx, buf.String())
	if err != nil {
		return "", err
	}
	
	// 清理可能的 Markdown 符号 (```sql ... ```)
	cleanSQL := strings.TrimSpace(resp)
	cleanSQL = strings.TrimPrefix(cleanSQL, "```sql")
	cleanSQL = strings.TrimPrefix(cleanSQL, "```")
	cleanSQL = strings.TrimSuffix(cleanSQL, "```")
	
	return cleanSQL, nil
}

// 初始化模拟数据
func initDB(db *sql.DB) {
	db.Exec("CREATE TABLE users (id INT, name TEXT, city TEXT)")
	db.Exec("INSERT INTO users VALUES (1, 'Alice', 'Beijing'), (2, 'Bob', 'Shanghai'), (3, 'Charlie', 'Beijing')")
}

const dbSchema = `
CREATE TABLE users (
    id INT,
    name TEXT, 
    city TEXT
);
`

const text2sqlTemplate = `
你是一个 SQL 专家。请根据以下数据库 Schema,将用户的自然语言问题转化为可执行的 SQLite 查询语句。
只返回 SQL,不要包含任何解释。

【Schema】:
{{ .Schema }}

【用户问题】: {{ .Query }}

【SQL】:
`

3. 进阶技巧

  1. Few-Shot Prompting: 在 Prompt 中提供几个 (Question, SQL) 的示例,能显著提高准确率。
  2. Schema Pruning: 如果数据库表非常多(如 100 张表),不能把所有 DDL 都塞进 Prompt。需要先用 Embedding 检索出相关的表,再生成 SQL。
  3. Self-Correction: 如果执行 SQL 报错(如“列不存在”),将错误信息喂回给 LLM,让它修正 SQL。

总结

Text2SQL 打通了 LLM 与 传统关系型数据库 的界限,是 RAG 系统处理统计类、报表类问题的最佳方案。

🚀 学习遇到瓶颈?想进大厂?

看完这篇技术文章,如果还是觉得不够系统,或者想在实战中快速提升?
王中阳的就业陪跑训练营,提供定制化学习路线 + 企业级实战项目 + 简历优化 + 模拟面试。

了解训练营详情