20行Python代碼，利用Streamlit和Vanna，讓SQL數據庫擁有智能聊天機器人

作者：學研君 2024-01-31 18:04:58

數據庫其他數據庫

在這里，我們將使用Vanna API密鑰，將Vanna Retrieval Augmentation模型設置為“chinook”，該模型具有此數據庫的元數據，并連接到SQLite數據庫，這樣我們就可以運行生成的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 模型設置

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

責任編輯：武曉燕來源： Python學研大本營

SQL 語句元數據