Build a RAG Agent with LangGraph

Build a document Q&A agent with LangGraph, FAISS, and MLflow tracing.

Introduction

30 Second Summary

When you ask a colleague a question, you expect them to know when to look something up versus answer off the top of their head. AI agents need that same decision-making ability.

In this project, you will build a document Q&A agent that intelligently decides whether to retrieve information from a knowledge base or answer directly. You will also set up a tracing dashboard to monitor every decision the agent makes.

What You'll Build

A LangGraph state graph agent with conditional routing, FAISS vector search for document retrieval, Google Gemini as the LLM, and MLflow tracing for production observability.

By the end of this project, you'll have:

  • A LangGraph agent with typed state and conditional routing that decides when to retrieve documents vs. answer directly.
  • A FAISS vector store loaded with sample documents for semantic retrieval using the RAG pattern.
  • An MLflow tracing dashboard showing agent execution traces, latency, and token usage for LLMOps observability.
  • Secret Mission: Add a human-in-the-loop approval step using LangGraph interrupts before the agent answers sensitive questions.

Are there any prerequisites?

No prior experience with LangGraph, FAISS, or MLflow is needed. This guide includes setup instructions for all required tools, including Python and your free Gemini API key.

Before We Start

Before We Start...

Set Up Your Environment

Before you can build an AI agent, you need the right tools in place. This project uses Python libraries for agentic AI, vector search, and observability.

In this step, you will install Python, create an isolated project environment, install all the required packages, and connect to Google's Gemini LLM.

In this step, get ready to:
  • Install Python and create a virtual environment.
  • Install all project dependencies.
  • Set up your Gemini API key and verify the connection.
Install Python and create your project environment

This project requires Python 3.10 or higher. LangGraph and all the other packages in this project need at least this version to work.

  • Press Cmd+Space (macOS) or the Windows key (Windows) to open your search bar.
  • Type Terminal and press Enter to open it.
  • Check your Python version by running this command:
python3 --version

✔️ I see version 3.10 or higher

Your Python version is compatible. Move on to creating your project folder.

ⓧ I see an older version

Your system has an older Python that won't work with this project's dependencies. You need to install Python 3.10 or higher.

  • Download Python 3.10+ from the official Python downloads page.
  • Run the installer and follow the prompts.
  • On macOS, the system python3 may still point to the old version after installing. Use python3.10 (or python3.11, python3.12) directly, or use a version manager like pyenv to set the default.
  • After installing, run python3 --version again to confirm you see 3.10 or higher.

ⓧ Command not found

Python is not installed on your machine yet.

  • Download Python 3.10+ from the official Python downloads page.
  • Run the installer and follow the prompts.
  • On Windows, check the box that says Add Python to PATH during installation.
  • After installing, close and reopen your terminal, then run python3 --version to confirm it prints 3.10 or higher.

Now create your project folder on your Desktop and set up a virtual environment.

  • Move to your Desktop by running this command:
cd ~/Desktop
  • Create your project folder and move into it by running these commands:
mkdir rag-agent
cd rag-agent
  • Create a virtual environment by running this command:
python3 -m venv venv

What is a virtual environment?

A virtual environment is an isolated space for your Python project's packages. It keeps the libraries you install here separate from your system Python, so different projects don't conflict with each other.

  • Activate the virtual environment by running this command:

macOS

source venv/bin/activate

Windows

venv\Scripts\activate
  • Confirm the files exist by running:
ls

You should see a venv folder listed in the output. Your terminal prompt should also show (venv) at the beginning, confirming the virtual environment is active.

Don't see (venv) in your prompt?

Make sure you ran the activation command for your operating system. On macOS, use source venv/bin/activate. On Windows, use venv\Scripts\activate.

If you see a permissions error on Windows, try running your terminal as Administrator.

Help me activate my Python virtual environment.

Install all required packages

With your virtual environment active, you can install the Python packages that power your agent. These include the agent framework, vector search, LLM integration, and observability tools.

  • Install the project dependencies by running this command:
pip install langgraph langchain-google-genai langchain-community faiss-cpu mlflow langchain-text-splitters

What are these packages?

  • langgraph is the agent framework that lets you build stateful AI workflows as graphs.
  • langchain-google-genai connects your agent to Google's Gemini LLM.
  • langchain-community provides integrations for vector stores and document loaders.
  • faiss-cpu is Meta's library for fast vector similarity search.
  • mlflow provides tracing and observability for your LLM calls.
  • langchain-text-splitters provides utilities for splitting documents into chunks for indexing.

