Text-to-SQL 工具Vanna进阶 | 数据库对话机器人的多轮对话 - 文章 - 开发者社区

跟数据库对话机器人对话，我可不止一个问题。
可能基于第一句问话，还有第二句、第三句问话。。。第N句对话。所以本文测试了多轮对话功能。

基础配置2. 多轮生成SQL3. 多轮对话

单轮对话的环境搭建参考博客 [Text-to-SQL 工具Vanna + MySQL本地部署 | 数据库对话机器人]https://blog.csdn.net/sinat\_29950703/article/details/136686608?spm=1001.2014.3001.5501

我的数据是这样

picture.image

基础配置 =======

  
import vanna  
from vanna.remote import VannaDefault  
from vanna.flask import VannaFlaskApp  
import pandas as pd  
import mysql.connector  
  
api_key = '7acxxx68c'  
vanna_model_name = 'community'  
vn = VannaDefault(model=vanna_model_name, api_key=api_key)  
  
def run\_sql(sql: str) -> pd.DataFrame:  
    cnx = mysql.connector.connect(user='root',password='111000',host='localhost',database='Community')  
    cursor = cnx.cursor()  
    cursor.execute(sql)  
    result = cursor.fetchall()  
    columns = cursor.column_names  
    # print('columns:',columns)  
    df = pd.DataFrame(result, columns=columns)  
    return df  
  
# 将函数设置到vn.run\_sql中  
vn.run_sql = run_sql  
vn.run_sql_is_set = True

多轮生成SQL ==========

用 vn.generate 生成SQL
使用方法看官方文档 https://vanna.ai/docs/vanna/remote.html

picture.image

多轮SQL的代码如下：

  
first_sql = vn.generate_sql('查询30岁以上的姓名、年龄、身份证？')  
if first_sql is not None:  
    print('first\_sql:\n',first_sql,)  
  
    second_sql = vn.generate_sql(f'基于SQL语句{first\_sql}修改,添加查询性别？')  
    print('\nsecond\_sql:\n',second_sql)  
  
    third_sql = vn.generate_sql(f'基于SQL语句{second\_sql}修改,查询女性、苗族？')  
    print('\nthird\_sql:\n',third_sql)  
  
    forth_sql = vn.generate_sql(f'基于SQL语句{third\_sql}修改,要查询男性的少数民族信息？')  
    print('\nforth\_sql:\n',forth_sql)  
  
    fivth_sql = vn.generate_sql(f'基于SQL语句{forth\_sql}修改,不要身份证号，要手机号？')  
    print('\nfivth\_sql:\n',fivth_sql)

无任何训练语句，SQL生成结果如下：

  
first_sql:  
 SELECT name, age, id_card  
FROM customer  
WHERE age > 30;  
  
second_sql:  
 SELECT name, age, id_card, gender  
FROM customer  
WHERE age > 30;  
  
third_sql:  
 SELECT name, age, id_card, gender  
FROM customer  
WHERE age > 30 AND gender = 2 AND nation = '苗族';  
  
forth_sql:  
 SELECT name, age, id_card, gender, nation  
FROM customer  
WHERE gender = 1 AND nation NOT IN ('汉族', '壮族', '回族', '满族', '维吾尔族', '苗族', '彝族', '藏族', '布依族', '侗族', '瑶族', '朝鲜族', '白族', '哈尼族', '哈 萨克族', '黎族', '傣族', '畲族', '傈僳族', '仡佬族', '东乡族', '高山族', '拉祜族', '水族', '佤族', '纳西族', '羌族', '土族', '仫佬族', '锡伯族', '柯尔克孜族', '达斡尔族', '景颇族', '毛南族', '撒拉族', '布朗族', '塔吉克族', '阿昌族', '普米族', '鄂温克族', '怒族', '京族', '基诺族', '德昂族', '保安族', '俄罗斯族', '裕固族', '乌孜别克族', '门巴族', '鄂伦春族', '独龙族', '塔塔尔族', '赫哲族', '珞巴族')  
  
fivth_sql:  
 SELECT name, age, mobile, gender, nation  
FROM customer  
WHERE gender = 1 AND nation NOT IN ('汉族', '壮族', '回族', '满族', '维吾尔族', '苗族', '彝族', '藏族', '布依族', '侗族', '瑶族', '朝鲜族', '白族', '哈尼族', '哈 萨克族', '黎族', '傣族', '畲族', '傈僳族', '仡佬族', '东乡族', '高山族', '拉祜族', '水族', '佤族', '纳西族', '羌族', '土族', '仫佬族', '锡伯族', '柯尔克孜族', '达斡尔族', '景颇族', '毛南族', '撒拉族', '布朗族', '塔吉克族', '阿昌族', '普米族', '鄂温克族', '怒族', '京族', '基诺族', '德昂族', '保安族', '俄罗斯族', '裕固族', '乌孜别克族', '门巴族', '鄂伦春族', '独龙族', '塔塔尔族', '赫哲族', '珞巴族')

添加训练语句 ，让vn学习优化SQL语句

  
vn.train(question='哪些是少数民族？',sql="SELECT name FROM std\_people WHERE nation !='汉族'")

SQL生成结果如下：

截图

picture.image

原输出：

  
first_sql:  
 SELECT name, age, id_card  
FROM customer  
WHERE age > 30;  
  
second_sql:  
 SELECT name, age, id_card, gender  
FROM customer  
WHERE age > 30;  
  
third_sql:  
 SELECT name, age, id_card, gender  
FROM customer  
WHERE age > 30 AND gender = 2 AND nation = '苗族';  
  
forth_sql:  
 SELECT name, age, id_card, gender, nation  
FROM customer  
WHERE gender = 1 AND nation != 'hanNat';  
  
fivth_sql:  
 SELECT name, age, gender, nation, mobile  
FROM customer  
WHERE gender = 1 AND nation != 'hanNat';

多轮对话 =======

用 vn.ask 对话：生成SQL、执行SQL、展示结果df和网页
使用方法参考Vanna源代码

question：接收的问题
print_results：是否在终端打印结果df
auto_train：是否参加vn的训练，默认把问题和SQL回答再放入vn训练
visualize：是否可视化，是否要生成可视化的图在网页端打开

picture.image

代码如下：

  
first_result = vn.ask('查询30岁以上的姓名、年龄、身份证？')  
if first_result is not None:  
    first_sql,_,_ = first_result  
    print('first\_sql:\n',first_sql,)  
  
    second_sql,_,_ = vn.ask(f'基于SQL语句{first\_sql}修改,添加查询性别？')  
    print('\nsecond\_sql:\n',second_sql)  
  
    third_sql,_,_ = vn.ask(f'基于SQL语句{second\_sql}修改,查询女性、苗族？')  
    print('\nthird\_sql:\n',third_sql)  
  
    forth_sql,_,_ = vn.ask(f'基于SQL语句{third\_sql}修改,查询男性的少数民族信息？')  
    print('\nforth\_sql:\n',forth_sql)  
  
    fivth_sql,_,_ = vn.ask(f'基于SQL语句{forth\_sql}修改,不要身份证号，要手机号？')  
    print('\nfivth\_sql:\n',fivth_sql)

ask 和 generate_sql的区别

generate_sql ：只给出SQL语句
ask : 给出SQL语句和执行结果(dataframe格式)。

如果把print_results = False, 则vn.ask返回结果取第一项，则与generate_sql无异。