Text to SQL without sharing database rows

Originally published on Medium on 2024-07-18. This tutorial reflects the libraries and APIs used at the time of publication.

Data privacy matters when connecting language models to company databases. This approach uses LangChain, OpenAI, and SQLAlchemy to generate SQL from a database schema, without sending the database rows to the model.

The model receives table names, column names, data types, and the user's question. It returns a SQL query rather than executing that query or returning database results. The schema and question still leave your application, so they must be appropriate to share with your model provider.

Understanding the Core Technologies

• Retrieval-Augmented Generation (RAG): RAG is an AI framework that enhances language models with external knowledge. It combines information retrieval with text generation. RAG first retrieves relevant information from a knowledge base. Then, it uses this information to augment the input to a language model improving the accuracy and relevance of AI-generated responses.
• LangChain: LangChain is a framework for developing applications powered by language models. It provides a set of tools and components for building complex AI systems. LangChain facilitates the integration of language models with external data sources. It offers modules for prompt management, memory, and chains of operations supporting various use cases like chatbots, Q&A systems, and text summarization.
• SQLAlchemy: A robust SQL toolkit and Object Relational Mapper (ORM) for Python. It facilitates database interaction by abstracting common tasks into Pythonic operations, reducing the need for direct SQL handling.

Normal RAG Implementation

My earlier database integration gave the language-model workflow access to database results. That can be useful when you want a direct answer, but it also means deciding which data the model may receive.

Our approach

Here, the application connects to the database locally and extracts only its schema. It sends that schema with the question to the model, which generates SQL. Review the returned query and enforce database permissions before executing it separately.

The original article includes the architecture diagrams and app screenshots.

Code Explanation

Let’s examine the key components of our privacy-preserving RAG system.

Importing the libraries

from langchain import OpenAI, LLMChain
from langchain.prompts import PromptTemplate
from langchain.utilities import SQLDatabase
from sqlalchemy import create_engine, MetaData, Table, Column, inspect
from langchain_experimental.sql import SQLDatabaseChain
# we kept the temp=0 as we dont want LLM to use creativity and randomness
llm = OpenAI(temperature=0, openai_api_key="your_openai_api_key")

Extracting Database Schema

This function connects to the database, retrieves table and column information, and formats it in a readable manner.

def extract_schema(db_url):
    engine = create_engine(db_url)
    inspector = inspect(engine)
    schema_info = []
    for table_name in inspector.get_table_names():
        columns = inspector.get_columns(table_name)
        schema_info.append(f"Table: {table_name}")
        for column in columns:
            schema_info.append(f"  - {column['name']} ({column['type']})")
    return "\n".join(schema_info)

Creating Prompt Templates with LangChain

LangChain uses a concept called PromptTemplate to structure the interaction with the language model. We pass the entire database schema in the prompt template itself. The Prompt Template guides the model on how to understand the user's input and how to format the output:

prompt_template = """
You are an AI assistant that generates SQL queries based on user requests.
You have access to the following database schema:
{schema}
Based on this schema, generate a SQL query to answer the following question:
{question}
SQL Query:
"""
prompt = PromptTemplate(
    input_variables=["schema", "question"],
    template=prompt_template,
)

Generating SQL Queries

This function passes the user’s question and the database schema to the language model, which then generates the appropriate SQL query.

schema = extract_schema("sqlite:///new.db")  # Replace with your database URL
chain = LLMChain(llm=llm, prompt=prompt)
def generate_sql_query(question):
    return chain.run(schema=schema, question=question)

Example Usage

Let’s see how everything comes together with a practical example:

user_question = "Find me the registration id of the hackathon"
sql_query = generate_sql_query(user_question)
print(f"Generated SQL Query: {sql_query}")

Full Code

You can find the full code to this here: View the code on GitHub

Try Here

I deployed this app on Streamlit and you can try it yourself here using your own DB and Open AI Key: Open the Streamlit app

Download Sample Database for testing out (new.db) : Download the sample database

Drawbacks and Scope of Improvement

Two limitations showed up in this implementation:

1. The model cannot see stored values. In the sample database, a query using hackathon did not match the stored value Hackathon. Whether case matters depends on the database and its comparison rules. Schema alone does not tell the model which values are present.
2. Large schemas consume context. Sending every table and column becomes expensive as the schema grows. A larger context window can help, or you can retrieve only the relevant schema sections. The row count itself does not expand this prompt because rows are not included.

If you like this blog you should also check out the videos I make on Instagram: Instagram

In case of any queries, feel free to reach out to me on paras@varnan.tech