Building a Q&A style SQL bot

Unless you have been living under a rock, its hard not to notice all the hype surrounding Generative AI. OpenAI, DeepSeek, Anthropic, Claude, Perplexity, Stable Diffusion all have become household names . YouTube and X (formerly Twitter) have seen a siege of AI related content or AI generated content, and it can be quite overwhelming to make sense of it all. While everyone on the AI hype train seems to be debating OpenAI O3 vs DeepSeek R1, I find myself trying to figure out the fundamentals.

Move Fast and Break Things

While I have been trying to understand the maths behind the transformer model that is foundation of all Language Models, I found myself asking “ What can I build that will be useful for my daily use? ” For those in analytics SQL is our bread and butter, which is why I thought it would be cool to build a bot that can answer SQL queries based on user input. A local LLM that will take my question as a natural language input and then based on the context provide me with custom SQL query. Essentially, this would a robust text-to-SQL solution that will generate complex queries.

This led to me spending my Sunday afternoon building a SQL bot that I can execute from my local Jupyter notebook. Below is a straight forward guide to build using completely open source tools. Before jumping in, I would recommend you to read up on the following concepts to gain a holistic understanding.

1. Chat Models

2. Agents

3. Chains

Setup

You can find the list of all packages to be installed, by navigating to the project root and running the following commands.

git clone https://github.com/alphadi/langchain.git

pip install -r requirements.txt

Architecture

Following are the high level steps

• Converting question to a SQL Query: Application converts natural language to query.

• Execute SQL query

• Provide the Answer: Application responds to user input using the query results.

Sample Data

We will be using LangChain, a framework for building applications powered by language models. Simply put, its makes working with AI models much easier. Additionally, we will be using SQLite connection with the Chinook database, a sample dataset that represents a digital media store. You can download and build the database via the command line.

curl -s https://raw.githubusercontent.com/lerocha/chinook-database/master/ChinookDatabase/DataSources/Chinook_Sqlite.sql | sqlite3 Chinook.db

Now that we have downloaded the .db file, we can use LangChain to interface with it.

if not os.environ.get("LANGSMITH_API_KEY"):

os.environ["LANGSMITH_API_KEY"] = getpass.getpass()

os.environ["LANGSMITH_TRACING"] = "true"

from langchain_community.utilities import SQLDatabase

db = SQLDatabase.from_uri("sqlite:///Chinook.db")

print(db.dialect)

print(db.get_usable_table_names())

db.run("SELECT * FROM Artist LIMIT 10;")

Above code should output 10 records

sqlite

['Album', 'Artist', 'Customer', 'Employee', 'Genre', 'Invoice', 'InvoiceLine', 'MediaType', 'Playlist', 'PlaylistTrack', 'Track']

"[(1, 'AC/DC'), (2, 'Accept'), (3, 'Aerosmith'), (4, 'Alanis Morissette'), (5, 'Alice In Chains'), (6, 'Antônio Carlos Jobim'), (7, 'Apocalyptica'), (8, 'Audioslave'), (9, 'BackBeat'), (10, 'Billy Cobham')]"

Awesome! now that our SQL database ready we will connect it to a LLM.

Chains

Now let’s get to the meat of the problem, chains are a composition of predictable steps, which we can use to create a sequence, that given a question will do the following actions.

• Convert natural language to SQL query.

• Run the query.

• Based on the output, answers the original question.

For this, we will be using LangGraph, a LangChain extension that helps us build multi entity applications with LLM. Word to the wise, building such a system will execute a SQL query for even the most mundane user inputs for e.g. “Ssup?” . It is also possible that some inputs will require multiple queries, it is upto us to address each scenario and ensure that we can get relevant answers.

Transform Question to SQL Query

Awesome! let’s get to the first step. We’’ll make use of LangChain’s structured output model module to obtain our SQL query. Below example uses OpenAI but feel free to use the chat model of your choice.

import getpass

import os

if not os.environ.get("OPENAI_API_KEY"):

os.environ["OPENAI_API_KEY"] = getpass.getpass("Enter API key for OpenAI: ")

from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model="gpt-4o-mini")

Any model needs a good instructor. Lucky for us, LangChain has a query prompt template which makes our job easy.

from langchain import hub

query_prompt_template = hub.pull("langchain-ai/sql-query-system-prompt")

assert len(query_prompt_template.messages) == 1

query_prompt_template.messages[0].pretty_print()

Now that our template is set, we will need to populate parameters such as the SQL dialect and table schema. We will use LangChain SQLDatabase object to create a function that will do just that.

from typing_extensions import TypedDict

from typing_extensions import Annotated

class State(TypedDict):

question: str

query: str

result: str

answer: str

class QueryOutput(TypedDict):

"""Generated SQL query."""

query: Annotated[str, ..., "Syntactically valid SQL query."]

def write_query(state: State):

"""Generate SQL query to fetch information."""

prompt = query_prompt_template.invoke(

{

"dialect": db.dialect,

"top_k": 10,

"table_info": db.get_table_info(),

"input": state["question"],

}

)

structured_llm = llm.with_structured_output(QueryOutput)

result = structured_llm.invoke(prompt)

return {"query": result["query"]}

Now lets see if our bot can generate queries

write_query({"question": "How many Employees are there?"})

Kudos! ask and you shall receive, in SQL :

{'query': 'SELECT COUNT(EmployeeId) AS EmployeeCount FROM Employee;'}

Transform Question to SQL Query

Now that we have got that out of our way let’s see if we can get our bot to also run the query. Essentially, we are building a custom SQL chain. Here is where we need to be careful, and not make unwanted changes to our dataset by running automated queries. I would recommend keep a tight control on database connector permissions and minimizing write privileges.

from langchain_community.tools.sql_database.tool import QuerySQLDatabaseTool

def execute_query(state: State):

"""Execute SQL query."""

execute_query_tool = QuerySQLDatabaseTool(db=db)

return {"result": execute_query_tool.invoke(state["query"])}

Moment of truth, let’s see if our chain works

execute_query({"query": "SELECT COUNT(EmployeeId) AS EmployeeCount FROM Employee;"})

Cool ! we got our query and our answer :

{'result': '[(8,)]'}

Now to wrap it all up nicely

def generate_answer(state: State):

"""Answer question using retrieved information as context."""

prompt = (

"Given the following user question, corresponding SQL query, "

"and SQL result, answer the user question.\n\n"

f'Question: {state["question"]}\n'

f'SQL Query: {state["query"]}\n'

f'SQL Result: {state["result"]}'

)

response = llm.invoke(prompt)

return {"answer": response.content}

That was quite a lot to cover, it can be pretty confusing at first but with practice you can wrap your head around it. It helps to visualize the system.

We have everything set, let’s test our application. We will compile our application by orchestrating with LangGraph. Not to worry, we are simply compiling the three steps into a single sequence.

from langgraph.graph import START, StateGraph

graph_builder = StateGraph(State).add_sequence(

[write_query, execute_query, generate_answer]

)

graph_builder.add_edge(START, "write_query")

graph = graph_builder.compile()

for step in graph.stream(

{"question": "How many employees are there?"}, stream_mode="updates"

):

print(step)

Since we are getting the bot to answer and execute queries, we should be careful. It is prudent to interrupt our bot before the execution of a query for human review. LangGraph provides a persistence layer which enables us to review the query before providing the approval to execute it. LangChain calls it human in the loop, check documentation for more details.

from langgraph.checkpoint.memory import MemorySaver

memory = MemorySaver()

graph = graph_builder.compile(checkpointer=memory, interrupt_before=["execute_query"])

# Now that we're using persistence, we need to specify a thread ID

# so that we can continue the run after review.

config = {"configurable": {"thread_id": "1"}}

for step in graph.stream(

{"question": "How many employees are there?"},

config,

stream_mode="updates",

):

print(step)

try:

user_approval = input("Do you want to go to execute query? (yes/no): ")

except Exception:

user_approval = "no"

if user_approval.lower() == "yes":

# If approved, continue the graph execution

for step in graph.stream(None, config, stream_mode="updates"):

print(step)

else:

print("Operation cancelled by user.")

If you would like to visualize the metadata associated you can use trace.

This is a pretty neat bot which can answer simple queries. If you want to build more complex query creation, we can leverage few shot prompting. Here is a list of resources:

1. Advanced Prompting strategies

2. Query Validation

Agents

Language models cannot take actions, they can only output text. Agents are tools that can understand a high level task and leverage a language model as a reasoning brain to decide what action to take and execute. LangGraph is a great tool to build resourceful and customizable agents. Here’s a few reasons why agents are suited for this task better than chains.

• Agents can answer the user question by querying the database multiple times.

• Agents can ‘understand’ errors from wrong queries, then traceback and regenerating correct queries.

• Agents can answer user questions based on database schema for e.g. describing a specific table.

We can use … you guessed it, LangChain to equip the agent with a set of tools to create and execute queries, check query syntax, retrieve description and more..

from langchain_community.agent_toolkits import SQLDatabaseToolkit

toolkit = SQLDatabaseToolkit(db=db, llm=llm)

tools = toolkit.get_tools()

tools

Let’s load a system prompt for our agent, this will have the instructions for the agent.

from langchain import hub

prompt_template = hub.pull("langchain-ai/sql-agent-system-prompt")

assert len(prompt_template.messages) == 1

prompt_template.messages[0].pretty_print()

We will populate the parameters like how we used in Chains

system_message = prompt_template.format(dialect="SQLite", top_k=5)

For simplicity, we will use a pre-built LangGraph agent to customize our agent.

from langchain_core.messages import HumanMessage

from langgraph.prebuilt import create_react_agent

agent_executor = create_react_agent(llm, tools, prompt=system_message)

question = "Which country's customers spent the most?"

for step in agent_executor.stream(

{"messages": [{"role": "user", "content": question}]},

stream_mode="values",

):

step["messages"][-1].pretty_print()

You can see how the agent runs with detailed steps. The agent can execute multiple queries until the user has the information. The agent is able to use the final query result and generate an answer. Our agent is also to able to handle qualitative questions.

question = "Describe the playlisttrack table"

for step in agent_executor.stream(

{"messages": [{"role": "user", "content": question}]},

stream_mode="values",

):

step["messages"][-1].pretty_print()

Conclusion

To wrap up, we built a chat application that uses language models to connect with SQL database via LangChain framework. For a user friendly interface we can build and deploy a custom application using Streamlit, that delivers complex queries based on conversational exchanges. It’s completely open source and works great .

There are other ways to enhance the bot’s capabilities, you can try and compare results with different chat models. This could make the system much better at working with databases, making it better at figuring out complex queries.