DB_GPT源码梳理
引言
随着数据量的不断增长和数据分析的需求日益增多,将自然语言文本转化为结构化查询语言(SQL)的能力变得越来越重要。Text to SQL方案是一种将自然语言查询转化为SQL查询的技术,它可以帮助用户更轻松、更高效地从文本中提取所需的信息。
博主近期在研究text2sql的项目应用,从大模型和传统的深度学习模型两方面入手,未来会持续发布相关文章,本文将介绍Text to SQL的一种基于大模型的方案:DB-GPT。安装与部署可以参考博主这篇文章:部署DB-GPT
DB-GPT项目简介
DB-GPT项目(项目地址)是在github上发布的,为解决使用大模型和数据库交互的过程中,私密数据以及环境是否能掌握自己的手里,完全自主可控的问题。项目支持为所有以数据库为基础的场景,构建一套完整的私有大模型解决方案。 此方案因为支持本地部署,所以不仅仅可以应用于独立私有环境,而且还可以根据业务模块独立部署隔离,让大模型的能力绝对私有、安全、可控。
概要
对于DBGPT的源码梳理,主要是针对于与数据库对话的部分源码分析
1后端
主要梳理提示词,以及用户的输入和大模型的响应是怎么被处理
1.1整体流程分析
首先我们选择聊天场景
根据前端抓包,获得的地址中可以看到http://localhost:5670/chat/?scene=chat_with_db_execute&id=8c88bf60-1f17-11ef-80b2-7413eaf68a78,scene=chat_with_db_execute这个参数,我们对应选择的聊天场景应该为ChatWithDbExecute
在这种聊天模式下,选定提示词模板
self.prompt_template: AppScenePromptTemplateAdapter = (
CFG.prompt_template_registry.get_prompt_template(
self.chat_mode.value(),
language=CFG.LANGUAGE,
model_name=self.llm_model,
proxyllm_backend=CFG.PROXYLLM_BACKEND,
)
)
英文
_DEFAULT_TEMPLATE_EN = """
Please answer the user's question based on the database selected by the user and some of the available table structure definitions of the database.
Database name:
{db_name}
Table structure definition:
{table_info}
Constraint:
1.Please understand the user's intention based on the user's question, and use the given table structure definition to create a grammatically correct {dialect} sql. If sql is not required, answer the user's question directly..
2.Always limit the query to a maximum of {top_k} results unless the user specifies in the question the specific number of rows of data he wishes to obtain.
3.You can only use the tables provided in the table structure information to generate sql. If you cannot generate sql based on the provided table structure, please say: "The table structure information provided is not enough to generate sql queries." It is prohibited to fabricate information at will.
4.Please be careful not to mistake the relationship between tables and columns when generating SQL.
5.Please check the correctness of the SQL and ensure that the query performance is optimized under correct conditions.
6.put the type name into the name parameter value that returns the required format. If you cannot find the most suitable one, use 'Table' as the display method. , the available data display methods are as follows: {display_type}
User Question:
{user_input}
Please think step by step and respond according to the following JSON format:
{response}
Ensure the response is correct json and can be parsed by Python json.loads.
"""
中文
_DEFAULT_TEMPLATE_ZH = """
请根据用户选择的数据库和该库的部分可用表结构定义来回答用户问题.
数据库名:
{db_name}
表结构定义:
{table_info}
约束:
1. 请根据用户问题理解用户意图,使用给出表结构定义创建一个语法正确的 {dialect} sql,如果不需要sql,则直接回答用户问题。
2. 除非用户在问题中指定了他希望获得的具体数据行数,否则始终将查询限制为最多 {top_k} 个结果。
3. 只能使用表结构信息中提供的表来生成 sql,如果无法根据提供的表结构中生成 sql ,请说:“提供的表结构信息不足以生成 sql 查询。” 禁止随意捏造信息。
4. 请注意生成SQL时不要弄错表和列的关系
5. 请检查SQL的正确性,并保证正确的情况下优化查询性能
6.类型名称放入返回要求格式的name参数值中,如果找不到最合适的则使用'Table'作为展示方式,可用数据展示方式如下: {display_type}
用户问题:
{user_input}
请一步步思考并按照以下JSON格式回复:
{response}
确保返回正确的json并且可以被Python json.loads方法解析.
"""
效果最好的是英文提问加英文模板----》中文提问中文模板-----》中英混杂
需要传入的参数
经过处理之后的参数
部分参数解释
对于模板中的{response}参数采用默认的回复格式
RESPONSE_FORMAT_SIMPLE = {
"thoughts": "和用户解释我们应该如何解决这个问题",
"sql": "要运行的sql语句查询",
"display_type": "数据展示方式",
}
对于6.类型名称放入返回要求格式的name参数值中,如果找不到最合适的则使用’Table’作为展示方式,可用数据展示方式如下: {display_type}
display_type可选的类型如下
response_line_chart:used to display comparative trend analysis data
response_pie_chart:suitable for scenarios such as proportion and distribution statistics
response_table:suitable for display with many display columns or non-numeric columns
response_scatter_plot:Suitable for exploring relationships between variables, detecting outliers, etc.
response_bubble_chart:Suitable for relationships between multiple variables, highlighting outliers or special situations, etc.
response_donut_chart:Suitable for hierarchical structure representation, category proportion display and highlighting key categories, etc.
response_area_chart:Suitable for visualization of time series data, comparison of multiple groups of data, analysis of data change trends, etc.
response_heatmap:Suitable for visual analysis of time series data, large-scale data sets, distribution of classified data, etc.
table_info就是对应我们所选择的数据库dbgpt_test下表的结构
当传入这些参数之后我们可以看到模板发生的变化
role='system'
content='\n请根据用户选择的数据库和该库的部分可用表结构定义来回答用户问题.
数据库名:
dbgpt_test
表结构定义: [(\'riskcontroller(measure_type,risk_level,harm,owner_id,order_number,safe_measure,id,procedures)\',), (\'security(id,owner_id,execute_content,measure_type,execute_situation,order_number)\',), (\'transaction_order(id,order_no,product_name,product_category,amount,pay_status,user_id,user_name,create_time,update_time)\',), (\'user(id,name,email,mobile,gender,birth,country,city,create_time,update_time)\',)]
约束:
1. 请根据用户问题理解用户意图,使用给出表结构定义创建一个语法正确的 mysql sql,如果不需要sql,则直接回答用户问题。\n
2. 除非用户在问题中指定了他希望获得的具体数据行数,否则始终将查询限制为最多 50 个结果。\n
3. 只能使用表结构信息中提供的表来生成 sql,如果无法根据提供的表结构中生成 sql ,请说:“提供的表结构信息不足以生成 sql 查询。” 禁止随意捏造信息。\n
4. 请注意生成SQL时不要弄错表和列的关系\n
5. 请检查SQL的正确性,并保证正确的情况下优化查询性能\n
6.将类型名称放入返回要求格式的name参数值中,如果找不到最合适的则使用\'Table\'作为展示方式,可用数据展示方式如下:
response_line_chart:used to display comparative trend analysis data
response_pie_chart:suitable for scenarios such as proportion and distribution statistics
response_table:suitable for display with many display columns or non-numeric columns
response_scatter_plot:Suitable for exploring relationships between variables, detecting outliers, etc.
response_bubble_chart:Suitable for relationships between multiple variables, highlighting outliers or special situations, etc.
response_donut_chart:Suitable for hierarchical structure representation, category proportion display and highlighting key categories, etc.
response_area_chart:Suitable for visualization of time series data, comparison of multiple groups of data, analysis of data change trends, etc.
response_heatmap:Suitable for visual analysis of time series data, large-scale data sets, distribution of classified data, etc.
用户问题:
查询每个用户的订单数量,并用饼图来进行表示
请一步步思考并按照以下JSON格式回复:
"{\\n \\"thoughts\\": \\"和用户解释我们应该如何解决这个问题\\",\\n \\"sql\\": \\"要运行的sql语句查询\\",\\n \\"display_type\\": \\"数据展示方式\\"\\n}"
确保返回正确的json并且可以被Python json.loads方法解析.\n\n' round_index=0
传入模型之后,得到模型的原始输出
得到输出结果如下
'{
"thoughts": "我们可以通过在transaction_order表中按照user_id进行分组,并计算每个用户的订单数量来解决这个问题。",
"sql": "SELECT user_id, COUNT(*) AS order_count FROM transaction_order GROUP BY user_id",
"display_type": "response_pie_chart"
}'
最后转换为一个用户可以理解和使用的格式,通常是前端显示所需的格式
view_message = await blocking_func_to_async(
self._executor,
self.prompt_template.output_parser.parse_view_response,
speak_to_user,
result,
prompt_define_response,
)
通过输出,得知结构如下
我们可以通过在transaction_order表中按照user_id进行分组,并计算每个用户的订单数量来解决这个问题。
<chart-view content="{"type": "response_pie_chart", "sql": "SELECT user_id, COUNT(*) AS order_count FROM transaction_order GROUP BY user_id", "data": [{"user_id": 111, "order_count": 5}, {"user_id": 112, "order_count": 13}, {"user_id": 113, "order_count": 7}, {"user_id": 114, "order_count": 5}, {"user_id": 121, "order_count": 9}, {"user_id": 122, "order_count": 11}, {"user_id": 123, "order_count": 16}, {"user_id": 124, "order_count": 10}, {"user_id": 125, "order_count": 8}, {"user_id": 211, "order_count": 10}, {"user_id": 212, "order_count": 7}, {"user_id": 213, "order_count": 8}, {"user_id": 214, "order_count": 10}, {"user_id": 221, "order_count": 6}, {"user_id": 222, "order_count": 6}, {"user_id": 223, "order_count": 7}, {"user_id": 311, "order_count": 5}, {"user_id": 312, "order_count": 7}, {"user_id": 313, "order_count": 10}, {"user_id": 314, "order_count": 3}, {"user_id": 315, "order_count": 4}, {"user_id": 321, "order_count": 11}, {"user_id": 322, "order_count": 8}, {"user_id": 323, "order_count": 5}, {"user_id": 324, "order_count": 9}]}" />
如果在我们的问题中不指定要用什么形式进行展示就选择默认形式
"{
"thoughts": "我们可以通过查询'user'表中的记录数量来获取用户数量。",
"sql": "SELECT COUNT(*) AS user_count FROM user",
"display_type": "response_table"
}"
模型输出如下
"我们可以通过查询'user'表中的记录数量来获取用户数量。
<chart-view content="{"type": "response_table", "sql": "SELECT COUNT(*) AS user_count FROM user", "data": [{"user_count": 25}]}" />"
1.2sql执行过程分析
查询模型原始输出
根据追踪代码,可以得知,ModelOutput.text就是模型的原始输出
@dataclass
@PublicAPI(stability="beta")
class ModelOutput:
"""A class to represent the output of a LLM.""" ""
text: str
"""The generated text."""
RESPONSE_FORMAT_SIMPLE = {
"thoughts": "和用户解释我们应该如何解决这个问题",
"sql": "要运行的sql语句查询",
"display_type": "数据展示方式",
}
model_output.text = "{
"thoughts": "我们应该查询'user'表中的用户数量。",
"sql": "SELECT COUNT(*) AS user_count FROM user",
"display_type": "response_table"
}"
提取SQL语句
● 从解析出的AI回复中提取SQL语句。
● “sql”: “SELECT COUNT(*) AS user_count FROM user”,
● 在代码中找到执行sql的地方
路径:D:\Code\Code_pycharm\DB_gpt\DB-GPT\dbgpt\app\scene\base_chat.py 下的380行
调用之后执行D:\Code\Code_pycharm\DB_gpt\DB-GPT\dbgpt\app\scene\chat_db\auto_execute\chat.py下的do_action()方法
执行SQL语句
执行sql的核心代码
核心sql执行代码
def run_to_df(self, command: str, fetch: str = "all"):
"""Execute sql command and return result as dataframe."""
import pandas as pd
# Pandas has too much dependence and the import time is too long
# TODO: Remove the dependency on pandas
result_lst = self.run(command, fetch)
colunms = result_lst[0]
values = result_lst[1:]
result_final = pd.DataFrame(values, columns=colunms)
return result_final
执行完sql语句之后结果封装成DataFrame类型
解析视图响应
# 定义 parse_view_response 方法,用于解析视图响应
def parse_view_response(self, speak, data, prompt_response) -> str:
# 初始化一个空字典 param 和一个 XML 元素 api_call_element,类型为 "chart-view"#
param = {}
api_call_element = ET.Element("chart-view")
err_msg = None
success = False
try:
if not prompt_response.sql or len(prompt_response.sql) <= 0:
raise AppActionException("Can not find sql in response", speak)
#调用 data 函数执行 SQL 语句,获取结果数据框 df
# ['user_count'] [0 25]
df = data(prompt_response.sql)
param["type"] = prompt_response.display
param["sql"] = prompt_response.sql
param["data"] = json.loads(
df.to_json(orient="records", date_format="iso", date_unit="s")
)
#param 字典转换为 JSON 字符串
view_json_str = json.dumps(param, default=serialize, ensure_ascii=False)
success = True
except Exception as e:
logger.error("parse_view_response error!" + str(e))
err_param = {
"sql": f"{prompt_response.sql}",
"type": "response_table",
"data": [],
}
# err_param["err_msg"] = str(e)
err_msg = str(e)
view_json_str = json.dumps(err_param, default=serialize, ensure_ascii=False)
# 这里将 JSON 字符串设置为 XML 元素的内容
api_call_element.set("content", view_json_str)
# 将XML 元素对象按照utf-8的编码方式,转化为字符串
result = ET.tostring(api_call_element, encoding="utf-8")
if not success:
view_content = (
f'{speak} \\n <span style="color:red">ERROR!</span>'
f"{err_msg} \n {result.decode('utf-8')}"
)
raise AppActionException("Generate view content failed", view_content)
else:
return speak + "\n" + result.decode("utf-8")
然后就能看到之前呈现的效果
就是这个效果
<chart-view content="{"type": "response_table", "sql": "SELECT COUNT(*) AS user_count FROM user", "data": [{"user_count": 25}]}" />'
2前端
前端主要梳理请求参数的构造和响应的处理,尤其是图表展示(折线图,饼图,柱状图等等如何展示)(如果有需要的话,评论区留言)