3. Agent ReAct Docstore

Build a ReAct agent that queries a local document store via a RAG retrieval chain tool

In this section, you will learn how to integrate a complete RAG (Retrieval-Augmented Generation) pipeline as a custom query tool inside a ReAct agent.

Objectives

1. Load a pre-existing Chroma vector database using Chroma and OpenAIEmbeddings.
2. Configure history-aware retrievers using create_history_aware_retriever.
3. Wrap RAG retrieval chains in a Tool constructor and execute them within an agent loop.

Implementation Plan

Goal

Set up a RAG pipeline querying a local database, wrap it inside a custom tool, and bind it to a ReAct agent to answer domain-specific questions in an interactive chat session.

Sample Input

"You: What is Apple Intelligence?"

Sample Output

An AI assistant response generated using context retrieved from the database files.

Plan

1. Retrieve path configuration to load the Chroma vector database.
2. Initialize OpenAIEmbeddings and load the persistent vector database instance.
3. Configure search settings (e.g. similarity search with $k=3$) to build a retriever.
4. Establish a contextualization prompt to convert chat inputs into standalone queries.
5. Create a history_aware_retriever and link it to a document-stuffing QA chain using create_retrieval_chain to obtain rag_chain.
6. Define a custom Tool wrapper mapping tool invocation to the rag_chain.invoke endpoint.
7. Pull the ReAct prompt hwchase17/react from the hub, instantiate the ReAct agent, and build the interaction loops.

Step-by-Step Implementation

Step 1: Load Vector Store and Embeddings

We configure the directory variables and load our pre-populated vector database.

import os
from langchain_community.vectorstores import Chroma
from langchain_google_genai import GoogleGenAIEmbeddings

current_dir = os.path.dirname(os.path.abspath(__file__))
db_dir = os.path.join(current_dir, "..", "..", "4_rag", "db")
persistent_directory = os.path.join(db_dir, "chroma_db_with_metadata")

embeddings = GoogleGenAIEmbeddings(model="models/text-embedding-004")
db = Chroma(persist_directory=persistent_directory, embedding_function=embeddings)
retriever = db.as_retriever(search_type="similarity", search_kwargs={"k": 3})

Step 2: Set Up Retrieval Q&A Chain

We set up a history-aware retriever to formulate standalone questions and combine it with a documents chain.

from langchain.chat_models import init_chat_model
from langchain.chains import create_history_aware_retriever, create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder

llm = init_chat_model("llama-3.3-70b-versatile", model_provider="groq")

contextualize_q_prompt = ChatPromptTemplate.from_messages([
    ("system", "Formulate a standalone question based on history..."),
    MessagesPlaceholder("chat_history"),
    ("human", "{input}"),
])
history_aware_retriever = create_history_aware_retriever(llm, retriever, contextualize_q_prompt)

qa_prompt = ChatPromptTemplate.from_messages([
    ("system", "Answer questions based on context: {context}"),
    MessagesPlaceholder("chat_history"),
    ("human", "{input}"),
])
question_answer_chain = create_stuff_documents_chain(llm, qa_prompt)
rag_chain = create_retrieval_chain(history_aware_retriever, question_answer_chain)

Step 3: Wrap RAG Chain as a Tool

We define a custom tool named "Answer Question" that maps its execution function to run the RAG chain pipeline.

from langchain_core.tools import Tool

tools = [
    Tool(
        name="Answer Question",
        func=lambda input, **kwargs: rag_chain.invoke(
            {"input": input, "chat_history": kwargs.get("chat_history", [])}
        ),
        description="useful for when you need to answer questions about the context",
    )
]

Step 4: Create ReAct Agent and Execute Loop

We pull the prompt, construct the agent using create_react_agent, and run the interactive loop.

from langchain import hub
from langchain.agents import AgentExecutor, create_react_agent
from langchain_core.messages import HumanMessage, AIMessage

react_docstore_prompt = hub.pull("hwchase17/react")
agent = create_react_agent(llm=llm, tools=tools, prompt=react_docstore_prompt)
agent_executor = AgentExecutor.from_agent_and_tools(agent=agent, tools=tools, handle_parsing_errors=True, verbose=True)

# chat_history = []
# query = "What is Apple Intelligence?"
# response = agent_executor.invoke({"input": query, "chat_history": chat_history})
# chat_history.append(HumanMessage(content=query))
# chat_history.append(AIMessage(content=response['output']))

