Text-to-SQL with Vanna-AI: A Developer's Guide

vanna-ai/vanna

Imagine you're working on an application and you constantly need to query your database. Sometimes, writing complex SQL queries can be time-consuming, or maybe you're dealing with a database you're not entirely familiar with. This is where vanna-ai/vanna shines!

At its core, Vanna is a Text-to-SQL Generation tool powered by Large Language Models (LLMs) using Retrieval Augmented Generation (RAG). In simpler terms, it allows you to "chat" with your SQL database using plain English. You ask a question in natural language, and Vanna translates that into the correct SQL query, executes it, and gives you the results.

Here's how it can be incredibly useful for a software engineer

Increased Productivity
Instead of manually crafting SQL queries, you can simply ask Vanna for the data you need. This is especially useful for ad-hoc queries, data exploration, or debugging.

Reduced Context Switching
You can stay within your natural language flow without needing to switch your brain to "SQL mode" as frequently.

Accessibility for Non-SQL Users
While you're a software engineer and likely know SQL, imagine a scenario where a product manager or a business analyst needs to pull some data. With Vanna, they could potentially get the data themselves without needing to bother a developer.

Learning and Exploration
If you're new to a database schema, Vanna can help you explore it by allowing you to ask questions about the data without having to meticulously inspect table definitions.

Data Validation and Debugging
Quickly run checks on your data by asking questions like "What are the top 10 customers by sales last month?"

"Smart" Query Generation
Because it uses LLMs and RAG, Vanna can potentially understand more nuanced requests and generate more accurate queries than a simple keyword-to-SQL converter. The RAG part is crucial because it means Vanna can retrieve relevant information (like your database schema or example queries) to help it generate better SQL.

Getting started with vanna-ai/vanna is usually quite straightforward. Here's a general outline, but always refer to the official Vanna documentation for the most up-to-date and specific instructions.

Installation

You'll typically install Vanna using pip

pip install vanna

You might also need to install specific database connectors depending on the database you're using (e.g., psycopg2 for PostgreSQL, mysql-connector-python for MySQL, sqlite3 is usually built-in).

Choosing a Model and Database

Vanna needs to know two main things

Which LLM to use
This is what translates your natural language to SQL. Vanna supports various LLMs, including OpenAI's models, Hugging Face models, and potentially local models.

Which database to connect to
This is where Vanna will execute the generated SQL.

Training Vanna (Crucial for Accuracy)

This is where the RAG part comes in! For Vanna to be truly effective, you need to "train" it on your database schema and provide it with some example questions and SQL queries. This helps the LLM understand your specific database context.

You can train Vanna by

Providing DDL (Data Definition Language)
This includes CREATE TABLE statements, which give Vanna the schema of your tables, columns, and relationships.

Adding documentation
You can provide natural language descriptions of your tables and columns.

Supplying example questions and SQL pairs
This is very powerful. If you show Vanna "What are the total sales for product X?" and the corresponding SELECT SUM(sales) FROM orders WHERE product = 'X';, it learns to associate that type of question with that type of query.

Let's look at a simplified Python example. Keep in mind that specific API calls might vary slightly with Vanna updates, so always check the official docs.

import vanna as vn
import pandas as pd # Often useful for displaying results

# --- 1. Initialize Vanna with an LLM and Database ---
# Replace with your actual API key and database connection details
# For demonstration, let's assume we're using a local SQLite database and a mock LLM.
# In a real scenario, you'd use a real LLM like OpenAI's.

# Mock LLM for local testing (Vanna has built-in mock LLMs for quick setup)
# For a real application, you would configure an actual LLM, e.g.:
# from vanna.openai import OpenAI
# vn = OpenAI(api_key="YOUR_OPENAI_API_KEY", model="gpt-4")

# Using the built-in `vanna.chromadb` for vector storage and `vanna.openai` for LLM
# This is a common setup if you're using OpenAI's models.
# For simplicity, let's assume a basic Vanna setup for this example.
# We'll use a placeholder for a 'real' Vanna instance connected to an LLM.

# Placeholder Vanna instance (you would replace this with your actual Vanna setup)
class MyVanna(vn.VannaBase):
    def __init__(self, config=None):
        super().__init__(config=config)
        # In a real scenario, you'd initialize your LLM here
        # For this example, we'll just simulate the behavior

    def generate_sql(self, question: str, auto_train: bool = False, **kwargs) -> str:
        # Simplified mock SQL generation for demonstration
        if "sales" in question.lower() and "product a" in question.lower():
            return "SELECT SUM(quantity * price) FROM orders WHERE product_name = 'Product A';"
        elif "customers" in question.lower() and "new york" in question.lower():
            return "SELECT * FROM customers WHERE city = 'New York';"
        return "SELECT 'Sorry, I couldn\'t generate SQL for that.' AS result;"

    def run_sql(self, sql: str) -> pd.DataFrame:
        # In a real scenario, you would connect to your database and execute the SQL.
        # For this example, we'll return mock data.
        if "orders" in sql.lower():
            return pd.DataFrame({'product_name': ['Product A', 'Product B'], 'quantity': [100, 150], 'price': [10.0, 20.0]})
        elif "customers" in sql.lower():
            return pd.DataFrame({'customer_id': [1, 2], 'name': ['Alice', 'Bob'], 'city': ['New York', 'Los Angeles']})
        return pd.DataFrame({'result': ['Mock data for: ' + sql]})

    def train(self, ddl: str = None, documentation: str = None, question: str = None, sql: str = None):
        print(f"Simulating training: DDL={ddl}, Docs={documentation}, Q={question}, SQL={sql}")
        # In a real Vanna setup, this would store information for RAG.

