In many data-driven companies, the main challenge isn’t a lack of data. Instead, it’s the gap between people who have questions and the engineers who write SQL queries. Most business data is stored in structured databases, so whenever a manager needs a quick report, an engineer has to pause their work to write a SELECT statement. In this guide, we’ll tackle that problem by building an AI SQL Assistant with LangChain.
This project is more than just interesting; it’s a practical way to make data more accessible using open-source, local Large Language Models (LLMs) without any extra cost for API credits.
AI SQL Assistant with LangChain
An AI SQL assistant works like a translator. You can ask it a question in plain language, such as “What was our best-selling product last month?” It then checks your database schema and creates a correct SQL query for you.
When you use LangChain for this task, three main steps happen behind the scenes:
- Schema Contextualization: The assistant reads your table and column names, along with their data types.
- Query Generation: The LLM uses that schema to write the SQL.
- Execution & Interpretation: The chain executes the query against your database and returns the result in a readable format.
I’ve seen teams try to build this by hard-coding every possible question, but that approach doesn’t work well. LLMs are good at understanding how things are connected, as long as you give them the right metadata.
Make sure to install these libraries before moving forward:
pip install langchain
pip install langchain-community
pip install langchain-experimental
pip install langchain-ollama
pip install streamlit
Be sure to download and install Ollama. Once that’s done, you can pull the CodeLlama model:
ollama pull codellama
Let’s get started and build the AI SQL Assistant step by step.
If you want to build more real-world AI applications like this, I’ve covered it step-by-step in my book: Hands-On GenAI, LLMs & AI Agents.
Step 1: Setting Up the Foundation
Before building the user interface, we need to make sure the LLM can connect to the database. We’ll use SQLDatabaseChain to set up this connection.
Create a file called setup_database.py. This script acts as a smoke test. It connects to a local SQLite file named sales.db and asks a simple question to check if the connection works:
from langchain_community.utilities import SQLDatabase
from langchain_ollama import OllamaLLM
from langchain_experimental.sql import SQLDatabaseChain
# Connect SQLite database
# Ensure you have a sales.db file in the same directory
db = SQLDatabase.from_uri("sqlite:///sales.db")
# Load Ollama model
# We use temperature=0 because we want deterministic, precise SQL, not creative writing
llm = OllamaLLM(
model="codellama",
temperature=0
)
# Create SQL chain
db_chain = SQLDatabaseChain.from_llm(
llm=llm,
db=db,
verbose=True,
return_direct=True
)
# User question
question = "Which product generated the highest revenue?"
# Generate and execute SQL query
response = db_chain.invoke({
"query": question
})
# Print final response
print("\nFinal Response:")
print(response)Notice that the temperature is set to 0. For creative writing, you want the model to explore, but for SQL, accuracy is key. Even a little creativity can cause the model to invent a column name that breaks your query.
Step 2: Refining the Assistant Logic
After the connection is working, put the logic into a reusable script. Create a file named sql_assistant.py:
from langchain_community.utilities import SQLDatabase
from langchain_ollama import OllamaLLM
from langchain_experimental.sql import SQLDatabaseChain
db = SQLDatabase.from_uri("sqlite:///sales.db")
llm = OllamaLLM(
model="codellama",
temperature=0
)
db_chain = SQLDatabaseChain.from_llm(
llm=llm,
db=db,
verbose=True,
return_direct=True
)
question = "Which product generated the highest revenue?"
response = db_chain.invoke(question)
print(response)> Entering new SQLDatabaseChain chain...
Which product generated the highest revenue?
SQLQuery:Question: Which product generated the highest revenue?
SQLQuery: SELECT product, SUM(revenue) AS total_revenue FROM sales GROUP BY product ORDER BY total_revenue DESC LIMIT 1;
SQLResult: [('Laptop', 1200)]
> Finished chain.
This script uses the same approach but focuses on the main interaction.
Step 3: Building the Interactive UI
A terminal script works well for engineers, but the real benefit comes when non-technical users can use it too. Create a file named app.py:
import streamlit as st
from langchain_community.utilities import SQLDatabase
from langchain_ollama import OllamaLLM
from langchain_experimental.sql import SQLDatabaseChain
# Page Config
st.set_page_config(page_title="AI SQL Assistant", layout="centered")
# Connect SQLite database
db = SQLDatabase.from_uri("sqlite:///sales.db")
# Load Ollama model
llm = OllamaLLM(
model="codellama",
temperature=0
)
# Create SQL chain
db_chain = SQLDatabaseChain.from_llm(
llm=llm,
db=db,
verbose=False, # Set to False for a cleaner UI
return_direct=True
)
# Streamlit UI
st.title("📊 AI SQL Assistant")
st.markdown("Ask questions about your sales data in plain English.")
question = st.text_input(
"Ask a database question",
placeholder="e.g., How many units of 'Widget A' were sold?"
)
if question:
with st.spinner("Analyzing database..."):
try:
# Generate and execute SQL query
response = db_chain.invoke({
"query": question
})
st.subheader("AI Response")
st.success(f"Result: {response}")
except Exception as e:
st.error(f"I encountered an error: {e}")
st.info("Tip: Make sure your question refers to tables or columns that exist in the database.")With Streamlit, you can build a simple interface that lets anyone type a question and get an answer based on your data.
Here’s the output you will see with an example query:
In a production environment, this tool is a lifesaver for internal dashboards. It’s perfect for:
- Sales Teams: Checking monthly targets without waiting for a BI developer.
- Inventory Managers: Quickly identifying low-stock items.
- Data Analysts: Rapidly prototyping complex queries before refining them.
Never give an AI Assistant read/write or delete permissions on a production database. Use a Read-Only user. You don’t want a “Delete all users” question to actually execute.
Closing Thoughts
Building an AI SQL Assistant is really about creating a good interface, not just the AI itself. We’re moving toward a time when it’s easier than ever to turn a question into a data insight.
Often, good AI engineering is really just good data engineering. Clear and consistent data will always work better than a clever prompt. Keep building, stay curious, and let machines handle the routine queries.
I hope you found this article on building an AI SQL Assistant with LangChain helpful.
For more AI and machine learning tips, follow me on Instagram. My book, Hands-On GenAI, LLMs & AI Agents, can also help you grow your AI career.
