开源新宠：RAG2SQL工具，超越Text2SQL的7K Star之作

news2024/9/22 10:01:30

查询数据库离不开SQL，那如何快速构建符合自己期望的SQL呢？AI发展带来了Text2SQL的能力，众多产品纷纷提供了很好的支持。

今天我们分享一个开源项目，它在Text2SQL的基础上还要继续提高，通过加入RAG的能力进一步增强，它就是：Vanna

Vanna 是什么

从本质上讲，Vanna 是一个 Python 包，它使用检索增强来帮助您使用 LLM 为数据库生成准确的 SQL 查询。

Vanna 的工作过程分为两个简单步骤 - 在您的数据上训练 RAG“模型”，然后提出问题，这些问题将返回 SQL 查询，这些查询可以设置为在您的数据库上自动运行。

Vanna有三种模式，适应不同的需求。

第一种，开源模式。

你可以用它来集成任何你喜欢的语言模型，完全自定义。

就像自己搭积木，想怎么玩就怎么玩。

第二种，免费模式。

这里有日限额的模型使用，还有免费的存储服务。

用GPT 3.5，你可以不花一分钱就开始探索数据。

第三种，付费模式。

如果你需要更多，这里有无限制的使用，GPT-4的支持，还有服务保障。

想象一下，不管你是开发者还是分析师，都能用Vanna来发现数据的秘密。

使用 Vanna

你可以从Jupyter Notebook、Slackbot、网页应用或Streamlit等开始使用Vanna。当然也可以将Vanna集成到你的网页应用里。

如果你只是想简单体验一下Vanna，可以使用已经训练好的在线大模型和Vanna 提供的向量数据库，最简单的方式是通过Colab notebook。

!pip install vanna
import vanna
from vanna.remote import VannaDefault
vn = VannaDefault(model='chinook', api_key=vanna.get_api_key('my-email@example.com'))
vn.connect_to_sqlite('<https://vanna.ai/Chinook.sqlite>')
vn.ask("What are the top 10 albums by sales?")
from vanna.flask import VannaFlaskApp
VannaFlaskApp(vn).run()

这里需要使用到vanna 的apikey，所以需要注册vanna 的账号。

AI 数据库机器人

可以通过Streamlit 和 Vanna的组合来快速创建一个数据库的SQL聊天机器人，首先，我们将导入我们将使用的 2 个 Python 包：Streamlit 和 Vanna。

Streamlit提供用户界面
Vanna 提供Text2SQL的能力

import vanna as vn
import streamlit as st

设置

在这里，我们将使用 Vanna API 密钥，将 Vanna 检索增强模型设置为“chinook”，它具有该数据库的元数据，并且我们将连接到 SQLite 数据库，这将允许我们运行生成的 SQL。

vn.set_api_key(st.secrets[ "vanna_api_key" ]) 
vn.set_model( 'chinook' ) 
vn.connect_to_sqlite( '<https://vanna.ai/Chinook.sqlite>' )

对话框

my_question = st.text_input( "问我一个可以转化为 SQL 的问题" , key= "my_question" )

使用 AI 生成 SQL

sql  = vn.generate_sql(my_question) 
st.code( sql , language = 'sql' )

显示表格

df = vn.run_sql(sql)     
st.dataframe(df, use_container_width= True )

使用AI生成图表并显示

fig = vn.get_plotly_figure(plotly_code=vn.generate_plotly_code(question=my_question, sql=sql, df=df), df=df)
st.plotly_chart(fig, use_container_width=True)

整合

接下来我们将几部分内容最终整合成一个完整的示例。如下所示：

import vanna as vn
import streamlit as st

vn.set_api_key(st.secrets["vanna_api_key"])
vn.set_model('chinook')
vn.connect_to_sqlite('<https://vanna.ai/Chinook.sqlite>')

my_question = st.session_state.get("my_question", default=None)
if my_question is None:
    st.image("chinook-schema.png", use_column_width=True)
    my_question = st.text_input("Ask me a question that I can turn into SQL", key="my_question")
else:
    st.title(my_question)
    sql = vn.generate_sql(my_question)
    st.code(sql, language='sql')
    df = vn.run_sql(sql)    
    st.dataframe(df, use_container_width=True)
    fig = vn.get_plotly_figure(plotly_code=vn.generate_plotly_code(question=my_question, sql=sql, df=df), df=df)
    st.plotly_chart(fig, use_container_width=True)
    st.button("Ask another question", on_click=lambda: st.session_state.clear())