🤖 Chat with your SQL database 📊. Accurate Text-to-SQL Generation via LLMs using RAG 🔄.
GitHub | PyPI | Documentation | Gurubase |
---|---|---|---|
Vanna is an MIT-licensed open-source Python RAG (Retrieval-Augmented Generation) framework for SQL generation and related functionality.
https://github.com/vanna-ai/vanna/assets/7146154/1901f47a-515d-4982-af50-f12761a3b2ce
Vanna works in two easy steps - train a RAG “model” on your data, and then ask questions which will return SQL queries that can be set up to automatically run on your database.
If you don’t know what RAG is, don’t worry – you don’t need to know how this works under the hood to use it. You just need to know that you “train” a model, which stores some metadata and then use it to “ask” questions.
See the base class for more details on how this works under the hood.
These are some of the user interfaces that we’ve built using Vanna. You can use these as-is or as a starting point for your own custom interface.
See the documentation for specifics on your desired database, LLM, etc.
If you want to get a feel for how it works after training, you can try this Colab notebook.
pip install vanna
There are a number of optional packages that can be installed so see the documentation for more details.
See the documentation if you’re customizing the LLM or vector database.
# The import statement will vary depending on your LLM and vector database. This is an example for OpenAI + ChromaDB
from vanna.openai.openai_chat import OpenAI_Chat
from vanna.chromadb.chromadb_vector import ChromaDB_VectorStore
class MyVanna(ChromaDB_VectorStore, OpenAI_Chat):
def __init__(self, config=None):
ChromaDB_VectorStore.__init__(self, config=config)
OpenAI_Chat.__init__(self, config=config)
vn = MyVanna(config={'api_key': 'sk-...', 'model': 'gpt-4-...'})
# See the documentation for other options
You may or may not need to run these vn.train
commands depending on your use case. See the documentation for more details.
These statements are shown to give you a feel for how it works.
DDL statements contain information about the table names, columns, data types, and relationships in your database.
vn.train(ddl="""
CREATE TABLE IF NOT EXISTS my-table (
id INT PRIMARY KEY,
name VARCHAR(100),
age INT
)
""")
Sometimes you may want to add documentation about your business terminology or definitions.
vn.train(documentation="Our business defines XYZ as ...")
You can also add SQL queries to your training data. This is useful if you have some queries already laying around. You can just copy and paste those from your editor to begin generating new SQL.
vn.train(sql="SELECT name, age FROM my-table WHERE name = 'John Doe'")
vn.ask("What are the top 10 customers by sales?")
You’ll get SQL
SELECT c.c_name as customer_name,
sum(l.l_extendedprice * (1 - l.l_discount)) as total_sales
FROM snowflake_sample_data.tpch_sf1.lineitem l join snowflake_sample_data.tpch_sf1.orders o
ON l.l_orderkey = o.o_orderkey join snowflake_sample_data.tpch_sf1.customer c
ON o.o_custkey = c.c_custkey
GROUP BY customer_name
ORDER BY total_sales desc limit 10;
If you’ve connected to a database, you’ll get the table:
CUSTOMER_NAME | TOTAL_SALES | |
---|---|---|
0 | Customer#000143500 | 6757566.0218 |
1 | Customer#000095257 | 6294115.3340 |
2 | Customer#000087115 | 6184649.5176 |
3 | Customer#000131113 | 6080943.8305 |
4 | Customer#000134380 | 6075141.9635 |
5 | Customer#000103834 | 6059770.3232 |
6 | Customer#000069682 | 6057779.0348 |
7 | Customer#000102022 | 6039653.6335 |
8 | Customer#000098587 | 6027021.5855 |
9 | Customer#000064660 | 5905659.6159 |
You’ll also get an automated Plotly chart:
RAG
Fine-Tuning
Vanna is designed to connect to any database, LLM, and vector database. There’s a VannaBase abstract base class that defines some basic functionality. The package provides implementations for use with OpenAI and ChromaDB. You can easily extend Vanna to use your own LLM or vector database. See the documentation for more details.
https://github.com/vanna-ai/vanna/assets/7146154/eb90ee1e-aa05-4740-891a-4fc10e611cab