Complete Combined Code

Below is the complete, consolidated Python script uniting all of the steps above:

import os

from dotenv import load_dotenv
from langchain import hub
from langchain.agents import AgentExecutor, create_react_agent
from langchain.chains import create_history_aware_retriever, create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_community.vectorstores import Chroma
from langchain_core.messages import AIMessage, HumanMessage
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.tools import Tool
from langchain.chat_models import init_chat_model
from langchain_google_genai import GoogleGenAIEmbeddings

# Load environment variables from .env file
load_dotenv()

# Load the existing Chroma vector store
current_dir = os.path.dirname(os.path.abspath(__file__))
db_dir = os.path.join(current_dir, "..", "..", "4_rag", "db")
persistent_directory = os.path.join(db_dir, "chroma_db_with_metadata")

# Check if the Chroma vector store already exists
if os.path.exists(persistent_directory):
    print("Loading existing vector store...")
    db = Chroma(persist_directory=persistent_directory,
                embedding_function=None)
else:
    raise FileNotFoundError(
        f"The directory {persistent_directory} does not exist. Please check the path."
    )

# Define the embedding model
embeddings = GoogleGenAIEmbeddings(model="models/text-embedding-004")

# Load the existing vector store with the embedding function
db = Chroma(persist_directory=persistent_directory,
            embedding_function=embeddings)

# Create a retriever for querying the vector store
retriever = db.as_retriever(
    search_type="similarity",
    search_kwargs={"k": 3},
)

# Create a Groq model using core abstractions
llm = init_chat_model("llama-3.3-70b-versatile", model_provider="groq")

# Contextualize question prompt
contextualize_q_system_prompt = (
    "Given a chat history and the latest user question "
    "which might reference context in the chat history, "
    "formulate a standalone question which can be understood "
    "without the chat history. Do NOT answer the question, just "
    "reformulate it if needed and otherwise return it as is."
)

# Create a prompt template for contextualizing questions
contextualize_q_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", contextualize_q_system_prompt),
        MessagesPlaceholder("chat_history"),
        ("human", "{input}"),
    ]
)

# Create a history-aware retriever
history_aware_retriever = create_history_aware_retriever(
    llm, retriever, contextualize_q_prompt
)

# Answer question prompt
qa_system_prompt = (
    "You are an assistant for question-answering tasks. Use "
    "the following pieces of retrieved context to answer the "
    "question. If you don't know the answer, just say that you "
    "don't know. Use three sentences maximum and keep the answer "
    "concise."
    "\n\n"
    "{context}"
)

# Create a prompt template for answering questions
qa_prompt = ChatPromptTemplate.from_messages(
    [
        ("system", qa_system_prompt),
        MessagesPlaceholder("chat_history"),
        ("human", "{input}"),
    ]
)

# Create a chain to combine documents for question answering
question_answer_chain = create_stuff_documents_chain(llm, qa_prompt)

# Create a retrieval chain that combines the history-aware retriever and the question answering chain
rag_chain = create_retrieval_chain(
    history_aware_retriever, question_answer_chain)


# Set Up ReAct Agent with Document Store Retriever
react_docstore_prompt = hub.pull("hwchase17/react")

tools = [
    Tool(
        name="Answer Question",
        func=lambda input, **kwargs: rag_chain.invoke(
            {"input": input, "chat_history": kwargs.get("chat_history", [])}
        ),
        description="useful for when you need to answer questions about the context",
    )
]

# Create the ReAct Agent with document store retriever
agent = create_react_agent(
    llm=llm,
    tools=tools,
    prompt=react_docstore_prompt,
)

agent_executor = AgentExecutor.from_agent_and_tools(
    agent=agent, tools=tools, handle_parsing_errors=True, verbose=True,
)

chat_history = []
while True:
    query = input("You: ")
    if query.lower() == "exit":
        break
    response = agent_executor.invoke(
        {"input": query, "chat_history": chat_history})
    print(f"AI: {response['output']}")

    # Update history
    chat_history.append(HumanMessage(content=query))
    chat_history.append(AIMessage(content=response["output"]))

Practice & Exercises

To practice querying local document stores using agents, open the interactive notebook:

Practice & Exercises

Practice loading vector database retrievers, writing history-aware retrievers, and setting up retriever-agent wrappers.

💻 VS Code | 🚀 Colab | 📥 Download