How to Build a Custom AI Agent Using LangGraph with RAG, NL2SQL and Web Search

How to Build a Custom AI Agent Using LangGraph with RAG, NL2SQL and Web Search

Rounak Show's photo
·

10 min read

Introduction

In this blog post, we will walk you through the process of creating a custom AI agent with three powerful tools: Web Search, Retrieval-Augmented Generation (RAG), and Natural Language to SQL (NL2SQL), all integrated within the LangGraph framework. This guide is designed to provide you with a practical, step-by-step approach to building a fully functional AI agent capable of performing complex tasks such as retrieving real-time data from the web, generating responses based on retrieved information from the knowledge base, and translating natural language queries into SQL database queries. By the end of this tutorial, you will have a working AI agent equipped to handle these diverse functionalities seamlessly.

In the previous LangGraph Tutorial blog, we explored the foundational concepts behind LangGraph, focusing on creating simple AI agents by defining nodes and edges. This time, we’re taking things a step further. We will incorporate additional tools that extend the agent’s capabilities—allowing it to interact with the web, retrieve specific information from knowledge-base, and query SQL databases through natural language commands. This comprehensive guide is aimed at developers looking to build customizable, task-oriented AI agents that can be deployed in various real-world applications.

We have already built a similar AI Agent using OpenAI Swarm. This time, we are going to do this using LangGraph.

Let’s dive into the setup process and start by building the environment for our AI agent!

LangGraph Agent with WebSearch, RAG, NL2SQL Tool

Setting up the Environment

To get started, you need to install the below dependencies

!pip install -U langchain langchain-chroma pypdf sentence-transformers langgraph langchain_openai langchain_community

Setting up API Keys

Before diving into building your LangGraph AI agent, it's crucial to set up your API keys. These keys allow your agent to interact with external tools like Tavily Search and OpenAI GPT models securely. Without them, the tools cannot function effectively.

Here we are using OpenAI model, but you can use any LLM of your choice.

import getpass
import os

def _set_env(var: str):
    if not os.environ.get(var):
        os.environ[var] = getpass.getpass(f"{var}: ")

_set_env("TAVILY_API_KEY")
_set_env("OPENAI_API_KEY")

Creating the LLM Object

Here’s how to initialize the LLM using LangChain ChatOpenAI

from langchain_openai import ChatOpenAI

llm = ChatOpenAI(model_name="gpt-4o")

WebSearch Tool Integration

Integrating a WebSearch tool allows your AI agent to fetch real-time information from the web. At FutureSmart AI, we generally use TavilySearch tool for quick and accurate web results. But, you can use any web search tool of your choice.

from langchain_community.tools.tavily_search import TavilySearchResults

web_search_tool = TavilySearchResults(max_results=2)

Agentic RAG Tool Integration

The Retrieval-Augmented Generation (RAG) tool enhances your AI agent by enabling it to fetch relevant documents to the user's questions. As a result, users will get more accurate and context-rich responses.

Steps to Integrate the RAG Tool

  1. Load Documents: Use the PyPDFLoader and Docx2txtLoader to load documents from a folder.

     from langchain_community.document_loaders import PyPDFLoader, Docx2txtLoader
    
     def load_documents(folder_path: str) -> List[Document]:
         documents = []
         for filename in os.listdir(folder_path):
             file_path = os.path.join(folder_path, filename)
             if filename.endswith('.pdf'):
                 loader = PyPDFLoader(file_path)
             elif filename.endswith('.docx'):
                 loader = Docx2txtLoader(file_path)
             else:
                 print(f"Unsupported file type: {filename}")
                 continue
             documents.extend(loader.load())
         return documents
    
     folder_path = "/content/docs"
     documents = load_documents(folder_path)
     print(f"Loaded {len(documents)} documents from the folder.")
    
     # Output
     Loaded 2 documents from the folder.
    
  2. Split Text into Chunks: Prepare documents for vectorization by splitting them into manageable chunks.

     text_splitter = RecursiveCharacterTextSplitter(
         chunk_size=1000,
         chunk_overlap=200,
         length_function=len
     )
    
     splits = text_splitter.split_documents(documents)
     print(f"Split the documents into {len(splits)} chunks.")
    
     # Output
     Split the documents into 4 chunks.
    
  3. Generate Embeddings: Use SentenceTransformers to create embeddings for efficient similarity searches.

     from langchain_community.embeddings.sentence_transformer import SentenceTransformerEmbeddings
    
     embedding_function = SentenceTransformerEmbeddings(model_name="all-MiniLM-L6-v2")
    
  4. Create and Persist a Vector Store:

     from langchain_chroma import Chroma
    
     collection_name = "my_collection"
     vectorstore = Chroma.from_documents(
         collection_name=collection_name,
         documents=splits,
         embedding=embedding_function,
         persist_directory="./chroma_db"
     )
    
  5. Build the Retriever Tool:

     from langchain.tools import tool
     from pydantic import BaseModel
    
     class RagToolSchema(BaseModel):
         question: str
    
     @tool(args_schema=RagToolSchema)
     def retriever_tool(question):
       """Tool to Retrieve Semantically Similar documents to answer User Questions related to FutureSmart AI"""
       print("INSIDE RETRIEVER NODE")
       retriever = vectorstore.as_retriever(search_kwargs={"k": 2})
       retriever_result = retriever.invoke(question)
       return "\n\n".join(doc.page_content for doc in retriever_results)
    

This tool allows your AI agent to retrieve relevant chunks of information from your document database, making it highly effective for knowledge-based tasks.

You can test the retriever right now to see how it is performing

retriever = vectorstore.as_retriever(search_kwargs={"k": 2})
# pass question
retriever_results = retriever.invoke("Who is the founder of Futuresmart AI?")
print(retriever_results)
# Output
[Document(metadata={'page': 1, 'source': '/content/docs/FutureSmart AI .pdf'}, page_content='FutureSmart AI provides customized speech to text services, employing cutting-\nedge speech recognition technologies to cater to specific client needs. Ideal for \ncreating efficient documentation and enabling voice-driven commands, this \nsolution boosts productivity and accessibility.'), 
Document(metadata={'page': 0, 'source': '/content/docs/FutureSmart AI .pdf'}, page_content='FutureSmart AI provides custom Natural Language Processing (NLP) \nsolutions for companies looking to get ahead of the future. Our \ndedicated team of Data Scientists and ML Engineers provides an end-\nto-end solution from data labeling to modeling and deploying an ML \nmodel tailored to your specific use case. \nFounder: Pradip Nichite \n \nServices: \nText Classification \nAt FutureSmart AI, we develop custom text classification solutions using \nadvanced NLP techniques tailored to your specific business requirements. \nLeveraging Python, Pytorch, and Hugging Face transformers, we enable precise \ndata categorization across applications such as intent detection, document \ncategorization, and sentiment analysis, enhancing your decision-making \nprocesses and operational efficiency. \n \nChatbots \nWe specialize in creating custom chatbots that integrate seamlessly with your \nbusiness environment. Using semantic search and large language models, our')]

If you're interested in building RAG from scratch, this video is for you.

NL2SQL Tool Integration

Now we have Web Search and RAG tool ready for the AI Agent. We only need to build the NL2SQL tool.

The SQL Agent bridges the gap between human language and SQL databases by generating and executing SQL queries from natural language questions. It enables your AI agent to answer database-related queries efficiently.

Steps to Integrate the NL2SQL Tool

  1. Set Up the Database
    We are using the Chinook SQLite database as a sample dataset for testing SQL queries. Check this blog to learn how to set up your own database for testing.

     !wget https://github.com/lerocha/chinook-database/raw/master/ChinookDatabase/DataSources/Chinook_Sqlite.sqlite
     !mv Chinook_Sqlite.sqlite Chinook.db
    
  2. Initialize the Database Connection
    Connect to the SQLite database using the LangChain SQLDatabase utility:

     from langchain_community.utilities import SQLDatabase
    
     db = SQLDatabase.from_uri("sqlite:///Chinook.db")
    
  3. Clean SQL Queries
    This function is very important. We often see this problem in many of our client’s projects that the SQL query generated from the LLM consists of some unnecessary symbols, texts, backticks, etc. As a result, you will get an error while executing this query. Thus, we require a function to clean up these additional texts.