The installation may take a minute or two as pip downloads all the packages and their dependencies.

You should see a success message at the end listing all installed packages.

Seeing installation errors?

Make sure your virtual environment is active. You should see (venv) at the start of your terminal prompt.

If you see permission errors, do not use sudo. Instead, deactivate and recreate your virtual environment.

Help me fix my pip install errors.

Get your Gemini API key and verify the connection

Your agent needs access to Google's Gemini LLM to generate responses. The free tier is more than enough for this project.

  • Open Google AI Studio in your browser.
  • Sign in with your Google account.
  • Click Get API Key in the left sidebar.
  • Click Create API Key.
  • Copy the generated key and store it somewhere safe. You will not be able to see it again.

Why is this free?

Google offers a free tier for Gemini API access that includes enough requests for learning projects like this one. No credit card is required to create a key.

Now set the API key as an environment variable so your code can access it.

  • Set the environment variable by running this command (replace your-api-key-here with the key you copied):

macOS

export GEMINI_API_KEY=your-api-key-here

Windows

set GEMINI_API_KEY=your-api-key-here

Why an environment variable?

Storing your API key in an environment variable keeps it out of your source code. This prevents accidentally sharing your key if you push code to a repository.

  • Verify the connection works by running this command:
python3 -c "from langchain_google_genai import ChatGoogleGenerativeAI; print('OK')"

✔️ I see OK printed

Your environment is fully configured. Python can import the Gemini integration and your API key is ready to use.

ⓧ I see an error

If you see ModuleNotFoundError, the package did not install correctly. Make sure your virtual environment is active (you should see (venv) in your prompt) and re-run the pip install command from the previous substep.

If you see an authentication or API error, double-check that you set the GEMINI_API_KEY environment variable correctly. Run echo $GEMINI_API_KEY (macOS) or echo %GEMINI_API_KEY% (Windows) to check its value.

Your environment is ready to go. Next up, you will build a FAISS vector store that your agent can search through to find answers.

Build the FAISS Vector Store

Your environment is ready and your Gemini connection is verified. Now it's time to build the knowledge base that your agent will search through when answering questions.

A RAG agent needs a way to find relevant information fast. Instead of scanning every document word-by-word, you'll use FAISS (Facebook AI Similarity Search) to convert documents into numerical vectors and search them by meaning.

In this step, get ready to:
  • Create a DocumentStore class that encapsulates document loading, splitting, and FAISS indexing.
  • Load sample company policy documents into the vector store and persist the index to disk.
  • Test retrieval by querying the vector store and verifying relevant chunks are returned.
Create the DocumentStore class

Your vector store logic will live in a dedicated class called DocumentStore. This follows the Object-Oriented design pattern of encapsulation. All the complexity of loading documents, splitting them into chunks, and managing the FAISS index stays inside one class.

  • In VS Code, create a new file called knowledge_base.py in your project directory.
  • Add the imports and class initialization by pasting this code:
import os
from langchain_google_genai import GoogleGenerativeAIEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_text_splitters import CharacterTextSplitter
from langchain_community.document_loaders import TextLoader, DirectoryLoader


class DocumentStore:
    """Encapsulates document loading, embedding, and FAISS vector store operations."""

    def __init__(self, persist_directory: str = "./faiss_index"):
        self.persist_directory = persist_directory
        self.embeddings = GoogleGenerativeAIEmbeddings(
            model="models/embedding-001",
            google_api_key=os.getenv("GEMINI_API_KEY"),
        )
        self.vector_store = None

What does this code do?

The imports bring in the tools you need: GoogleGenerativeAIEmbeddings converts text into vectors using Gemini's embedding model, FAISS handles the vector index, CharacterTextSplitter breaks documents into smaller chunks, and DirectoryLoader reads all text files from a folder.

The __init__ method sets up the embedding model (models/embedding-001) and a directory path where the FAISS index will be saved. The vector_store starts as None until you load or create an index.

  • Below the __init__ method, add the load_and_index method that handles document loading and FAISS index creation:
    def load_and_index(self, documents_path: str) -> None:
        """Load documents from a directory, split into chunks, and create FAISS index."""
        loader = DirectoryLoader(documents_path, glob="**/*.txt", loader_cls=TextLoader)
        documents = loader.load()

        text_splitter = CharacterTextSplitter(
            chunk_size=500,
            chunk_overlap=50,
            separator="\n",
        )
        chunks = text_splitter.split_documents(documents)

        self.vector_store = FAISS.from_documents(chunks, self.embeddings)
        self.vector_store.save_local(self.persist_directory)
        print(f"Indexed {len(chunks)} chunks from {len(documents)} documents.")

