基于大模型实现 SQL Agent - 文章 - 开发者社区

背景

相信大部分研发同学都遇到过这样的场景：

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

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

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

实现过程

2.1 数据准备

为了演示功能，我们首先需要准备一个 SQL 数据库和相关数据。方便期间，我们采用 SQLite（https://www.sqlite.org/）和它官方提供的示例数据库（https://www.sqlitetutorial.net/sqlite-sample-database/）来完成我们的需求。

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

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

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

picture.image

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 利用大模型的 tool-calling 生成 SQL

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

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

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

picture.image

其他大模型的工具调用接口定义基本与 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}")

最终执行结果如下：

picture.image

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

生产环境最佳实践

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

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

随着人工智能的发展，现在的大模型已经具有了非常强大的推理能力，并且随着上下文长度的扩展，再结合函数、工具调用、插件、Agent 等功能，大模型一定可以越来越多地覆盖产品设计、研发、测试、部署等各个生命周期。

本文中介绍的利用 LLM 将自然语言查询翻译成结构化的标准 SQL，虽然功能并不复杂，但是其实是一种编程范式上的转变。人类和计算机本来就是一对互相合作的伙伴，但是它们之间无法直接交流，而编程语言本质上就是人类与计算机的沟通工具。以前计算机的理解能力较差，需要人类持续开发更高级的编程语言，来不断靠近计算机。

而现在有了大模型，计算机的理解、推理能力越来越强，终于到了计算机可以主动走近人类的时代，它已经慢慢可以理解人类的语言了。也许在不久的将来，我们将不再需要任何编程语言，或者说自然语言本身就是编程语言。

附完整代码：https://gitee.com/zhangshenao/happy-llm/blob/master/openai\_api/tool\_call/sql\_generator\_zhipu.py