Posted inAI

Building a Production-Ready Text-to-SQL Engine with Vanna.ai and Python

AI tutorial - IT technology blog
AI tutorial - IT technology blog

The Problem with Ad-Hoc Queries and Generic LLMs

Most engineering teams run into the same friction. Product managers, support leads, and executives need quick numbers for weekly syncs.

Within months, senior engineers spend 4 to 6 hours every week writing throwaway SQL queries instead of shipping core features. Handing non-technical teammates direct access to a database client like DBeaver rarely works. Without deep SQL knowledge, users struggle to join tables correctly, and granting direct access to production replicas invites serious stability and compliance risks.

Plugging off-the-shelf foundation models straight into a database creates a different set of headaches. Feed GPT-4o a prompt like “You are a Postgres expert, answer this user question with SQL”, and cracks appear immediately.

Base models have no concept of your specific business semantics. They hallucinate non-existent column names, confuse legacy soft-delete flags, and output syntactically valid queries that calculate metrics wrong. Stuffing an entire 80-table schema into a prompt context window also wastes tens of thousands of tokens per call while confusing the model with irrelevant table definitions.

Core Concepts: How Vanna.ai Bridges Natural Language and SQL

Vanna.ai solves this by treating Text-to-SQL as an information retrieval problem rather than raw generative guessing. Instead of sending your entire catalog on every API call, Vanna isolates your domain knowledge into indexed vector embeddings. When a question arrives, it pulls only the relevant schema fragments and sample queries before building the final generation prompt.

The Three Layers of Context

Clean queries require domain-specific reference material. Vanna organizes your training data across three distinct layers:

  • DDL (Data Definition Language): Exact CREATE TABLE statements complete with primary keys, foreign keys, default values, and column constraints.
  • Documentation: Plain-language business glossaries and domain quirks (for instance: “Active MRR includes accounts where status = ‘active’ and billing_plan != ‘trial'”).
  • Golden Query Pairs: Curated questions paired with human-audited SQL. These examples serve as few-shot demonstrations, anchoring the model to your team’s preferred dialect and join patterns.

Modular Architecture

Vanna keeps its components decoupled, letting you assemble your production stack like building blocks:

  1. Vector Store: ChromaDB, FAISS, Pinecone, Qdrant, or PGVector for storing and querying your semantic context.
  2. Language Model: Hosted endpoints like OpenAI (GPT-4o), Anthropic Claude, and Google Gemini, or self-hosted models running through Ollama and vLLM.
  3. Database Driver: Native connectors for PostgreSQL, MySQL, Snowflake, BigQuery, SQLite, and ClickHouse.

Hands-On Practice: Building Your Custom Text-to-SQL Pipeline

You can assemble a working local prototype in about fifteen minutes. Here is the architecture I deploy when building internal query assistants for data teams.

Step 1: Environment Setup

Install the core package along with local vector storage and database drivers:

pip install "vanna[chromadb,openai]" psycopg2-binary

Step 2: Initializing the Vanna Client

Vanna uses Python mixins to pair your chosen vector database with an LLM provider. This setup creates a persistent local ChromaDB instance on disk and connects directly to the OpenAI API:

import os
from vanna.openai.openai_chat import OpenAI_Chat
from vanna.chromadb.chromadb_vector import ChromaDB_VectorStore

class CustomVanna(ChromaDB_VectorStore, OpenAI_Chat):
    def __init__(self, config=None):
        ChromaDB_VectorStore.__init__(self, config=config)
        OpenAI_Chat.__init__(self, config=config)

vn = CustomVanna(config={
    'api_key': os.environ.get("OPENAI_API_KEY"),
    'model': 'gpt-4o',
    'path': './vanna_chroma_db'  # Persistent directory for local embeddings
})

Step 3: Training the Pipeline with Schema and Context

Real-world accuracy hinges on this step. Indexing raw DDL is rarely enough on its own. Adding business rules and vetted queries lifts generation accuracy from roughly 60% up into the 90%+ range on complex relational schemas.

# 1. Train with DDL statements
vn.train(ddl="""
CREATE TABLE users (
    id SERIAL PRIMARY KEY,
    email VARCHAR(255) UNIQUE NOT NULL,
    created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW(),
    is_active BOOLEAN DEFAULT TRUE
);

CREATE TABLE orders (
    id SERIAL PRIMARY KEY,
    user_id INT REFERENCES users(id),
    amount NUMERIC(10, 2) NOT NULL,
    status VARCHAR(50) NOT NULL,
    created_at TIMESTAMP WITH TIME ZONE DEFAULT NOW()
);
""")

# 2. Train with business rules and glossary
vn.train(documentation="""
'Completed revenue' refers only to orders where status = 'paid'. 
Orders with status 'pending' or 'refunded' must always be excluded from revenue aggregations.
""")

# 3. Train with vetted question-to-SQL pairs
vn.train(
    question="Who are the top 5 customers by total spending?",
    sql="""
    SELECT u.id, u.email, SUM(o.amount) AS total_spent
    FROM users u
    JOIN orders o ON u.id = o.user_id
    WHERE o.status = 'paid'
    GROUP BY u.id, u.email
    ORDER BY total_spent DESC
    LIMIT 5;
    """
)

Step 4: Connecting to the Database and Running Queries

With context indexed, connect Vanna to your database instance to convert natural language into executed pandas DataFrames:

vn.connect_to_postgres(
    host='localhost',
    dbname='analytics_db',
    user='readonly_user',
    password=os.environ.get("DB_PASSWORD"),
    port=5432
)

# Ask a natural language question
question = "How much revenue did we generate last month from active users?"
sql = vn.generate_sql(question=question)
print("Generated SQL:\n", sql)

# Execute and inspect the resulting DataFrame
df = vn.run_sql(sql=sql)
print("Query Results:\n", df.head())

Step 5: Hardening Security in Production

Connecting an LLM directly to database execution warrants defensive design. Put these guardrails in place before exposing your agent to internal Slack channels or web UIs:

  • Dedicated Read-Only User: Create an isolated database role restricted strictly to SELECT privileges. Explicitly revoke write and DDL permissions across all production schemas.
  • AST Query Validation: Pass generated queries through an AST parser like sqlglot or sqlparse. Reject execution immediately if the syntax tree detects mutation keywords like DROP, TRUNCATE, UPDATE, or ALTER.
  • Strict Query Timeouts: Enforce hard session timeouts on your database role (such as SET statement_timeout = '3000ms'). This stops unindexed table scans or accidental Cartesian joins from degrading database performance.
  • Row Limits: Enforce an explicit ceiling on returned rows (e.g., LIMIT 1000) to prevent out-of-memory crashes when a user asks an open-ended question.

Conclusion: Making Self-Service Data Reliable

Text-to-SQL systems succeed or fail based on the context you feed them. By pairing vector retrieval with explicit business definitions, Vanna strips out the hallucinations that undermine basic LLM wrappers. Start modest: index three or four high-traffic tables, document critical calculation rules, and add verified queries whenever teammates request a custom report. You hand stakeholders instant answers while saving your engineering team hours of tedious report generation every week.

Share:
Tags: