Talk to your Database using RAG and LLMs

5 min readOct 12, 2023

What is RAG?

Retrieval Augmented Generation(RAG), involves the process of acquiring current or context-relevant information from an external database and presenting it to a Large Language Model (LLM) when requesting it to generate a response. This approach effectively addresses the issue of generating incorrect or misleading information. You have the capability to store exclusive business data or global knowledge and enable your application to retrieve this data for the LLM during the response generation phase.

How does it help?

LLMs lack domain knowledge, but we can fill this gap by utilizing RAGs to retrieve context information from the database and pass it along with the user input to the LLM and generate an enriched, relevant response.

In this article, we will look at how to use LangChain to connect to our SQL database, retrieve contextual information, pass the User Query along with the Context to the LLM and generate accurate responses.

Case Overview

We will build a custom QA ChatBot for an Employment Agency that helps users get relevant information on the job market. A user might want to know the trending jobs, popularity of a specific role in a particular area, etc.

We use Langchain as the framework, MySQL database and OpenAI’s LLM to build our app.

Our Database consists of a single table with the following columns:

Let’s start with the implementation.

Installing dependencies

pip install langchain langchain-experimental openai pymysql

Import necessary packages

from langchain.utilities import SQLDatabase
from langchain.llms import OpenAI
from langchain_experimental.sql import SQLDatabaseChain
from langchain.prompts import PromptTemplate
from langchain.prompts.chat import HumanMessagePromptTemplate
from langchain.chat_models import ChatOpenAI
from langchain.schema import HumanMessage, SystemMessage

Initialize the LLM

OPENAI_API_KEY = "your-key-here"
llm = ChatOpenAI(temperature=0, openai_api_key=OPENAI_API_KEY)

We use ChatOpenAI from Langchain which is uses gpt3.5-turbo model from OpenAI.

Database Setup

host = 'localhost'
port = '3306'
username = 'root'
password = 'password'
database_schema = 'agency_db'
mysql_uri = f"mysql+pymysql://{username}:{password}@{host}:{port}/{database_schema}"

db = SQLDatabase.from_uri(mysql_uri, include_tables=['job_details'],sample_rows_in_table_info=2)

db_chain = SQLDatabaseChain.from_llm(llm, db, verbose=True)

We first construct a URI to connect to the database, here I am using pymysql as the driver to connect to mysql, but you can change it according to your db dialect, for example- connection string for a postgresql DB would be exactly the same as given above, you just need to replace “mysql+pymysql” with “postgresql+psycopg2” and install psycopg2 via pip.

When initializing the DB, we use the SQLDatabase.from_uri() method from langchain, and pass the URI along with a few other parameters, here is what they do:

include_tables: We only refer to the job_details table and query over this table only
sample_rows_in_table_info: When interacting with the DB, langchain will receive the top 2 rows of this table as a sample for additional context and better retrieval

The final step in this is to setup a database chain which would help our LLM interact with the database, we do so using SQLDatabaseChain.

Setting up retrieval functions

def retrieve_from_db(query: str) -> str:
    db_context = db_chain(query)
    db_context = db_context['result'].strip()
    return db_context

This function would take the query provided by the user as an input, pass it to our db chain and retrieve the relevant information from out database and return it as a string.

Building the Prompts

We would be using Prompt Templates to keep it clean and understandable. A Prompt Template helps us build a skeletal frame for our prompts and we can fill in our template based on the user input and use case.

system_message = """You are a professional representative of an employment agency.
        You have to answer user's queries and provide relevant information to help in their job search. 
        Example:
        
        Input:
        Where are the most number of jobs for an English Teacher in Canada?
        
        Context:
        The most number of jobs for an English Teacher in Canada is in the following cities:
        1. Ontario
        2. British Columbia
        
        Output:
        The most number of jobs for an English Teacher in Canada is in Toronto and British Columbia
        """
human_qry_template = HumanMessagePromptTemplate.from_template(
        """Input:
        {human_input}
        
        Context:
        {db_context}
        
        Output:
        """
    )
messages = [
      SystemMessage(content=system_message),
      human_qry_template.format(human_input=query, db_context=db_context)
    ]
response = llm(messages).content

Since we are using the chat model, we provide instructions and an example in the system message, making it a one-shot inference task.

In the Human Message, along with the user query, we feed the data we retrieved from our database to the model.

Finally, format the human message template with the query input by the user and the context retrieved from our database and hit the chat model to generate a response.

Piecing everything together

def generate(query: str) -> str:
    db_context = retrieve_from_db(query)
    
    system_message = """You are a professional representative of an employment agency.
        You have to answer user's queries and provide relevant information to help in their job search. 
        Example:
        
        Input:
        Where are the most number of jobs for an English Teacher in Canada?
        
        Context:
        The most number of jobs for an English Teacher in Canada is in the following cities:
        1. Ontario
        2. British Columbia
        
        Output:
        The most number of jobs for an English Teacher in Canada is in Toronto and British Columbia
        """
    
    human_qry_template = HumanMessagePromptTemplate.from_template(
        """Input:
        {human_input}
        
        Context:
        {db_context}
        
        Output:
        """
    )
    messages = [
      SystemMessage(content=system_message),
      human_qry_template.format(human_input=query, db_context=db_context)
    ]
    response = llm(messages).content
    return response

Sample Input and Output

Here, we can see the db chain fetch the most popular job in Toronto and it is passed as context to our LLM giving us the correct answer in Natural Language.

Conclusion

In this article, we explored how Retrieval Augmented Generation (RAG) can be a game-changer when building domain-specific chat bots. RAG bridges the gap between Large Language Models (LLMs) and external databases, ensuring more informed and accurate responses.

We discussed RAG’s ability to compensate for LLMs’ lack of domain knowledge by dynamically retrieving context information from databases. The practical example, a ChatBot for an Employment Agency, demonstrated Langchain’s role in connecting with an SQL database and utilizing OpenAI’s LLM for precise responses.

By adopting RAG and frameworks like Langchain, we pave the way for more context-aware AI applications. This enables better user experiences and decision-making by delivering accurate, context-rich information. Embrace RAG’s potential to unlock the wealth of data-driven conversational AI.