← Back to DevBytes

Building a Document Q&A Agent with LangGraph: Step-by-Step Tutorial

Introduction to Document Q&A Agents and LangGraph

Building a question-answering system over your own documents is one of the most practical applications of large language models (LLMs). A basic RAG (Retrieval-Augmented Generation) pipeline can answer simple factual queries, but real-world questions often require multi-step reasoning, comparison across documents, or clarification. This is where an agentic approach shines. In this tutorial, you'll learn how to build a powerful, stateful Document Q&A agent using LangGraph, a framework designed for orchestrating LLM workflows as cyclical graphs.

What is LangGraph?

LangGraph is a library from the LangChain ecosystem that allows you to define computation as a directed, potentially cyclic graph. Unlike a simple DAG (directed acyclic graph) where data flows in one direction, LangGraph supports cycles and conditional branching. This makes it ideal for agents that need to iterate, reflect, and call tools multiple times before arriving at a final answer. At its core, LangGraph models an agent as a state machine, where each node updates a shared state, and edges determine which node executes next based on the current state.

Why LangGraph for Document Q&A?

A traditional RAG chain retrieves documents once and generates an answer. If the retrieved context is insufficient or irrelevant, the chain has no recourse. A LangGraph agent can:

By structuring the Q&A flow as a state graph, you gain fine-grained control and the ability to handle complex questions that stump simpler pipelines.

Setting Up Your Environment

Before we write any code, make sure you have Python 3.10+ and a working virtual environment. We'll need several packages: langgraph, langchain-core, langchain-openai (or your preferred LLM provider), chromadb for vector storage, and langchain-text-splitters. Run the following command:

pip install langgraph langchain-core langchain-openai chromadb langchain-text-splitters

You'll also need an OpenAI API key (or equivalent) set as an environment variable:

export OPENAI_API_KEY="your-key-here"

All code examples assume this setup. We'll build everything step by step in a single Python script or notebook.

Core Components of a Document Q&A Agent

Before diving into LangGraph specifics, let's outline the essential pieces our agent will need:

Document Ingestion and Indexing

We'll keep this part straightforward. Use a text loader to read documents, split them into chunks with overlap, compute embeddings using OpenAI, and store them in a Chroma vector store. This is a one-time setup that you can run separately.

from langchain_community.document_loaders import TextLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_openai import OpenAIEmbeddings
from langchain_community.vectorstores import Chroma

# Load a sample document (replace with your own)
loader = TextLoader("knowledge_base.txt")
documents = loader.load()

# Split into chunks
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=1000,
    chunk_overlap=200,
    separators=["\n\n", "\n", " ", ""]
)
chunks = text_splitter.split_documents(documents)

# Create vector store
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")
vector_store = Chroma.from_documents(
    documents=chunks,
    embedding=embeddings,
    collection_name="my_knowledge_base"
)

Now vector_store can be used as a retriever. In production, you might persist the Chroma database to disk, but for this tutorial we'll keep it in memory.

Retrieval and RAG Workflow

A simple RAG chain retrieves documents based on user input and passes them to the LLM for answer generation. Our agent will wrap the retrieval in a tool call, but the core logic remains the same. We'll use the vector store's as_retriever() method with similarity search.

Agentic Decision-Making with LangGraph

The agent's "brain" is a graph that can decide to stop or continue based on the quality of retrieved context. We'll implement a cycle: retrieve β†’ check relevance β†’ if relevant, generate answer; else, rewrite the question β†’ retrieve again. This loop continues until the answer is found or a maximum number of iterations is reached.

Building the Agent: Step-by-Step

Now we'll construct the LangGraph agent. The code will be presented in logical blocks, each building on the previous. By the end, you'll have a fully functional script.

Step 1: Imports and LLM Initialization

Start by importing the required modules and initializing your LLM and embeddings. We'll use LangChain's ChatOpenAI for both generation and query rewriting.

import os
from typing import TypedDict, List, Annotated
from langgraph.graph import StateGraph, END
from langgraph.graph.message import add_messages
from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_core.tools import tool
from langchain_community.vectorstores import Chroma
from langgraph.prebuilt import ToolNode
import operator

# Initialize the LLM
llm = ChatOpenAI(model="gpt-4o", temperature=0)
# Embeddings are already defined, but we'll reuse the same model
embeddings = OpenAIEmbeddings(model="text-embedding-3-small")

