Создайте чатбота для вашей базы данных SQL за 20 строк Python с помощью Streamlit и Vanna

В этом руководстве оптимизировано количество строк кода для наглядности. Код более многофункционального чатбота можно посмотреть в этом репозитории GitHub.

Сначала давайте посмотрим на конечный продукт. Учитывая пример базы данных SQLite с информацией о продажах записей, мы можем сделать текстовое поле, которое позволит вам задать любой вопрос на естественном языке:

• Генерировать SQL на основе вопроса
• Показать результаты в виде таблицы
• Показать результаты в виде диаграммы

Streamlit и Vanna

Теперь давайте разберемся с этим. Сначала импортируем 2 пакета Python, которые мы будем использовать: Streamlit и Vanna.

• Streamlit предоставляет пользовательский интерфейс.
• Vanna предоставляет инфраструктуру для использования LLM (т.е. искусственного интеллекта) для генерации SQL

import vanna as vn
import streamlit as st

Настройка

Здесь мы используем наш ключ Vanna API, установим модель Vanna Retrieval Augmentation Model на “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("Ask me a question that I can turn into SQL", key="my_question")

Генерация SQL с помощью искусственного интеллекта

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)

Собираем все вместе

На самом деле это полный код, показанный в примере!

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

Начните с этого