What does this method do?

This is the core indexing pipeline. DirectoryLoader grabs every .txt file from the specified folder. CharacterTextSplitter then breaks each document into chunks of 500 characters with a 50-character overlap so context isn't lost at chunk boundaries.

FAISS.from_documents converts each chunk into a vector embedding and stores them in a searchable index. Finally, save_local persists the index to disk so you don't have to rebuild it every time.

  • Below load_and_index, add the final two methods that handle loading a saved index and performing searches:
    def load_existing(self) -> None:
        """Load a previously persisted FAISS index from disk."""
        self.vector_store = FAISS.load_local(
            self.persist_directory,
            self.embeddings,
            allow_dangerous_deserialization=True,
        )

    def search(self, query: str, k: int = 3) -> list:
        """Perform similarity search and return top-k relevant chunks."""
        if self.vector_store is None:
            raise ValueError("Vector store not initialized. Call load_and_index() or load_existing() first.")
        return self.vector_store.similarity_search(query, k=k)

What do these methods do?

load_existing reloads a previously saved FAISS index from disk. The allow_dangerous_deserialization=True flag is required because FAISS uses pickle under the hood. This is safe for indexes you created yourself.

search takes a query string and returns the top k most semantically similar chunks from the index. It raises an error if you haven't loaded or created an index yet.

  • Save knowledge_base.py.

✔️ Awesome, I've got everything!

Great. Double check you've saved knowledge_base.py before moving on.

ⓧ I'd like to double check the full code

import os
from langchain_google_genai import GoogleGenerativeAIEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_text_splitters import CharacterTextSplitter
from langchain_community.document_loaders import TextLoader, DirectoryLoader


class DocumentStore:
    """Encapsulates document loading, embedding, and FAISS vector store operations."""

    def __init__(self, persist_directory: str = "./faiss_index"):
        self.persist_directory = persist_directory
        self.embeddings = GoogleGenerativeAIEmbeddings(
            model="models/embedding-001",
            google_api_key=os.getenv("GEMINI_API_KEY"),
        )
        self.vector_store = None

    def load_and_index(self, documents_path: str) -> None:
        """Load documents from a directory, split into chunks, and create FAISS index."""
        loader = DirectoryLoader(documents_path, glob="**/*.txt", loader_cls=TextLoader)
        documents = loader.load()

        text_splitter = CharacterTextSplitter(
            chunk_size=500,
            chunk_overlap=50,
            separator="\n",
        )
        chunks = text_splitter.split_documents(documents)

        self.vector_store = FAISS.from_documents(chunks, self.embeddings)
        self.vector_store.save_local(self.persist_directory)
        print(f"Indexed {len(chunks)} chunks from {len(documents)} documents.")

    def load_existing(self) -> None:
        """Load a previously persisted FAISS index from disk."""
        self.vector_store = FAISS.load_local(
            self.persist_directory,
            self.embeddings,
            allow_dangerous_deserialization=True,
        )

    def search(self, query: str, k: int = 3) -> list:
        """Perform similarity search and return top-k relevant chunks."""
        if self.vector_store is None:
            raise ValueError("Vector store not initialized. Call load_and_index() or load_existing() first.")
        return self.vector_store.similarity_search(query, k=k)
Create sample documents and build the index

Your vector store needs documents to index. You'll create a set of fictional company policy documents that your agent will search through when answering questions.

  • Create a documents folder inside your project directory by running this command:
mkdir documents
  • Confirm the folder exists by running:
ls

You should see documents and knowledge_base.py listed in the output.

  • In VS Code, create a new file at documents/vacation_policy.txt and paste in the following content:
Company Vacation Policy

All full-time employees are entitled to 20 days of paid vacation per year.
Vacation days accrue at a rate of 1.67 days per month.
Unused vacation days can be carried over to the next year, up to a maximum of 5 days.
Employees must submit vacation requests at least 2 weeks in advance through the HR portal.
Manager approval is required for all vacation requests exceeding 5 consecutive days.
  • Create a second file at documents/remote_work_policy.txt and paste in this content:
Remote Work Policy

Employees may work remotely up to 3 days per week with manager approval.
A dedicated workspace with reliable internet is required for remote work days.
Core hours of 10am-3pm must be observed regardless of work location.
All remote work arrangements must be documented in the HR system.
Equipment provided by the company must be returned if remote work privileges are revoked.
  • Create a third file at documents/expense_policy.txt and paste in this content:
Expense Reimbursement Policy

