[AI For DB] 基于 Dify 实现 SQL 补全功能

785 阅读7分钟

前言

Dify 是一个大模型的应用开发平台,提供了许多 LLM 开箱即用的能力。本篇文章主要分享如何通过 Dify 平台去打造 SQL 代码补全功能。

Pasted Graphic 6.png

Pasted Graphic 7.png

Pasted Graphic 8.png

基本设计

要想实现一个 SQL 补全,最简单的办法就是将用户当前的 SQL 原封不动的发给大模型,然后让模型来推断后续的内容。但是这样做意义不大,因为实际的 SQL 往往会涉及到当前环境的表,视图,函数等信息,失去了这些关键信息,大模型无法提供准确的信息,同时也可能会出现幻觉,生成数据库中不存在的表。

因此,我们需要尽可能的将当前数据库信息提供给大模型,如何提供全面,够用的信息是最终 SQL 补全效果好坏的关键。 我们可以把 SQL 补全功能分为三部分

  1. 获取当前 SQL 中关联到的表
  2. 去数据库中查询这些表的详细结构
  3. 将用户 SQL,依赖的表信息组装起来发送给大模型进行补全。

在 Dify 应用类型上,我们采用工作流将上述流程串联在一起。

获取关联表

我们可以把用户的 SQL 分为几类:DDL,非DDL,注释。不同的类型处理方式不一样。

对于非 DDL 类型的处理比较简单,我们只需要提取语句中访问的数据库表名即可。 比如:

select * from db.table;
INSERT INTO ccc1(`a`, `b`) VALUES(expr, expr);

这条语句经过提取后,可以得到 db.table, ccc1两个表。那么如何提取呢? 比较准确并且复杂的方式可以通过 SQL Parser 进行解析,然后通过分析语法树来得到对应的表。乍一听好像也不是特别难,那是因为上面提供的 SQL 比较简单,例如下面这条稍微复杂一点的

WITH tb1 as (SELECT * FROM a1) SELECT * FROM (
    SELECT * FROM ccc
) tb2, tb1 WHERE tb2.id = tb1.id; |

想要提取这条语句中访问的数据库表就比较麻烦了,需要一些专业能力。并且,这种方式无法解决多种 SQL 方言的问题,需要针对每一种方言准备一个 SQL Parser。

因此我们需要跳出传统思维,借助大模型的能力,将这一串 SQL 全部丢给大模型,让大模型去提取里面的信息,这种方式比较依赖模型能力,用 DeepSeek-R1 和 qwen-max 测下来准确率比较高,目前还没遇到提取出错的场景。

Pasted Graphic.png

对于 DDL 的场景,处理会稍微复杂一点,比如在一条查询语句之前执行了一个 alter 语句,我们需要考虑这条 alter 语句是否对后续的查询造成影响。

ALTER TABLE a1 DROP COLUMN `col7`;

SELECT * FROM a1 where | ...

在这篇分享中,我们暂时不考虑这种场景,先简单的忽略 DDL。

解决了 SQL 之后,我们开始考虑如何解决注释,注释中往往记录着用户的目的

-- 当年销售增长情况
select

我们需要从注释中提取出数据库里有关联的表,例如上述 SQL 中,我们要提取销售记录相关的表。在这里,我们需要通过 RAG 来实现。

首先,将用户数据库中所有的表,通过一些方法将其向量化之后存入向量数据库中,然后将注释信息向量化之后,去向量数据库中去匹配相似度最高的表。整个过程的核心就是如何将表进行向量化。

本分享主要说明整体流程,所以暂时采用比较简便的方式绕过向量化数据库,后续分享会详细说明数据库向量化的实践方法。我们首先获取数据库的所有表名,然后将其和注释一起发送给大模型,让大模型判断哪些表可能会引用到。这个方法实测下来准确度并不高,并且十分依赖模型能力。

Pasted Graphic 1.png

获取数据库表结构

我们已经获取到了一系列关联的数据库表,接下来我们就要通过 dify 去获取这些表的详细信息。要实现这种能力,可以通过扩展或者HTTP API 的方式实现,个人比较偏向于 HTTP API 的方式,因为扩展需要熟悉 dify 的扩展设计,并且对于 python 也要有一定的了解,而 HTTP API 可以使用任何技术栈,同时也能直接复用现有的业务平台能力,心智负担比较小,假如从长远的角度来看,HTTP API也是一种更好的方式,因为具备更好的跨平台复用能力,比如也能在 langchain 等 proCode 场景下使用。

同时考虑到于数据库交互获取信息的能力是通用的,我们可以将这块功能单独封装到一个工作流中 Pasted Graphic 2.png

Pasted Graphic 3.png

大模型补全

至此,我们已经拿到了需要提供给大模型的表列信息,接下来就是组装 prompt 发送,在 prompt 中我们需要明确指出让大模型根据我们提供的信息进行下一个 token 的推测,比如 SQL 内容,光标位置,数据库表信息。并且要求大模型只输出 SQL

整体的 Dify 工作流如下

Pasted Graphic 4.png

工作流调用

完成工作流的搭建之后,接下来的任务就是用编辑器去调用工作流来获取结果,我们可以很方便的通过 Dify 提供的 API 接口来调用,并且可以通过日志模块来查看调用情况,

Pasted Graphic 9.png

Pasted Graphic 10.png

如何选择模型

考虑过qwen-max,qwen-turbo,gpt,claude,deepseek-r1 等模型,最终选择了 qwen-max。 选择模型的时候差不多会考虑几个因素

  1. 模型的能力如何,在我的场景下是否够用
  2. 这个模型我们能通过合法渠道稳定使用吗
  3. 模型的token输出速度能满足我的场景吗
  4. 模型的价格

qwen-turbo 的效果不佳,准确率比较低,需要对 prompt 进行非常精细的控制。gpt 和 claude 系列由于无法合法稳定调用,所以也不考虑。deepseek-r1 从价格和能力上来说都非常符合要求,但是实测下来,token 输出速度实在是太慢了,实在是不适合代码编辑的场景。qwen-max 从能力,速度,稳定性来说都还行,因此选择了他。后续若是 deepseek-r1 的速度问题得到改善,优先考虑切换到 deepseek-r1 上

Agent Or Workflow

工作流做的事情完全可以通过 Agent 来自动完成,但是 Agent 实在是太依赖模型能力了,实测下来,准确性不如 Workflow,并且相较于 Workflow, Agent 并没有发挥出额外的惊喜效果,只是在应用实现上比较简单。

Agent 目前能想到的主要是在无法预测用户行为的场景,比如数据库助手,用户可能是提问,也有可能是想要创建一个表,也有可能是查询数据或者生成可视化图表,这个场景用 workflow 去实现会非常繁琐,用 agent 只需要将这几个场景依赖的能力通过 tools 传给 agent,让 agent 自己判断如何去组合工作去处理用户请求。而 Workflow 非常适合流程比较简单的场景。

Pasted Graphic 11.png

Next

接下来的分享将会介绍如何结合数据库 RAG 来进行更加精确的 SQL 补全