Chat with Your SQL Database Using Vanna AI

Sanjjushri Varshini R
2 min readJul 22, 2024

--

In this article, we’ll explore leveraging Vanna AI to seamlessly interact with your SQL database. Vanna AI combines the power of Large Language Models (LLMs) and ChromaDB Vector Stores to provide an intuitive interface for querying databases. Below is a step-by-step guide to setting up and using Vanna AI.

Step 1: Install Vanna AI

First, we need to install Vanna AI and other necessary libraries. This can be done using pip:

pip install vanna==0.6.2 ollama==0.1.6

Step 2: Import Vanna Modules

Import the necessary modules from Vanna AI:

from vanna.ollama import Ollama
from vanna.chromadb import ChromaDB_VectorStore

Step 3: Create a Custom Vanna Class

Define a custom class that inherits from both ChromaDB_VectorStore and Ollama:

class MyVanna(ChromaDB_VectorStore, Ollama):
def __init__(self, config=None):
ChromaDB_VectorStore.__init__(self, config=config)
Ollama.__init__(self, config=config)

vn = MyVanna(config={'model': 'llama3'})

Step 4: Connect to Your SQL Database

Uncomment and fill in your database credentials to connect:

# type in your creds here and uncomment to connect
"""
vn.connect_to_mysql(host='', dbname='', user='', password='', port=)
"""

Step 5: Query the Information Schema

Run an SQL query to retrieve the database schema:

df_information_schema = vn.run_sql("SELECT * FROM INFORMATION_SCHEMA.COLUMNS")

Step 6: Generate a Training Plan

Break the information schema into manageable chunks for the LLM:

# This will break up the information schema into bite-sized chunks that can be referenced by the LLM
plan = vn.get_training_plan_generic(df_information_schema)
plan

Step 7: Train the Model

If satisfied with the training plan, uncomment and run the training command:

# If you like the plan, then uncomment this and run it to train
# vn.train(plan=plan)

Step 8: Ask Questions

Use the trained model to ask questions about your data:

vanna_ans = vn.ask(question="how many artists are present?")

Step 9: Extract the Exact Answer

Retrieve just the exact answer from the response:

# getting the exact answer alone:
vanna_ans[1]

Step 10: Deploy as a Flask App

Set up a Flask application to interact with Vanna AI via a web interface:

from vanna.flask import VannaFlaskApp
app = VannaFlaskApp(vn, allow_llm_to_see_data=True)
app.run()

The VannaFlaskApp class creates a web server that allows users to interact with the LLM and the database through a web browser.

This guide provides a comprehensive overview of setting up and using Vanna AI to interact with your SQL database. By following these steps, you can leverage the power of LLMs to query and analyze your data more intuitively.

GitHub: https://github.com/Sanjjushri/vanna-ai/tree/main

--

--