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!
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")
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)
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.
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.")
Loaded 2 documents from the folder.
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.")
Split the documents into 4 chunks.
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")
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"
)
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})
retriever_results = retriever.invoke("Who is the founder of Futuresmart AI?")
print(retriever_results)
[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.
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.
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
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")
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
"""
block_pattern = r"```(?:sql|SQL|SQLQuery|mysql|postgresql)?\s*(.*?)\s*```"
text = re.sub(block_pattern, r"\1", text, flags=re.DOTALL)
prefix_pattern = r"^(?:SQL\s*Query|SQLQuery|MySQL|PostgreSQL|SQL)\s*:\s*"
text = re.sub(prefix_pattern, "", text, flags=re.IGNORECASE)
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)
text = re.sub(r'`([^`]*)`', r'\1', text)
text = re.sub(r'\s+', ' ', text)
keywords = ['SELECT', 'FROM', 'WHERE', 'GROUP BY', 'HAVING', 'ORDER BY',
'LIMIT', 'JOIN', 'LEFT JOIN', 'RIGHT JOIN', 'INNER JOIN',
'OUTER JOIN', 'UNION', 'VALUES', 'INSERT', 'UPDATE', 'DELETE']
pattern = '|'.join(r'\b{}\b'.format(k) for k in keywords)
text = re.sub(f'({pattern})', r'\n\1', text, flags=re.IGNORECASE)
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 š
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']
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}")
INSIDE NL2SQL TOOL
Question: How many employees are there?
Answer: [(8,)]
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
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]
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)
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")
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
You can implement methods like Contextual Retrieval, reranking, etc to improve the performance of the Agentic RAG system.
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.
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.