import re

def clean_sql_query(text: str) -> str:
    """
    Clean SQL query by removing code block syntax, various SQL tags, backticks,
    prefixes, and unnecessary whitespace while preserving the core SQL query.

    Args:
        text (str): Raw SQL query text that may contain code blocks, tags, and backticks

    Returns:
        str: Cleaned SQL query
    """
    # Step 1: Remove code block syntax and any SQL-related tags
    # This handles variations like ```sql, ```SQL, ```SQLQuery, etc.
    block_pattern = r"```(?:sql|SQL|SQLQuery|mysql|postgresql)?\s*(.*?)\s*```"
    text = re.sub(block_pattern, r"\1", text, flags=re.DOTALL)

    # Step 2: Handle "SQLQuery:" prefix and similar variations
    # This will match patterns like "SQLQuery:", "SQL Query:", "MySQL:", etc.
    prefix_pattern = r"^(?:SQL\s*Query|SQLQuery|MySQL|PostgreSQL|SQL)\s*:\s*"
    text = re.sub(prefix_pattern, "", text, flags=re.IGNORECASE)

    # Step 3: Extract the first SQL statement if there's random text after it
    # Look for a complete SQL statement ending with semicolon
    sql_statement_pattern = r"(SELECT.*?;)"
    sql_match = re.search(sql_statement_pattern, text, flags=re.IGNORECASE | re.DOTALL)
    if sql_match:
        text = sql_match.group(1)

    # Step 4: Remove backticks around identifiers
    text = re.sub(r'`([^`]*)`', r'\1', text)

    # Step 5: Normalize whitespace
    # Replace multiple spaces with single space
    text = re.sub(r'\s+', ' ', text)

    # Step 6: Preserve newlines for main SQL keywords to maintain readability
    keywords = ['SELECT', 'FROM', 'WHERE', 'GROUP BY', 'HAVING', 'ORDER BY',
               'LIMIT', 'JOIN', 'LEFT JOIN', 'RIGHT JOIN', 'INNER JOIN',
               'OUTER JOIN', 'UNION', 'VALUES', 'INSERT', 'UPDATE', 'DELETE']

    # Case-insensitive replacement for keywords
    pattern = '|'.join(r'\b{}\b'.format(k) for k in keywords)
    text = re.sub(f'({pattern})', r'\n\1', text, flags=re.IGNORECASE)

    # Step 7: Final cleanup
    # Remove leading/trailing whitespace and extra newlines
    text = text.strip()
    text = re.sub(r'\n\s*\n', '\n', text)

    return text

We create this function to process the LLM-generated SQL Query Output. If you have a better one, feel free to share with everyone in the comment 🙂

  1. Create the NL2SQL Tool
    Define a tool that generates and executes SQL queries:

     from langchain.chains import create_sql_query_chain
     from langchain_community.tools.sql_database.tool import QuerySQLDataBaseTool
     from operator import itemgetter
     import re
     from langchain_core.output_parsers import StrOutputParser
     from langchain_core.prompts import PromptTemplate
     from langchain_core.runnables import RunnablePassthrough, RunnableLambda
    
     class SQLToolSchema(BaseModel):
         question: str
    
     @tool(args_schema=SQLToolSchema)
     def nl2sql_tool(question):
       """Tool to Generate and Execute SQL Query to answer User Questions related to chinook DB"""
       print("INSIDE NL2SQL TOOL")
       execute_query = QuerySQLDataBaseTool(db=db)
       write_query = create_sql_query_chain(llm, db)
    
       chain = (
           RunnablePassthrough.assign(query=write_query | RunnableLambda(clean_sql_query)).assign(
               result=itemgetter("query") | execute_query
           )
       )
    
       response = chain.invoke({"question": question})
       return response['result']
    
  2. Test the Tool
    Use a sample query to verify functionality:

     question = "How many employees are there?"
     result = nl2sql_tool.invoke({"question": question})
     print(f"Answer: {result}")
    
     # output
     INSIDE NL2SQL TOOL
     Question: How many employees are there?
     Answer: [(8,)]
    

Combining the Tools

To create a fully functional LangGraph AI agent, you need to integrate the WebSearch, RAG, and NL2SQL tools into the LLM. And then you will get a search agent, a rag agent, and an SQL agent - all in one.