Business expenses must be submitted within 30 days of the transaction.
Receipts are required for all expenses over $25.
Travel expenses including flights, hotels, and meals are reimbursable with prior approval.
The maximum daily meal allowance is $75 for domestic travel and $100 for international travel.
All expense reports require manager approval before processing by the finance team.
  • Save all three files.

Why three separate documents?

Each document covers a different company policy. When your agent receives a question like "What is the vacation policy?", the vector store needs to find the right chunks from the right document. Multiple documents test that your retrieval is actually semantic, not just returning everything.

  • Build the FAISS index from your documents by running this command in your terminal:
python3 -c "from knowledge_base import DocumentStore; ds = DocumentStore(); ds.load_and_index('./documents')"
  • You should see output like Indexed 3 chunks from 3 documents. confirming the index was built.
  • Verify the index was persisted by running:
ls faiss_index

You should see index.faiss and index.pkl listed. These are your saved vector index files.

Seeing an API error?

If you see an authentication or API key error, make sure your GEMINI_API_KEY environment variable is still set in this terminal session. Run echo $GEMINI_API_KEY to check. If it's empty, re-export it.

If you see a module not found error, confirm your virtual environment is still active (you should see (venv) in your terminal prompt).

Help me debug my FAISS indexing error

Test retrieval with a similarity search

The index is built. Now you need to verify that querying it returns the correct document chunks based on semantic similarity.

  • Run a test query against your vector store by running this command:
python3 -c "
from knowledge_base import DocumentStore
ds = DocumentStore()
ds.load_existing()
results = ds.search('What is the vacation policy?')
for r in results:
    print(r.page_content[:80])
    print('---')
"

You should see text from vacation_policy.txt appearing in the results. The vector store found the most relevant chunks based on the meaning of your query.

What is similarity search doing?

When you call search(), your query text gets converted into a vector using the same Gemini embedding model. FAISS then compares that query vector against all stored document vectors and returns the closest matches. This is how RAG agents find relevant context without keyword matching.

Results don't look relevant?

If you're seeing chunks from unrelated documents, try re-running the index build command. A corrupted index can return unexpected results.

If you see a "Vector store not initialized" error, make sure the faiss_index directory exists and contains index.faiss and index.pkl.

Help me fix my FAISS similarity search results

Your vector store is loaded, indexed, and returning relevant results. Next up, you'll build the LangGraph agent that uses this knowledge base to intelligently route and answer questions.

Build the LangGraph Agent

Your FAISS vector store is loaded and ready to answer semantic queries. The missing piece is an intelligent agent that decides when to use the knowledge base versus answering directly from general knowledge.

In this step, you will build a LangGraph state graph that routes questions through conditional edges. The agent uses a typed state object to track messages, retrieval context, and routing decisions as data flows between nodes.

In this step, get ready to:
  • Define a typed AgentState with LangGraph's state management.
  • Build a RAGAgent class with conditional routing nodes.
  • Test that the agent correctly routes questions to retrieval or direct answers.
Define the agent state and imports

A LangGraph agent needs a typed state object that defines what data flows between nodes. Each node reads from and writes to this shared state.

  • Create a new file called agent.py in your project folder by running:
touch agent.py
  • Open agent.py in VS Code and add the following imports and state definition:
import os
from typing import TypedDict, Annotated
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
from langgraph.graph import StateGraph, END
from langgraph.graph.message import add_messages
from knowledge_base import DocumentStore


class AgentState(TypedDict):
    """Typed state for the RAG agent graph."""
    messages: Annotated[list, add_messages]
    retrieved_context: str
    route_decision: str

What does this code do?

The imports bring in LangGraph's graph-building tools alongside LangChain's message types and your DocumentStore class from the previous step.

AgentState is a TypedDict that defines the shape of data flowing through the graph. The add_messages annotation on the messages field tells LangGraph to append new messages rather than overwrite them.

  • Save the file. You won't see output yet. The next chunk builds the class that uses this state.
Build the RAGAgent class with graph construction

The agent class encapsulates the LLM connection, the document store, and the graph topology. The _build_graph method wires nodes and conditional edges together.

  • In agent.py, add the following class below the AgentState definition:
class RAGAgent:
    """Object-oriented RAG agent built with LangGraph."""

    def __init__(self, documents_path: str = "./documents"):
        self.llm = ChatGoogleGenerativeAI(
            model="gemini-2.5-flash",
            temperature=0,
            google_api_key=os.getenv("GEMINI_API_KEY"),
        )
        self.doc_store = DocumentStore()
        try:
            self.doc_store.load_existing()
        except Exception:
            self.doc_store.load_and_index(documents_path)

        self.graph = self._build_graph()

    def _build_graph(self):
        """Construct the LangGraph state graph with conditional routing."""
        workflow = StateGraph(AgentState)

        workflow.add_node("router", self._router_node)
        workflow.add_node("retrieve", self._retrieve_node)
        workflow.add_node("generate", self._generate_node)

        workflow.set_entry_point("router")

        workflow.add_conditional_edges(
            "router",
            self._route_decision,
            {
                "retrieve": "retrieve",
                "direct": "generate",
            },
        )
        workflow.add_edge("retrieve", "generate")
        workflow.add_edge("generate", END)

        return workflow.compile()

What does this code do?

The __init__ method connects to Gemini, loads the FAISS index you built earlier, and compiles the graph.

_build_graph defines the agent's decision flow. It registers three nodes (router, retrieve, generate) and uses add_conditional_edges to branch after the router. If the router decides "retrieve", the graph sends the state to the retrieve node. If it decides "direct", the graph skips retrieval and goes straight to generate.

  • Save the file. The class references node methods that do not exist yet. You will add them next.
Implement the routing, retrieval, and generation nodes

Each node is a method that reads from the agent state, performs one job, and returns an update dictionary. The router asks the LLM to classify the question. The retrieve node queries FAISS. The generate node produces the final answer.

  • In agent.py, add the router node method inside the RAGAgent class, below _build_graph:
    def _router_node(self, state: AgentState) -> dict:
        """Decide whether to retrieve context or answer directly."""
        last_message = state["messages"][-1].content

        router_prompt = SystemMessage(
            content=(
                "You are a routing assistant. Determine if the user's question "
                "requires looking up company documents or can be answered with general knowledge. "
                "Respond with exactly one word: 'retrieve' or 'direct'."
            )
        )
        response = self.llm.invoke([router_prompt, HumanMessage(content=last_message)])
        decision = response.content.strip().lower()

        if "retrieve" in decision:
            return {"route_decision": "retrieve"}
        return {"route_decision": "direct"}

What does this code do?

The router node extracts the latest user message, then asks Gemini to classify the question as "retrieve" or "direct". The result is stored in route_decision in the agent state. The conditional edge you defined in _build_graph reads this field to decide the next node.

  • Below _router_node, add the retrieval node:
    def _retrieve_node(self, state: AgentState) -> dict:
        """Query the FAISS vector store for relevant document chunks."""
        last_message = state["messages"][-1].content
        results = self.doc_store.search(last_message, k=3)
        context = "\n\n".join([doc.page_content for doc in results])
        return {"retrieved_context": context}

What does this code do?

This node calls the search() method on your DocumentStore to find the top 3 most relevant chunks. It joins them into a single string and stores it in retrieved_context for the generate node to use.

  • Below _retrieve_node, add the generate node:
    def _generate_node(self, state: AgentState) -> dict:
        """Generate a response using the LLM with optional retrieved context."""
        last_message = state["messages"][-1].content
        context = state.get("retrieved_context", "")

        if context:
            system_prompt = SystemMessage(
                content=(
                    "You are a helpful assistant. Answer the user's question based on "
                    "the following context from company documents. If the context doesn't "
                    "contain the answer, say so.\n\n"
                    f"Context:\n{context}"
                )
            )
        else:
            system_prompt = SystemMessage(
                content="You are a helpful assistant. Answer the user's question directly."
            )

        response = self.llm.invoke([system_prompt, HumanMessage(content=last_message)])
        return {"messages": [AIMessage(content=response.content)]}

What does this code do?

The generate node checks whether retrieved_context exists in the state. If it does, the node includes that context in the system prompt so Gemini answers using your documents. If not, it answers from general knowledge.

The response is wrapped in an AIMessage and added to the messages list. Thanks to the add_messages annotation, LangGraph appends this to the conversation history rather than replacing it.

  • Finally, add the routing helper and public ask method below _generate_node:
    def _route_decision(self, state: AgentState) -> str:
        """Return the routing decision from state."""
        return state["route_decision"]

    def ask(self, question: str) -> str:
        """Run a question through the agent and return the response."""
        result = self.graph.invoke({"messages": [HumanMessage(content=question)]})
        return result["messages"][-1].content

What does this code do?

_route_decision is the condition function that add_conditional_edges calls after the router node. It reads the route_decision field from state and returns it so LangGraph knows which edge to follow.

The ask method is the public interface. It wraps the user's question in a HumanMessage, invokes the compiled graph, and returns the final AI response.

  • Save agent.py.