Step 2: Rebuild or Load the Vector Store

If you ran the ingestion code earlier, you can load an existing Chroma database. For simplicity, we'll assume the vector store is already populated and accessible as a variable. In practice, you'd persist it:

# If you already built it, reuse it. Otherwise build as shown earlier.
# Here we just reference the pre-built store.
# vector_store = ... (already exists from previous step)

Step 3: Define the Retrieval Tool

LangGraph agents often interact with external resources through tools. We'll define a simple retrieval tool using the @tool decorator. This tool will take a query string and return a list of document excerpts. Later we'll bind it to the graph.

@tool
def retrieve_documents(query: str) -> str:
    """Search the knowledge base and return relevant document excerpts."""
    retriever = vector_store.as_retriever(
        search_type="similarity",
        search_kwargs={"k": 3}
    )
    docs = retriever.invoke(query)
    # Format results as a single string
    formatted = "\n\n".join(
        [f"Source {i+1}:\n{doc.page_content}" for i, doc in enumerate(docs)]
    )
    return formatted

tools = [retrieve_documents]
tool_node = ToolNode(tools)  # LangGraph prebuilt node to execute tools

ToolNode is a special node that takes the last message (which contains tool calls) and returns the tool results. We'll integrate it shortly.

Step 4: Define the Agent State

The state is the central data structure that flows through the graph. We'll use a TypedDict to keep it well-typed. It will include the user's question, the chat history (messages), retrieved context, a final answer, and an iteration counter to prevent infinite loops.

class AgentState(TypedDict):
    question: str                          # original user question
    messages: Annotated[list, add_messages] # conversation history (including tool calls)
    retrieved_context: str                 # latest retrieved documents
    final_answer: str                     # the answer to return
    iteration_count: int                   # how many retrieval attempts so far

Annotated[list, add_messages] is a special annotation that tells LangGraph how to merge lists across node updatesβ€”here it appends new messages to the existing list.

Step 5: Implement the Node Functions

Each node is a Python function that takes the current state and returns an update to the state. We need at least four nodes:

We'll implement each one carefully. Note that retrieve_node will be handled by ToolNode in our final graph design, but for clarity we can also define a custom node.

def retrieve_node(state: AgentState) -> AgentState:
    """
    Execute the retrieval tool and update context.
    We'll use the tool directly.
    """
    question = state["question"]
    # Call the tool (we can invoke it like a function)
    context = retrieve_documents.invoke(question)
    return {
        "retrieved_context": context,
        "iteration_count": state.get("iteration_count", 0) + 1
    }

def rewrite_question_node(state: AgentState) -> AgentState:
    """
    Use the LLM to rewrite the question to be more search-friendly.
    """
    question = state["question"]
    context_summary = state.get("retrieved_context", "")[:500]  # brief summary
    prompt = f"""You are a helpful assistant. The user asked:
"{question}"

After searching, the retrieved information was:
{context_summary}

The context may not fully answer the question. Please rephrase the question to better target the knowledge base. 
Output only the rewritten question."""
    response = llm.invoke(prompt)
    new_question = response.content.strip()
    return {"question": new_question}

def relevance_check_node(state: AgentState) -> AgentState:
    """
    LLM decides if the retrieved context is sufficient to answer.
    Returns an update with a flag we'll use in routing.
    """
    question = state["question"]
    context = state.get("retrieved_context", "")
    prompt = f"""Given the question and the following retrieved document excerpts, decide whether the provided context contains enough information to answer the question fully.

Question: {question}

Context:
{context}

If the context is sufficient, respond with exactly "SUFFICIENT". If it is missing key details or completely irrelevant, respond with "INSUFFICIENT". Do not add any other text."""
    response = llm.invoke(prompt)
    decision = response.content.strip().upper()
    # We'll store the decision as a special key in the state.
    # Alternatively, we can just return a new key like 'next_action'.
    return {"messages": [("assistant", f"Relevance decision: {decision}")]}

def generate_answer_node(state: AgentState) -> AgentState:
    """
    Generate the final answer using the retrieved context.
    """
    question = state["question"]
    context = state.get("retrieved_context", "")
    prompt = f"""Answer the question based on the provided context. If the context does not contain the answer, say "I couldn't find the answer in the documents."

Question: {question}

Context:
{context}

Answer:"""
    response = llm.invoke(prompt)
    return {"final_answer": response.content.strip()}

