使用LLM实现自然语言的SQL查询代理

1. 背景

相信大部分研发同学都遇到过这样的场景：业务已经上线了一段时间，产品或运营同学想查看一下实时的业务数据。虽然对于成熟的系统来说，大多有完善的数据分析平台或运营管理后台，可以在各个维度展示业务数据，但是产品和运营同学的查询需求可能是频繁变更的，如“帮我查一下粉丝数>100的用户有哪些？”、“最近1小时内点赞数超过100的视频有多少条？”…。针对这种场景，现有的平台可能无法完全覆盖多样化的查询需求，所以通常需要研发同学手动跑 SQL，这也是很多后端同学经常自嘲为 “SQL Boy” 的原因。这种手动跑 SQL 的流程大致如下：

1. 产品/业务同学口头提出需求
2. 研发同学将需求翻译成 SQL 语句
3. 研发同学执行 SQL，获取查询结果
4. 研发同学将查询结果转化成口头描述，或者生成结构化的格式（json、csv 等），发送给产品/业务
5. 如果数据结果或格式不满足需求，可能还需要重复上述过程，直到获取满意的结果

可以看到，整个过程是是否繁琐的，会耗费研发同学的大量人力。那么，既然现在大语言模型的能力已经比较成熟，我们是否可以利用 LLM 来开发一些工具，可利用 LLM 的推理能力，将产品/业务同学的查询需求转化成标准的 SQL 语句，并自动执行获取结果呢？答案当然是肯定的，利用 LLM 的 tool calling 能力就可以轻松完成这个功能。下面我们就来实现一下吧！

2. 具体实现

2.1 数据准备

为了演示功能，我们首先需要准备一个 SQL 数据库和相关数据。方便期间，我们采用 SQLite 和它官方提供的示例数据库来完成我们的需求。

SQLite 是一款支持 SQL92 标准、高性能、轻量级的嵌入式数据库，它在使用前不需要安装设置，不需要进程来启动、停止或配置，只需要嵌入到应用程序中，基于一个本地的数据库文件，即可完成数据的 CRUD 操作。

SQLite 官方提供了一个叫做 Chinook 的示例数据库，主要用于面向音乐领域的相关数据，包括专辑、音频、媒体类型、播放列表、歌手等等。

我们可以在 https://www.sqlitetutorial.net/sqlite-sample-database/ 这里下载这个数据库文件 chinook.db，即可在本地操作 SQLite。

2.2 加载数据库 Schema

有了数据库文件，我们就可以在应用程序中操作数据库。这里我们使用 Python 语言开发，采用官方提供的 sqlite3 包即可完成。具体操作数据库的函数如下：

def connect_db() -> Connection:
    """创建sqlite数据库连接"""

    # 找到本地的chinook.db文件,创建数据库连接
    conn = sqlite3.connect("./chinook.db")
    print("connect to sqlite success!")
    return conn


def get_table_names(conn: Connection) -> List[str]:
    """返回一个包含所有表名的列表"""

    # 执行SQL查询,获取数据库中所有表的名字
    tables = conn.execute("SELECT name FROM sqlite_master WHERE type='table';")
    # 遍历查询table,返回table_name列表
    return [table[0] for table in tables]


def get_column_names(conn: Connection, table_name: str):
    """返回一个给定表的所有列名的列表"""

    # 执行SQL查询，获取表的所有列的信息
    columns = conn.execute(f"PRAGMA table_info('{table_name}');").fetchall()
    # 遍历columns,返回列名列表
    return [col[1] for col in columns]


def get_database_schema(conn: Connection) -> str:
    """获取数据库的Schema信息,包括表名及其包含的列表"""
    
    table_dicts = []  # 创建一个空的字典列表
    # 遍历数据库中的所有表
    for table_name in get_table_names(conn):
        columns_names = get_column_names(conn, table_name)  # 获取当前表的所有列名
        # 将表名和列名信息作为一个字典添加到列表中
        table_dicts.append({"table_name": table_name, "column_names": columns_names})

    # 格式化字典,转换成字符串返回
    database_schema_string = "\n".join(
        [
            f"Table: {table['table_name']}\nColumns: {', '.join(table['column_names'])}"
            for table in table_dicts
        ]
    )
    return database_schema_string
  
  def exec_sql(conn: Connection, sql: str) -> str:
    """执行SQL,获取结果"""
    
    try:
        # 执行SQL,并将结果转换为字符串
        results = str(conn.execute(sql).fetchall())
    except Exception as e:
        # 执行失败,捕获异常并返回错误信息
        results = f"query failed with error: {e}"

    # 返回查询结果
    return results