Seeing an IndentationError?

Every method from _router_node through ask must be indented one level inside the RAGAgent class. Make sure each def line starts with exactly 4 spaces.

Help me fix my agent.py indentation

✔️ Awesome, I've got everything!

Great. Double check you have saved agent.py before moving on.

ⓧ I'd like to double check the full code

import os
from typing import TypedDict, Annotated
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
from langgraph.graph import StateGraph, END
from langgraph.graph.message import add_messages
from knowledge_base import DocumentStore


class AgentState(TypedDict):
    """Typed state for the RAG agent graph."""
    messages: Annotated[list, add_messages]
    retrieved_context: str
    route_decision: str


class RAGAgent:
    """Object-oriented RAG agent built with LangGraph."""

    def __init__(self, documents_path: str = "./documents"):
        self.llm = ChatGoogleGenerativeAI(
            model="gemini-2.5-flash",
            temperature=0,
            google_api_key=os.getenv("GEMINI_API_KEY"),
        )
        self.doc_store = DocumentStore()
        try:
            self.doc_store.load_existing()
        except Exception:
            self.doc_store.load_and_index(documents_path)

        self.graph = self._build_graph()

    def _build_graph(self):
        """Construct the LangGraph state graph with conditional routing."""
        workflow = StateGraph(AgentState)

        workflow.add_node("router", self._router_node)
        workflow.add_node("retrieve", self._retrieve_node)
        workflow.add_node("generate", self._generate_node)

        workflow.set_entry_point("router")

        workflow.add_conditional_edges(
            "router",
            self._route_decision,
            {
                "retrieve": "retrieve",
                "direct": "generate",
            },
        )
        workflow.add_edge("retrieve", "generate")
        workflow.add_edge("generate", END)

        return workflow.compile()

    def _router_node(self, state: AgentState) -> dict:
        """Decide whether to retrieve context or answer directly."""
        last_message = state["messages"][-1].content

        router_prompt = SystemMessage(
            content=(
                "You are a routing assistant. Determine if the user's question "
                "requires looking up company documents or can be answered with general knowledge. "
                "Respond with exactly one word: 'retrieve' or 'direct'."
            )
        )
        response = self.llm.invoke([router_prompt, HumanMessage(content=last_message)])
        decision = response.content.strip().lower()

        if "retrieve" in decision:
            return {"route_decision": "retrieve"}
        return {"route_decision": "direct"}

    def _retrieve_node(self, state: AgentState) -> dict:
        """Query the FAISS vector store for relevant document chunks."""
        last_message = state["messages"][-1].content
        results = self.doc_store.search(last_message, k=3)
        context = "\n\n".join([doc.page_content for doc in results])
        return {"retrieved_context": context}

    def _generate_node(self, state: AgentState) -> dict:
        """Generate a response using the LLM with optional retrieved context."""
        last_message = state["messages"][-1].content
        context = state.get("retrieved_context", "")

        if context:
            system_prompt = SystemMessage(
                content=(
                    "You are a helpful assistant. Answer the user's question based on "
                    "the following context from company documents. If the context doesn't "
                    "contain the answer, say so.\n\n"
                    f"Context:\n{context}"
                )
            )
        else:
            system_prompt = SystemMessage(
                content="You are a helpful assistant. Answer the user's question directly."
            )

        response = self.llm.invoke([system_prompt, HumanMessage(content=last_message)])
        return {"messages": [AIMessage(content=response.content)]}

    def _route_decision(self, state: AgentState) -> str:
        """Return the routing decision from state."""
        return state["route_decision"]

    def ask(self, question: str) -> str:
        """Run a question through the agent and return the response."""
        result = self.graph.invoke({"messages": [HumanMessage(content=question)]})
        return result["messages"][-1].content
Test the agent with different question types

The true test of the routing logic is asking two different kinds of questions. A company policy question should route through retrieval. A general knowledge question should skip retrieval entirely.

  • Test the agent by running this command in your terminal:
python3 -c "
from agent import RAGAgent
agent = RAGAgent()

print('--- Policy Question ---')
print(agent.ask('What is the company vacation policy?'))

print('\n--- General Knowledge Question ---')
print(agent.ask('What is 2 + 2?'))
"

What should I see?

The first question should return an answer that references vacation days, accrual rates, or the HR portal. These details come from your vacation_policy.txt document.

The second question should return a simple answer like "4" without any reference to company documents. This confirms the router correctly identified it as a general knowledge question and skipped retrieval.

Getting an API error or timeout?

