Multi Agent

Build a Retrieval-Augmented Generation agent using Chroma and LangGraph

In this section, we will see how to build a RAG (Retrieval-Augmented Generation) agent that can load a PDF document, store embeddings in a vector database, and retrieve info using custom tools.

Objectives

1. Split document pages and create a Chroma database.
2. Wrap database query retrieval in a custom @tool.
3. Set up a StateGraph that invokes LLM reasoning and retrieval tools sequentially.

Agent IV: RAG Agent

Goal

Build a RAG agent querying details from the Stock_Market_Performance_2024.pdf document.

Sample Input

{"messages": [("user", "How did the tech sector perform in 2024?")]}

Sample Output

Outputs summarizing tech sector performance: In 2024, the tech sector gained 25% driven by AI breakthroughs. (cited from the document).

Plan

1. Load and chunk the PDF document (Stock_Market_Performance_2024.pdf).
2. Embed chunks using OpenAIEmbeddings and store them in a local Chroma vector database.
3. Define a retriever tool function (retriever_tool) and bind it to the LLM.
4. Set up a StateGraph containing an LLM call node, a retriever action node, and conditional routing edges.

Code Implementation

1. Load and Chunk PDF

We split document text to feed to embeddings model:

import os
from langchain_community.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter

pdf_path = "Stock_Market_Performance_2024.pdf"
loader = PyPDFLoader(pdf_path)
pages = loader.load()

text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
pages_split = text_splitter.split_documents(pages)

2. Setup ChromaDB Vector Store

We embed the split documents and store them in a Chroma DB:

from langchain_openai import OpenAIEmbeddings
from langchain_chroma import Chroma

embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
vectorstore = Chroma.from_documents(
    documents=pages_split,
    embedding=embeddings,
    collection_name="stock_market"
)
retriever = vectorstore.as_retriever(search_type="similarity", search_kwargs={"k": 5})

3. Define Retriever Tool

We wrap retrieval search in a custom tool and bind it to the model:

from langchain_core.tools import tool
from langchain_openai import ChatOpenAI

@tool
def retriever_tool(query: str) -> str:
    """This tool searches and returns the information from the Stock Market Performance 2024 document."""
    docs = retriever.invoke(query)
    results = [f"Document {i+1}:\n{doc.page_content}" for i, doc in enumerate(docs)]
    return "

".join(results)

tools = [retriever_tool]
tools_dict = {t.name: t for t in tools}
llm = ChatOpenAI(model="gpt-4o", temperature=0).bind_tools(tools)

4. Define Nodes and Build Graph

We define the graph node calls and setup the graph structure:

from typing import Annotated, Sequence, TypedDict
from langchain_core.messages import BaseMessage, SystemMessage, ToolMessage
from langgraph.graph.message import add_messages
from langgraph.graph import StateGraph, END

class AgentState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], add_messages]

def call_llm(state: AgentState) -> dict:
    sys = SystemMessage(content="You are an AI assistant answering questions about Stock Market Performance.")
    message = llm.invoke([sys] + list(state["messages"]))
    return {"messages": [message]}

def take_action(state: AgentState) -> dict:
    tool_calls = state["messages"][-1].tool_calls
    results = []
    for t in tool_calls:
        res = tools_dict[t["name"]].invoke(t["args"].get("query", ""))
        results.append(ToolMessage(tool_call_id=t["id"], name=t["name"], content=str(res)))
    return {"messages": results}

def should_continue(state: AgentState):
    last_msg = state["messages"][-1]
    if hasattr(last_msg, "tool_calls") and last_msg.tool_calls:
        return "retriever_agent"
    return "end"

graph = StateGraph(AgentState)
graph.add_node("llm", call_llm)
graph.add_node("retriever_agent", take_action)

graph.set_entry_point("llm")
graph.add_conditional_edges(
    "llm",
    should_continue,
    {
        "retriever_agent": "retriever_agent",
        "end": END
    }
)
graph.add_edge("retriever_agent", "llm")
app = graph.compile()

5. Invoke the Agent

result = app.invoke({"messages": [HumanMessage(content="How did the tech sector perform in 2024?")]})
print(result["messages"][-1].content)

Exercise: Multi-Document RAG with Document Type Router 📂

Goal

Add a second Corporate Email records search tool to the RAG Agent, allowing the model to choose the correct resource (Stock Performance DB or Email database) dynamically based on user query.

Sample Input

{"messages": [("user", "Check the meeting email: why was the user unable to attend?")]}

Sample Output

Outputs citing the email document (e.g., "The user was unable to attend due to a conflicting doctor's appointment.").

Plan

1. Create a retriever tool @tool named email_retriever returning content from corporate email records.
2. Bind both [retriever_tool, email_retriever] to the LLM.
3. Configure take_action to map and invoke both tools correctly.
4. Compile the graph, invoke with the sample query, and print the results.

Solution

from typing import Annotated, Sequence, TypedDict
from langchain_core.messages import BaseMessage, SystemMessage, ToolMessage
from langgraph.graph.message import add_messages
from langchain_openai import ChatOpenAI
from langchain_core.tools import tool
from langgraph.graph import StateGraph, START, END

# 1. State
class AgentState(TypedDict):
    messages: Annotated[Sequence[BaseMessage], add_messages]

# 2. Tools
@tool
def retriever_tool(query: str) -> str:
    """This tool searches the Stock Market Performance 2024 database."""
    return "In 2024, the tech sector gained 25% driven by AI breakthroughs."

@tool
def email_retriever(query: str) -> str:
    """This tool searches corporate email records."""
    return "Subject: Unable to Attend Meeting\n\nHi, I cannot make it to the meeting today because I have a doctor's appointment scheduled at that time."

tools = [retriever_tool, email_retriever]
tools_dict = {t.name: t for t in tools}
llm = ChatOpenAI(model="gpt-4o", temperature=0).bind_tools(tools)

# 3. Nodes
def call_llm(state: AgentState) -> dict:
    sys = SystemMessage(content="Use retriever_tool for stock performance, and email_retriever for email records.")
    msg = llm.invoke([sys] + list(state["messages"]))
    return {"messages": [msg]}

def take_action(state: AgentState) -> dict:
    tool_calls = state["messages"][-1].tool_calls
    results = []
    for t in tool_calls:
        res_content = tools_dict[t["name"]].invoke(t["args"].get("query", ""))
        results.append(ToolMessage(tool_call_id=t["id"], name=t["name"], content=str(res_content)))
    return {"messages": results}

def should_continue(state: AgentState):
    last_msg = state["messages"][-1]
    return "retriever_agent" if (hasattr(last_msg, "tool_calls") and last_msg.tool_calls) else "end"

# 4. Build
graph = StateGraph(AgentState)
graph.add_node("llm", call_llm)
graph.add_node("retriever_agent", take_action)
graph.set_entry_point("llm")
graph.add_conditional_edges(
    "llm",
    should_continue,
    {
        "retriever_agent": "retriever_agent",
        "end": END
    }
)
graph.add_edge("retriever_agent", "llm")
app = graph.compile()

# 5. Test
res = app.invoke({"messages": [("user", "Check the meeting email: why was the user unable to attend?")]})
print(res["messages"][-1].content)

Practice & Exercises

To reinforce what you’ve learned in this section, practice with the interactive notebook:

Practice & Exercises

Practice loading source documents, chunking text, setting up a local vector database retriever, and routing questions dynamically.

💻 VS Code | 🚀 Colab | 📥 Download