Here’s how to combine the tools:

tools = [web_search_tool, retriever_tool, nl2sql_tool]
llm_with_tools = llm.bind_tools(tools)

With this configuration, your LLM can invoke the appropriate tool based on the user’s query, ensuring dynamic and context-aware responses.

Building the LangGraph

LangGraph enables you to define a stateful workflow for your AI agent. By structuring nodes and edges, you can define how the agent processes user inputs and transitions between tools.

Steps to Build the LangGraph

  1. Define the State: Create a State dictionary to manage the agent’s inputs and outputs.

     from typing import Annotated
     from langgraph.graph import StateGraph
     from langgraph.graph.message import add_messages
    
     class State(TypedDict):
         messages: Annotated[list, add_messages]
    
  2. Add Nodes: Add nodes for the chatbot and tools to handle user queries and invoke the tools.

     def chatbot(state: State):
         return {"messages": [llm_with_tools.invoke(state["messages"])]}
    
     graph_builder = StateGraph(State)
     graph_builder.add_node("chatbot", chatbot)
    
     tool_node = ToolNode(tools=[web_search_tool, retriever_tool, nl2sql_tool])
     graph_builder.add_node("tools", tool_node)
    
  3. Define Edges: Use conditional edges to determine when the agent should switch between nodes.

     from langgraph.prebuilt import tools_condition
    
     graph_builder.add_conditional_edges("chatbot", tools_condition)
     graph_builder.add_edge("tools", "chatbot")
     graph_builder.set_entry_point("chatbot")
    
  4. Compile the Graph: Finalize the graph for execution.

     graph = graph_builder.compile()
    

Testing the AI Agent

Once the LangGraph is set up, you can test the agent by simulating user inputs. This ensures the tools and workflows are functioning as expected.

Interactive Testing

Run the following code to test your AI agent interactively:

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

while True:
    user_input = input("User: ")
    if user_input.lower() in ["quit", "exit", "q"]:
        print("Goodbye!")
        break

    for event in graph.stream({"messages": [("user", user_input)]}, config):
        for value in event.values():
            print("Assistant:", value["messages"][-1].content)

You can provide queries like:

  • You can ask questions based on your knowledge-base (Trigger Retriever or RAG tool)

  • "What is the current weather in Delhi?" (Trigger Websearch tool)

  • "Tell me the name of any two employees from the database" (Trigger NL2SQL tool)

The AI agent will invoke the appropriate tool to generate responses.

Visualizing the LangGraph

Visualization helps you understand the workflow of your AI agent and how it transitions between nodes and tools.

from IPython.display import Image, display

try:
    display(Image(graph.get_graph().draw_mermaid_png()))
except Exception:
    print("Error generating graph visualization.")

The generated diagram will showcase nodes (chatbot, tools) and the transitions between them, providing a clear overview of your AI agent’s workflow.

💡
Check out the full code on our GitHub

Improvements

  1. You can implement methods like Contextual Retrieval, reranking, etc to improve the performance of the Agentic RAG system.

  2. You can try implementing few-shot examples, dynamic table selection, custom SQL query generation chain, etc to improve the SQL Agent tool. We would recommend watching this Complete Natural Language to SQL video for a better understanding.

  3. Moreover, you can always improve the agent by modifying graph creation, writing better prompts, adding human-in-the-loop, etc.

Conclusion

This concludes our step-by-step guide to creating a custom AI agent using LangGraph with Web Search, RAG, and NL2SQL tools. By following this tutorial, you've built an AI agent capable of performing diverse tasks such as retrieving real-time data, answering questions based on document-based knowledge, and executing SQL queries directly from natural language commands. This hands-on approach has equipped you with practical knowledge of integrating multiple tools into a stateful AI workflow using LangGraph.

If you found this guide useful and want to explore more advanced techniques, don’t forget to check out our other tutorials. At FutureSmart AI, we help businesses develop state-of-the-art AI solutions tailored to their needs. For inquiries, feel free to reach out to us at contact@futuresmart.ai.

For real-world examples of our work, take a look at our case studies, where we showcase the practical value of our expertise.