vanna_instance = MyVanna() # Initialize our mock Vanna

# --- 2. Train Vanna (Crucial Step!) ---

# Provide DDL to give Vanna context about your database schema
vanna_instance.train(ddl="""
CREATE TABLE customers (
    customer_id INT PRIMARY KEY,
    name VARCHAR(255),
    email VARCHAR(255),
    city VARCHAR(255)
);

CREATE TABLE orders (
    order_id INT PRIMARY KEY,
    customer_id INT,
    product_name VARCHAR(255),
    quantity INT,
    price DECIMAL(10, 2),
    order_date DATE,
    FOREIGN KEY (customer_id) REFERENCES customers(customer_id)
);
""")

# Add some documentation for better understanding
vanna_instance.train(documentation="""
The 'orders' table contains information about customer purchases.
'quantity' is the number of units purchased. 'price' is the price per unit.
""")

# Provide example question/SQL pairs (very effective for training!)
vanna_instance.train(question="What are the total sales for 'Product A'?",
                     sql="SELECT SUM(quantity * price) FROM orders WHERE product_name = 'Product A';")

vanna_instance.train(question="List all customers from New York.",
                     sql="SELECT * FROM customers WHERE city = 'New York';")

# --- 3. Ask Questions! ---

# Question 1: Get sales for a specific product
user_question_1 = "What were the total sales for Product A?"
print(f"\nUser Question: {user_question_1}")

# Generate SQL
sql_query_1 = vanna_instance.generate_sql(question=user_question_1)
print(f"Generated SQL: {sql_query_1}")

# Run SQL and get results
if sql_query_1:
    results_1 = vanna_instance.run_sql(sql=sql_query_1)
    print("Results:")
    print(results_1)
else:
    print("Could not generate SQL for this question.")

# Question 2: List customers from a city
user_question_2 = "Show me all customers residing in New York."
print(f"\nUser Question: {user_question_2}")

# Generate SQL
sql_query_2 = vanna_instance.generate_sql(question=user_question_2)
print(f"Generated SQL: {sql_query_2}")

# Run SQL and get results
if sql_query_2:
    results_2 = vanna_instance.run_sql(sql=sql_query_2)
    print("Results:")
    print(results_2)
else:
    print("Could not generate SQL for this question.")

# Question 3: A question not explicitly trained on
user_question_3 = "How many products were sold in total?"
print(f"\nUser Question: {user_question_3}")

sql_query_3 = vanna_instance.generate_sql(question=user_question_3)
print(f"Generated SQL: {sql_query_3}")

if sql_query_3:
    results_3 = vanna_instance.run_sql(sql=sql_query_3)
    print("Results:")
    print(results_3)
else:
    print("Could not generate SQL for this question.")

Explanation of the Example

MyVanna Class (Mock)
In a real Vanna setup, you'd instantiate a Vanna class that connects to your chosen LLM (e.g., vanna.openai.OpenAI) and your database connector. Here, I've created a simplified MyVanna class to simulate the generate_sql and run_sql methods for demonstration purposes without requiring actual API keys or database connections.

train() Method
This is where you feed Vanna information about your database.

ddl
Provides the structural definition of your tables. This is critical for Vanna to understand column names, data types, and relationships.

documentation
Natural language descriptions can help Vanna understand the meaning of your data.

question and sql
These are powerful. By giving Vanna examples of a natural language question and its corresponding SQL, you teach it how to translate similar queries.

generate_sql()
You pass your natural language question to this method. Vanna uses its trained knowledge and the LLM to generate the most appropriate SQL query.

run_sql()
Once you have the SQL, you execute it against your database (simulated here) to get the actual results.

Output
Vanna will show you the generated SQL and then the results from executing that SQL.

Security
Be extremely mindful of security, especially if Vanna is exposed to end-users. Ensure proper authentication and authorization. Generated SQL queries should be carefully reviewed, and parameterization should be used to prevent SQL injection.

Performance
Generating SQL via LLMs can have latency. For highly performance-critical operations, direct SQL might still be preferred.

Cost
LLM API calls can incur costs. Monitor your usage.

Accuracy and Hallucination
While Vanna aims for high accuracy, LLMs can sometimes "hallucinate" or generate incorrect SQL. The quality of your training data (DDL, documentation, examples) is paramount to mitigate this. You might want to implement a review step for generated queries in production.

Integration
Vanna can be integrated into various applications – command-line tools, web dashboards, Slack bots, etc. Its Pythonic interface makes this relatively easy.

Customization
Vanna is designed to be extensible. You can potentially swap out different LLMs, vector databases (for RAG), and UI components.