We now have all the core node functions. Notice that relevance_check_node only adds a message; we'll extract the decision in the conditional edge function.

Step 6: Define Routing Logic and Build the Graph

LangGraph graphs are built using StateGraph. We add nodes and edges, specifying how the state flows. Conditional edges allow branching based on state. We'll design the following flow:

  1. Start at retrieve_node β†’ go to relevance_check_node.
  2. After relevance check: if "SUFFICIENT" β†’ go to generate_answer_node β†’ END.
  3. If "INSUFFICIENT" and iteration_count < max_iterations β†’ go to rewrite_question_node β†’ back to retrieve_node.
  4. If max iterations reached β†’ go to generate_answer_node anyway (force answer).

Let's define the conditional routing function:

def route_after_check(state: AgentState) -> str:
    """
    Examine the last assistant message to get the relevance decision,
    and decide next node.
    """
    last_message = state["messages"][-1]
    # The content is like "Relevance decision: SUFFICIENT"
    if "SUFFICIENT" in last_message.content:
        return "generate_answer"
    else:
        iteration = state.get("iteration_count", 0)
        max_iterations = 3
        if iteration >= max_iterations:
            print("Maximum iterations reached, forcing answer generation.")
            return "generate_answer"
        return "rewrite_question"

Now assemble the graph:

# Create the graph
workflow = StateGraph(AgentState)

# Add nodes
workflow.add_node("retrieve", retrieve_node)
workflow.add_node("relevance_check", relevance_check_node)
workflow.add_node("rewrite_question", rewrite_question_node)
workflow.add_node("generate_answer", generate_answer_node)

# Add edges
workflow.add_edge("retrieve", "relevance_check")
workflow.add_conditional_edges(
    "relevance_check",
    route_after_check,
    {
        "generate_answer": "generate_answer",
        "rewrite_question": "rewrite_question"
    }
)
workflow.add_edge("rewrite_question", "retrieve")
workflow.add_edge("generate_answer", END)

# Set the entry point
workflow.set_entry_point("retrieve")

# Compile the graph
agent = workflow.compile()

The compiled agent object is a runnable that you can invoke with an initial state.

Step 7: Run the Agent and Ask Questions

To test the agent, we create an initial state with a question and empty fields, then invoke it. The graph will execute nodes in the defined order until it reaches END. Let's run a query:

# Prepare initial state
initial_state = {
    "question": "What is the return policy for international orders?",
    "messages": [],
    "retrieved_context": "",
    "final_answer": "",
    "iteration_count": 0
}

# Execute the agent
final_state = agent.invoke(initial_state)

# Print the answer
print("Final Answer:")
print(final_state["final_answer"])
print("\n--- Iterations:", final_state["iteration_count"])

The output will show the final answer. If your documents contain the return policy, the agent should retrieve it and answer directly. If not, it may rewrite the question a couple of times before giving up and returning a fallback message.

Step 8: Adding Tool-Based Retrieval (Alternative Approach)

The above uses custom nodes. You can also leverage LangGraph's built-in ToolNode and let an LLM decide when to call tools. This is more agentic but requires careful prompt engineering. For completeness, here's a sketch using a "chat agent" pattern with tools:

from langgraph.prebuilt import create_react_agent

# Define tools list
tools = [retrieve_documents]

# Create a ReAct-style agent graph
agent = create_react_agent(llm, tools)

# Run it
result = agent.invoke({"messages": [("user", "What is the return policy?")]})
print(result["messages"][-1].content)

This approach works well but gives less control over the iteration logic. Our custom graph is more explicit and easier to debug for document Q&A.

Best Practices for Production-Ready Agents

Conclusion

You've now built a complete Document Q&A agent using LangGraph. Starting from document ingestion, through embedding and vector storage, to a stateful graph that retrieves, evaluates, rewrites, and generates answers, you've seen how to go beyond a static RAG pipeline. The agent can iterate and adapt, making it far more robust for ambiguous or complex queries. LangGraph's graph-based architecture gives you the flexibility to add more nodes, integrate additional tools, or even connect multiple agents. Experiment with the code, adapt it to your own documents, and explore the LangGraph documentation for advanced patterns like persistence, streaming, and multi-agent collaboration. With these fundamentals, you're ready to build intelligent, reliable document assistants that truly understand your knowledge base.

β€” Ad β€”

Google AdSense will appear here after approval

← Back to all articles