20行Python代码，利用Streamlit和Vanna，让SQL数据库拥有智能聊天机器人

一、前言

出于说明目的，本指南优化了代码行数来进行演示。对于功能更丰富的聊天机器人代码，可以参考如下GitHub存储库。

【GitHub Repo】：https://github.com/vanna-ai/vanna-streamlit

首先，让我们看看最终的产品。给定一个包含有关唱片销售信息的示例SQLite数据库，我们可以制作一个文本框，允许用户使用自然语言提问，然后：

• 基于问题生成SQL语句。
• 以表格形式显示结果。
• 以图表形式显示结果。

，时长00:15

聊天机器人的运行示例

二、Streamlit和Vanna

现在让我们来分解一下。首先，我们将导入两个Python软件包：Streamlit和Vanna。

• Streamlit提供用户界面。
• Vanna提供使用LLM（即人工智能）生成SQL的基础设施。

import vanna as vn
import streamlit as st

三、设置

3.1 模型设置

在这里，我们将使用Vanna API密钥，将Vanna Retrieval Augmentation模型设置为“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')

3.2 设置文本框

my_question = st.text_input("Ask me a question that I can turn into SQL", key="my_question")

图片

3.3 使用AI生成SQL语句

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

图片

3.4 显示表格

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

3.5 使用AI生成图表并显示

fig = vn.get_plotly_figure(plotly_code=vn.generate_plotly_code(questinotallow=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(questinotallow=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())

图片

五、更多内容

如果想了解Vanna更多内容及下载相关源码，可通过如下链接查询。

【GitHub】：https://github.com/vanna-ai/vanna