Make sure your GEMINI_API_KEY environment variable is still set in your current terminal session. If you opened a new terminal, you need to activate your virtual environment and export the key again.

If you see a rate limit error, wait 30 seconds and try again. The Gemini free tier has per-minute rate limits.

Help me debug my RAG agent API errors

Why conditional routing instead of always retrieving?

Always querying the vector store wastes time and API calls on questions that do not need document context. A routing node lets the agent skip retrieval for simple questions like math or general knowledge. This pattern mirrors how production agents handle diverse query types efficiently.

Your agent now intelligently routes questions through different paths in the graph. Next up, you will wire in MLflow tracing to observe every node execution, LLM call, and retrieval in a visual dashboard.

Add MLflow Tracing

Your RAG agent routes questions, retrieves document chunks, and generates answers. It works. But right now the entire process is a black box.

How long does each Gemini call take? How many tokens is the agent consuming per question? Which node is the bottleneck? MLflow tracing answers all of these questions by auto-instrumenting every LangChain and LangGraph call with zero manual code.

In this step, get ready to:
  • Enable MLflow auto-instrumentation with a single line of code.
  • Run your agent and generate trace data for multiple questions.
  • Explore the MLflow dashboard to inspect node transitions, LLM calls, and retrieval metadata.
Enable MLflow auto-instrumentation

In production LLMOps, observability means capturing every LLM call, retrieval, and routing decision as a structured trace. MLflow's langchain.autolog() hooks into LangChain's callback system to record all of this automatically.

  • Create a new file called main.py in your project folder. This is the entry point that runs your agent with tracing enabled.
  • Add the MLflow import and autolog activation by pasting this code:
import mlflow
from agent import RAGAgent

# Enable automatic tracing for all LangChain/LangGraph calls
mlflow.langchain.autolog()

What does this code do?

The import mlflow line loads the MLflow tracing library. The from agent import RAGAgent line imports the agent class you built in the previous step.

The key line is mlflow.langchain.autolog(). It injects a callback (MlflowLangchainTracer) into every LangChain invocation. This tracer automatically records span type, action name, input, output, latency, and token usage for each step in your graph.

  • Now add the main() function below the autolog line. This runs four questions through your agent to generate trace data:
def main():
    agent = RAGAgent(documents_path="./documents")

    questions = [
        "What is the company vacation policy?",
        "What is 2 + 2?",
        "How do I request time off?",
        "What is the capital of France?",
    ]

    for question in questions:
        print(f"\nQ: {question}")
        answer = agent.ask(question)
        print(f"A: {answer}")

    print("\n--- Done! Run 'mlflow ui' to view traces at http://localhost:5000 ---")


if __name__ == "__main__":
    main()

What does this code do?

The questions list mixes document-related questions ("What is the company vacation policy?") with general knowledge questions ("What is 2 + 2?"). This generates traces that show both retrieval and direct-answer routing paths.

The loop calls agent.ask() for each question and prints the response. Because autolog is active, MLflow silently records a complete trace for every invocation.

  • Save main.py.

✔️ Awesome, I've got everything!

Great. Double check you've saved your file before running it.

ⓧ I'd like to double check the full code

import mlflow
from agent import RAGAgent

mlflow.langchain.autolog()

def main():
    agent = RAGAgent(documents_path="./documents")

    questions = [
        "What is the company vacation policy?",
        "What is 2 + 2?",
        "How do I request time off?",
        "What is the capital of France?",
    ]

    for question in questions:
        print(f"\nQ: {question}")
        answer = agent.ask(question)
        print(f"A: {answer}")

    print("\n--- Done! Run 'mlflow ui' to view traces at http://localhost:5000 ---")


if __name__ == "__main__":
    main()
Run the agent and launch the MLflow UI

With tracing enabled, every call your agent makes is being recorded. Time to generate some data and visualize it.

  • Run your agent script by executing this command:
python3 main.py

You should see each question printed with its answer. The two document-related questions should return policy information, while the general knowledge questions should answer directly.

Seeing an import error?

  • Make sure your virtual environment is still activated. You should see (venv) at the start of your terminal prompt.
  • Confirm main.py is in the same folder as agent.py and knowledge_base.py.
  • If you see a Gemini API error, verify your GEMINI_API_KEY environment variable is still set in this terminal session.

Help me debug my MLflow tracing setup

  • Once the script finishes, open a new terminal tab (keep your current one available). Make sure you activate your virtual environment in the new tab too.
  • Launch the MLflow UI by running:
mlflow ui