首先通过 connect_db 函数连接 SQLite 数据库，之后最关键的方法就是 get_database_schema。因为我们希望大模型能理解人类的提问，并转换成 SQL 语句生成，前提就是需要大模型理解整个数据库的表结构，这样才能将自然语言与数据库的表名和列表做好映射。因此，get_database_schema 这个函数的作用就是获取整个数据库的表结构信息字符串。打印结果为：

Table: albums
Columns: AlbumId, Title, ArtistId
Table: sqlite_sequence
Columns: name, seq
Table: artists
Columns: ArtistId, Name
Table: customers
Columns: CustomerId, FirstName, LastName, Company, Address, City, State, Country, PostalCode, Phone, Fax, Email, SupportRepId
Table: employees
Columns: EmployeeId, LastName, FirstName, Title, ReportsTo, BirthDate, HireDate, Address, City, State, Country, PostalCode, Phone, Fax, Email
Table: genres
Columns: GenreId, Name
Table: invoices
Columns: InvoiceId, CustomerId, InvoiceDate, BillingAddress, BillingCity, BillingState, BillingCountry, BillingPostalCode, Total
Table: invoice_items
Columns: InvoiceLineId, InvoiceId, TrackId, UnitPrice, Quantity
Table: media_types
Columns: MediaTypeId, Name
Table: playlists
Columns: PlaylistId, Name
Table: playlist_track
Columns: PlaylistId, TrackId
Table: tracks
Columns: TrackId, Name, AlbumId, MediaTypeId, GenreId, Composer, Milliseconds, Bytes, UnitPrice
Table: sqlite_stat1
Columns: tbl, idx, stat

2.3 利用 LLM 生成 SQL

有了数据库的表结构信息，我们就可以利用 LLM 的 tool calling 工具调用能力，将用户的自然语言提问转化成标准 SQL 生成。

Tool Calling 工具调用是大模型相对较新的一种扩展能力，它支持在调用 API 时，传入指定的工具列表，并描述每个工具的具体功能、参数格式。大模型可以根据用户的输入，推理出是否需要调用某个或某些工具。如果需要调用工具，大模型会生成一个结构化的结果，包括需要调用的工具名称以及具体的调用参数，并且输出的格式为完全匹配用户的工具参数定义。这样一来，用户获取到生成的工具调用参数，就可以实际发起调用，并进行自定义的扩展实现。工具调用的一大优势是可以使得 LLM 生成完全结构化的输出，便于应用做自定义的扩展。同时，工具调研也是实现 LLM Agent 的基础。

OpenAI 的 Tool Calling 可以参考文档：https://platform.openai.com/docs/api-reference/chat/create ，通过 tools 参数即可传入预定义的工具列表。

其他大模型的工具调用接口定义基本与 OpenAI 类型。

在我们的项目里，采用智谱AI的 GLM-4 模型来实现。智谱AI是国内领先的大模型服务提供商，而且相较于 GPT 模型，GLM 对于中文有更好地支持，而且因为涉及到业务数据的传递，采用国产的大模型也更满足数据安全的需求。

我们实现函数 generate_sql ，利用 LLM 的 tool calling 能力，将用户原始的自然语言提问，转化成结构化的标准 SQL。具体函数实现如下：

