5行代码打造你的大模型数据库聊天机器人

看效果

首先，让我们看一下最终产品。假设有一个包含有关记录销售信息的样本 SQLite 数据库，允许以自然语言提出任何问题，然后程序：

• 基于问题生成 SQL
• 将结果显示为表格
• 将结果显示为图表

​上代码

from vanna.remote import VannaDefault
from vanna.flask import VannaFlaskApp

vn = VannaDefault(model='chinook', api_key="your_vanna_api_key")
vn.connect_to_sqlite('https://vanna.ai/Chinook.sqlite')
VannaFlaskApp(vn).run()

过程拆解

当然要让这个应用真正运行起来，你需要做如下工作：

• 准备 Python 环境（虚拟机）
• 安装必要的包

  pip install vanna flask
  

• 注册 vanna 账号^[1]，并找到你的 API Key
• 输入你的邮箱
  
• 填写收到的验证码
  
• API Key 在这
  

当然你也可以尝试这样的方法，

#记得把[e-mail] 替换为你的真实邮箱，用于接受验证码
api_key = vanna.get_api_key('[e-mail]')
print(api_key)
vn = VannaDefault(model='chinook', api_key=api_key)

顺利的话，会有这样的提示，

Check your email for the code and enter it here:

输入收到的验证码，最好把 API Key 记录下来，下次直接使用 API Key，不然多次这样访问，可能会返回这样的报错，

requests.exceptions.SSLError: HTTPSConnectionPool(host='ask.vanna.ai', port=443):
Max retries exceeded with url: /unauthenticated_rpc (Caused by SSLError(SSLEOFError(8, '[SSL: UNEXPECTED_EOF_WHILE_READING]
EOF occurred in violation of protocol (_ssl.c:1002)')))

vanna

通过 RAG 使用 LLMs 进行准确的文本到 SQL 生成 ，创建你的数据库聊天机器人 。

RAG：检索增强生成（Retrieval-augmented generation，简称 RAG）是一种自然语言处理（NLP）技术，它结合了信息检索和文本生成的能力，旨在生成更准确和相关的文本。RAG 的核心思想是通过引入检索机制，使生成模型能够从外部数据源中检索信息，然后将该信息用于生成文本。

vanna 的特点：

• 开源
• 支持多种 LLMs
• 自学习，你可以持续训练你的模型
• 准确度高
• 支持多种数据库
• 默认支持 SQLite，Postgres，BigQuery 和 Snowflake
• 可以方便扩展数据库支持，比如我就轻松实现了 DuckDB 的支持
• 支持多种前端
• Jupyter Notebook^[2]
• Slackbot^[3]
• web app^[4]
• Streamlit app^[5]

最后放一张 DuckDB 的实现

从这张图可以看出，使用 DuckDB 可以方便分析各种 DuckDB 支持的数据源，如 SQLite，CSV，Excel，Parquet 等等。

Streamlit 版本

如果更喜欢 Streamlit，这里也把代码给你，

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())

参考资料

[1]

注册 vanna 账号: https://vanna.ai/account/login?next=/account/profile

[2]

Jupyter Notebook: https://vanna.ai/docs/snowflake-openai-vanna-vannadb.html

[3]

Slackbot: https://github.com/vanna-ai/vanna-slack

[4]

web app: https://github.com/vanna-ai/vanna-flask

[5]

Streamlit app: https://github.com/vanna-ai/vanna-streamlit