You should see output indicating the server is listening at http://127.0.0.1:5000.

  • Open your browser and navigate to http://localhost:5000.

What is the MLflow UI?

The MLflow UI is a local web dashboard for inspecting your agent's execution history. It shows every trace (one per agent.ask() call) with a timeline of spans representing each node your graph executed.

In production LLMOps, this dashboard is how teams monitor latency, catch regressions, and debug unexpected agent behavior.

Explore traces in the MLflow dashboard

The autolog captures detailed metadata for every node in your LangGraph agent. This includes the retrieval step, LLM calls, and routing decisions.

  • In the MLflow UI, click on the Traces tab in the left sidebar.
  • You should see 4 traces listed (one per question you ran).
  • Click on the trace for "What is the company vacation policy?" to expand it.

Inside the trace, you can see a tree of spans showing the full execution path through your agent.

  • Look for spans labeled with your node names: router, retrieve, and generate. Each span shows its input, output, and duration.
  • Click on an LLM span (the Gemini call inside router or generate) to see token usage and latency for that specific call.
  • Compare the "What is 2 + 2?" trace against the vacation policy trace. The general knowledge question skips the retrieve span entirely, going straight from router to generate.

What does autolog capture for retrievals?

The mlflow.langchain.autolog() function auto-instruments all LangChain operations including retriever calls. For the retrieval span, MLflow captures the query text sent to FAISS, the documents returned, and the execution time.

This means you get full visibility into your RAG pipeline without writing any manual tracing code. Every invoke call, every similarity_search, and every LLM generation is recorded as a span in the trace.

Compare latency across nodes

Look at the duration column for each span. You will likely notice that LLM calls (Gemini) dominate the total latency, while FAISS retrieval is nearly instant. This is the kind of insight production teams use to decide where to optimize.

You now have full LLMOps observability over your RAG agent. Every routing decision, retrieval, and generation is traced and visible in a production-grade dashboard. Next up, you will push your agent further with a Secret Mission challenge.

Secret mission

Add Human-in-the-Loop with LangGraph Interrupts

Your agent answers any question without hesitation. In production, sensitive topics like salaries and terminations need human oversight before the AI responds. Can you add a governance layer that pauses the agent and waits for human approval?

Clean Up Your Resources

Clean Up Your Resources

This project runs entirely locally and uses the free Gemini API tier, so there are no ongoing costs. Decide whether to keep your resources running, pause them to come back later, or delete them entirely.

Resources you used:

  • MLflow tracking server (running at http://localhost:5000).
  • Python virtual environment (venv).
  • Project folder on your Desktop (contains knowledge_base.py, agent.py, main.py, documents/, and faiss_index/).
  • Gemini API key (created in Google AI Studio).
  • GEMINI_API_KEY environment variable.

Keep everything running

No action needed. Choose this if you're still actively building or plan to extend the agent with more features.

  • Your MLflow server will continue to collect traces as you experiment.
  • Your virtual environment and project files stay on disk ready to use.
  • Your Gemini API key remains active for future projects.

Pause - I'll come back to this later

Stop running processes to free up memory, but keep your files for later.

  • Stop the MLflow server by pressing Ctrl+C in the terminal where it is running.
  • Deactivate your virtual environment by running:
deactivate

Your project files, FAISS index, and API key all remain intact. When you come back, reactivate the environment and restart MLflow.

Delete - I don't want to use this again

Remove all project resources and start fresh.

  • Stop the MLflow server by pressing Ctrl+C in the terminal where it is running.
  • Deactivate your virtual environment by running:
deactivate
  • Unset the environment variable by running:
unset GEMINI_API_KEY
  • Delete the project folder by running:
cd ~/Desktop
rm -rf [[PROJECT_FOLDER="your project folder name"]]

On Windows, use this instead:

cd ~\Desktop
rmdir /s /q [[PROJECT_FOLDER="your project folder name"]]
  • Delete your Gemini API key from Google AI Studio. Go to https://aistudio.google.com, open the API keys page, and delete the key you created for this project.

Nice Work!

Nice Work!

You just built a production-grade RAG-powered document Q&A agent from scratch. This project brought together agentic AI development, vector database retrieval, and LLMOps observability in a single, modular codebase.

Here's what you accomplished:

  • Built a LangGraph agent with typed state, conditional routing, and modular OOP design patterns.
  • Created a FAISS vector store for semantic document retrieval using the RAG pattern.
  • Integrated MLflow tracing for production LLMOps observability across every agent invocation.
  • Secret Mission: Added human-in-the-loop governance using LangGraph interrupts and checkpointing.

Ready to quiz yourself?