def generate_sql(db_schema: str, query: str) -> str:
    """
    调用LLM,利用工具调用能力,生成SQL语句
    :param db_schema: 数据库表结构信息
    :param query: 用户的原始提问
    :return: 生成的结构化SQL
    """

    # 加载环境变量
    dotenv.load_dotenv()

    # 创建智谱AI客户端
    client = ZhipuAI(api_key=os.getenv("ZHIPUAI_API_KEY"))

    # 定义工具的详细描述,便于LLM理解用户的需求
    tool_desc = f"""根据用户提问，生成的SQL语句，用于回答用户的问题。
                生成的SQL语句基于如下的数据库表结构定义：
                {db_schema}
                最终的SQL语句以纯文本的格式输出，不要使用json或者其它的结构化格式。
                """

    # 定义工具
    tools = [
        {
            "type": "function",  # 工具类型为function函数调用
            "function": {  # 函数定义
                "name": "generate_sql",  # 函数名称
                "description": "该函数用于回答用户提出的关于音乐的相关问题。 "
                               "生成的结果是结构化的标准SQL语句。",
                # 函数描述
                "parameters": {  # 函数参数定义
                    "type": "object",
                    "properties": {
                        "sql": {  # 参数名称
                            "type": "string",  # 参数类型
                            "description": tool_desc,  # 参数描述
                        },
                    },
                    "required": ["sql"],  # 必需的参数
                },
            }
        }
    ]

    # 创建消息列表
    messages = [
        {"role": "system",
         "content": "请根据用户的提问，基于Chinook Music数据库的信息，生成SQL语句来回答用户的问题。"},
        {"role": "user", "content": f"{query}"},
    ]

    # 执行工具调用,获取结果
    completion = client.chat.completions.create(
        model="glm-4-flash",
        messages=messages,
        tools=tools,
        tool_choice="auto"  # 工具选择模式为auto,表示由LLM自行推理,觉得是生成普通消息还是进行工具调用
    )

    # 将工具调用结果解析成sql字符串,并返回
    return json.loads(completion.choices[0].message.tool_calls[0].function.arguments).get("sql")
    

在这里，我们定义了 generate_sql 这个工具，并添加了详细的描述信息，便于大模型理解用户的需求。此外，在工具描述中，我们将数据库的表结构信息 db_schema 也作为上下文信息传递给 LLM，这样 LLM 就可以将用户的提问翻译成具体的 SQL 语句。

2.4 实际效果演示

到这里，功能就实现完成了，是不是非常简单？我们简单演示下具体的效果：在控制台循环接收用户的原始提问，打印生成的 SQL 语句，并且获取最终的查询结果：

if __name__ == '__main__':
    # 创建sqlite数据库连接
    conn = connect_db()

    # 获取数据库的schema信息
    db_schema = get_database_schema(conn)

    # 在控制台循环获取用户输入
    while True:
        query = input("用户提问: ")
        if query == "bye":
            break

        # 将用户提问翻译成SQL
        sql = generate_sql(db_schema, query)
        print("--------------------------------------------------")
        print(f"生成的SQL语句: \n{sql}")

        # 执行SQL,获取结果
        answer = exec_sql(conn, sql)
        print("--------------------------------------------------")
        print(f"执行结果: {answer}")

最终执行结果如下：

可以看到，利用 LLM 的 Tool Calling 能力，大模型很好地理解了业务同学的自然语言提问，并将其翻译成结构化的标准 SQL 生成。应用程序获取到 SQL ，就可以直接在数据库执行，获取最终的查询结果。整个查询没有任何技术门槛、高度可定制化，并且完全不需要研发同学参与，极大地节省了人力和开发资源。

3. 生产环境最佳实践

这个工具整体上基本可以满足业务同学日常的查询需求，但是如果想实际用到生成环境，还需要进行进一步的完善和优化，我自己整理了以下的优化点，大家可以结合自己的业务场景进行适配：

1. 为该工具开发简单的前端页面，提升工具的易用性；
2. 由于 LLM 生成内容的随机性，输出的 SQL 并不一定 100% 是可执行的，程序中需要做好重试、异常捕获等容错机制；
3. 使用只读库或离线库执行查询，避免频繁的 SQL 执行对线上业务造成影响；
4. 对于执行查询的数据库，做好权限控制；
5. LLM 的 API 服务通常有限流机制，如果查询较为频繁，需要在程序中做好限流、限速等机制。

4. 总结

随着人工智能的发展，现在的大模型已经具有了非常强大的推理能力，并且随着上下文长度的扩展，再结合函数、工具调用、插件、Agent 等功能，大模型一定可以越来越多地覆盖产品设计、研发、测试、部署等各个生命周期。本文中介绍的利用 LLM 将自然语言查询翻译成结构化的标准 SQL，虽然功能并不复杂，但是其实是一种编程范式上的转变。人类和计算机本来就是一对互相合作的伙伴，但是它们之间无法直接交流，而编程语言本质上就是人类与计算机的沟通工具。以前计算机的理解能力较差，需要人类持续开发更高级的编程语言，来不断靠近计算机。而现在有了大模型，计算机的理解、推理能力越来越强，终于到了计算机可以主动走近人类的时代，它已经慢慢可以理解人类的语言了。也许在不久的将来，我们将不再需要任何编程语言，或者说自然语言本身就是编程语言。

（完整项目代码：https://gitee.com/zhangshenao/happy-llm/blob/master/openai_api/tool_call/sql_generator_zhipu.py）