Build a Self-Correcting RAG Agent with LangGraph and Milvus

This guide is adapted from the Milvus Bootcamp's langgraph-rag-agent-local.ipynb, published under the Apache-2.0 licence.

Retrieval-Augmented Generation works well for single-turn lookups, but it breaks down the moment a question requires the system to judge whether it found the right information, correct a hallucinated answer, or fall back to a different data source mid-flight. LangGraph solves this by treating the RAG pipeline as a stateful graph where each node — retrieve, grade, generate, web-search — is a discrete function connected by conditional edges. The graph loops until the answer is both grounded in evidence and genuinely responsive to the question, then terminates. The result is a self-correcting agent that behaves more like a careful researcher than a one-shot autocomplete call.

The architecture pulls together three published ideas: adaptive routing (sending questions to the right data source before retrieval starts), corrective RAG (falling back to web search when local documents fail the relevance test), and self-RAG (re-generating or re-searching when the produced answer contains hallucinations). All of it runs locally — no OpenAI API key required — which matters for organisations with data-residency constraints, teams experimenting on laptops without per-token costs, or anyone who wants to inspect every intermediate model call. If you are interested in the broader landscape of agent design patterns, the AI Mastery roundup on agentic APIs and MCP is a useful companion read.

Running this locally has real resource requirements. Llama 3 via Ollama requires at least 16 GB of RAM (32 GB is comfortable). Embedding with HuggingFaceEmbeddings adds the sentence-transformers model download on first use. The Milvus Lite backend stores the vector index in a single .db file on disk — no separate server process is needed. Tavily web search requires a free API key and counts against your monthly quota.

Prerequisites

Step 1: Install dependencies and configure the environment

Everything the agent needs is installable in one command. The package list covers the LangChain ecosystem, the Milvus client, LangGraph itself, the Tavily search tool, and the HuggingFace sentence-transformer library used for local embeddings.

! pip install -U langchain_community tiktoken langchainhub pymilvus langchain langgraph tavily-python sentence-transformers langchain-milvus langchain-huggingface

After installation, load environment variables from a .env file. At minimum that file must contain TAVILY_API_KEY. Turning on LangChain's debug and verbose modes is optional but valuable while learning the graph — every prompt, every intermediate output, and every routing decision is printed to stdout.

from dotenv import load_dotenv
import os

load_dotenv()

from langchain.globals import set_verbose, set_debug

set_debug(True)
set_verbose(True)

### LLM

local_llm = 'llama3'

The local_llm string is consumed by every ChatOllama instantiation later. Changing it here — to 'llama3:70b' or 'mistral', for instance — propagates through the whole graph without hunting for individual references.

Step 2: Build the Milvus vector store and retriever

The knowledge base for this agent is three long-form blog posts from Lilian Weng covering LLM agents, prompt engineering, and adversarial attacks against language models. WebBaseLoader fetches and parses each URL. RecursiveCharacterTextSplitter breaks the resulting documents into 250-token chunks with no overlap — shorter chunks improve retrieval precision at the cost of losing some surrounding context. The chunks are embedded with HuggingFaceEmbeddings and written into a local Milvus Lite database file named milvus_rag.db.

### Index

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import WebBaseLoader
from langchain_milvus import Milvus
from langchain_community.embeddings import HuggingFaceEmbeddings

urls = [
    "https://lilianweng.github.io/posts/2023-06-23-agent/",
    "https://lilianweng.github.io/posts/2023-03-15-prompt-engineering/",
    "https://lilianweng.github.io/posts/2023-10-25-adv-attack-llm/",
]

docs = [WebBaseLoader(url).load() for url in urls]
docs_list = [item for sublist in docs for item in sublist]
text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
    chunk_size=250, chunk_overlap=0
)
doc_splits = text_splitter.split_documents(docs_list)

# Add to Milvus
vectorstore = Milvus.from_documents(
    documents=doc_splits,
    collection_name="rag_milvus",
    embedding=HuggingFaceEmbeddings(),
    connection_args={"uri": "./milvus_rag.db"},

)
retriever = vectorstore.as_retriever()

The connection_args={"uri": "./milvus_rag.db"} line is what makes this truly local — Milvus Lite runs in-process and persists to a file, requiring no Docker container or separate server. Re-running this cell on an existing .db file appends duplicates rather than replacing them, so delete the file between experiments.

Step 3: Build the grader and generator chains

This is where the self-correcting intelligence lives. Four separate chains are constructed, each a prompt-plus-model pipeline. The retrieval grader and the two quality graders (hallucination_grader and answer_grader) all use format="json" and temperature=0 to force deterministic, structured outputs — a {"score": "yes"} or {"score": "no"} — that the graph's conditional edges can act on reliably.

### Retrieval Grader 

from langchain.prompts import PromptTemplate
from langchain_community.chat_models import ChatOllama
from langchain_core.output_parsers import JsonOutputParser

# LLM
llm = ChatOllama(model=local_llm, format="json", temperature=0)

prompt = PromptTemplate(
    template="""You are a grader assessing relevance 
    of a retrieved document to a user question. If the document contains keywords related to the user question, 
    grade it as relevant. It does not need to be a stringent test. The goal is to filter out erroneous retrievals. 
    
    Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question.
    Provide the binary score as a JSON with a single key 'score' and no premable or explaination.
     
    Here is the retrieved document: 
    {document}
    
    Here is the user question: 
    {question}
    """,
    input_variables=["question", "document"],
)

retrieval_grader = prompt | llm | JsonOutputParser()
question = "agent memory"
docs = retriever.invoke(question)
doc_txt = docs[1].page_content
print(retrieval_grader.invoke({"question": question, "document": doc_txt}))

The RAG generation chain uses a plain StrOutputParser because here the model is writing a prose answer rather than filling a structured schema. The three-sentence ceiling in the prompt keeps answers short enough to evaluate programmatically in the next stage.

### Generate

from langchain.prompts import PromptTemplate
from langchain import hub
from langchain_core.output_parsers import StrOutputParser

# Prompt
prompt = PromptTemplate(
    template="""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:
    Question: {question} 
    Context: {context} 
    Answer: 
    """,
    input_variables=["question", "document"],
)

llm = ChatOllama(model=local_llm, temperature=0)

# Post-processing
def format_docs(docs):
    return "\n\n".join(doc.page_content for doc in docs)

# Chain
rag_chain = prompt | llm | StrOutputParser()

# Run
question = "agent memory"
docs = retriever.invoke(question)
generation = rag_chain.invoke({"context": docs, "question": question})
print(generation)

The hallucination and answer graders each independently score a different property of the output:

### Hallucination Grader 

# LLM
llm = ChatOllama(model=local_llm, format="json", temperature=0)

# Prompt
prompt = PromptTemplate(
    template="""You are a grader assessing whether 
    an answer is grounded in / supported by a set of facts. Give a binary score 'yes' or 'no' score to indicate 
    whether the answer is grounded in / supported by a set of facts. Provide the binary score as a JSON with a 
    single key 'score' and no preamble or explanation.
    
    Here are the facts:
    {documents} 

    Here is the answer: 
    {generation}
    """,
    input_variables=["generation", "documents"],
)

hallucination_grader = prompt | llm | JsonOutputParser()
hallucination_grader.invoke({"documents": docs, "generation": generation})

### Answer Grader 

# LLM
llm = ChatOllama(model=local_llm, format="json", temperature=0)

# Prompt
prompt = PromptTemplate(
    template="""You are a grader assessing whether an 
    answer is useful to resolve a question. Give a binary score 'yes' or 'no' to indicate whether the answer is 
    useful to resolve a question. Provide the binary score as a JSON with a single key 'score' and no preamble or explanation.
     
    Here is the answer:
    {generation} 

    Here is the question: {question}
    """,
    input_variables=["generation", "question"],
)

answer_grader = prompt | llm | JsonOutputParser()
answer_grader.invoke({"question": question,"generation": generation})

The question router is a fifth chain that decides — before any retrieval happens — whether the question belongs in the vector store or should go directly to web search. Questions about LLM agents, prompt engineering, and adversarial attacks get routed to Milvus; everything else goes to Tavily.

### Router

from langchain.prompts import PromptTemplate
from langchain_community.chat_models import ChatOllama
from langchain_core.output_parsers import JsonOutputParser

# LLM
llm = ChatOllama(model=local_llm, format="json", temperature=0)

prompt = PromptTemplate(
    template="""You are an expert at routing a 
    user question to a vectorstore or web search. Use the vectorstore for questions on LLM  agents, 
    prompt engineering, and adversarial attacks. You do not need to be stringent with the keywords 
    in the question related to these topics. Otherwise, use web-search. Give a binary choice 'web_search' 
    or 'vectorstore' based on the question. Return the a JSON with a single key 'datasource' and 
    no premable or explaination. 
    
    Question to route: 
    {question}""",
    input_variables=["question"],
)

question_router = prompt | llm | JsonOutputParser()
question = "llm agent memory"
docs = retriever.get_relevant_documents(question)
doc_txt = docs[1].page_content
print(question_router.invoke({"question": question}))

### Search

from langchain_community.tools.tavily_search import TavilySearchResults
web_search_tool = TavilySearchResults(k=3)

Step 4: Assemble the LangGraph state machine

All five chains now need to be wired into a graph. LangGraph requires a typed state dictionary that every node reads from and writes to. Each node function accepts the full state, does its work, and returns only the keys it changed — LangGraph merges those changes back into the shared state automatically.

from typing_extensions import TypedDict
from typing import List

### State

class GraphState(TypedDict):
    """
    Represents the state of our graph.

    Attributes:
        question: question
        generation: LLM generation
        web_search: whether to add search
        documents: list of documents 
    """
    question : str
    generation : str
    web_search : str
    documents : List[str]

from langchain.schema import Document

### Nodes

def retrieve(state):
    """
    Retrieve documents from vectorstore

    Args:
        state (dict): The current graph state

    Returns:
        state (dict): New key added to state, documents, that contains retrieved documents
    """
    print("---RETRIEVE---")
    question = state["question"]

    # Retrieval
    documents = retriever.invoke(question)
    return {"documents": documents, "question": question}

def generate(state):
    """
    Generate answer using RAG on retrieved documents

    Args:
        state (dict): The current graph state

    Returns:
        state (dict): New key added to state, generation, that contains LLM generation
    """
    print("---GENERATE---")
    question = state["question"]
    documents = state["documents"]
    
    # RAG generation
    generation = rag_chain.invoke({"context": documents, "question": question})
    return {"documents": documents, "question": question, "generation": generation}

def grade_documents(state):
    """
    Determines whether the retrieved documents are relevant to the question
    If any document is not relevant, we will set a flag to run web search

    Args:
        state (dict): The current graph state

    Returns:
        state (dict): Filtered out irrelevant documents and updated web_search state
    """

    print("---CHECK DOCUMENT RELEVANCE TO QUESTION---")
    question = state["question"]
    documents = state["documents"]
    
    # Score each doc
    filtered_docs = []
    web_search = "No"
    for d in documents:
        score = retrieval_grader.invoke({"question": question, "document": d.page_content})
        grade = score['score']
        # Document relevant
        if grade.lower() == "yes":
            print("---GRADE: DOCUMENT RELEVANT---")
            filtered_docs.append(d)
        # Document not relevant
        else:
            print("---GRADE: DOCUMENT NOT RELEVANT---")
            # We do not include the document in filtered_docs
            # We set a flag to indicate that we want to run web search
            web_search = "Yes"
            continue
    return {"documents": filtered_docs, "question": question, "web_search": web_search}
    
def web_search(state):
    """
    Web search based based on the question

    Args:
        state (dict): The current graph state

    Returns:
        state (dict): Appended web results to documents
    """

    print("---WEB SEARCH---")
    question = state["question"]
    documents = state["documents"]

    # Web search
    docs = web_search_tool.invoke({"query": question})
    web_results = "\n".join([d["content"] for d in docs])
    web_results = Document(page_content=web_results)
    if documents is not None:
        documents.append(web_results)
    else:
        documents = [web_results]
    return {"documents": documents, "question": question}

### Conditional edge

def route_question(state):
    """
    Route question to web search or RAG.

    Args:
        state (dict): The current graph state

    Returns:
        str: Next node to call
    """

    print("---ROUTE QUESTION---")
    question = state["question"]
    print(question)
    source = question_router.invoke({"question": question})  
    print(source)
    print(source['datasource'])
    if source['datasource'] == 'web_search':
        print("---ROUTE QUESTION TO WEB SEARCH---")
        return "websearch"
    elif source['datasource'] == 'vectorstore':
        print("---ROUTE QUESTION TO RAG---")
        return "vectorstore"

def decide_to_generate(state):
    """
    Determines whether to generate an answer, or add web search

    Args:
        state (dict): The current graph state

    Returns:
        str: Binary decision for next node to call
    """

    print("---ASSESS GRADED DOCUMENTS---")
    question = state["question"]
    web_search = state["web_search"]
    filtered_documents = state["documents"]

    if web_search == "Yes":
        # All documents have been filtered check_relevance
        # We will re-generate a new query
        print("---DECISION: ALL DOCUMENTS ARE NOT RELEVANT TO QUESTION, INCLUDE WEB SEARCH---")
        return "websearch"
    else:
        # We have relevant documents, so generate answer
        print("---DECISION: GENERATE---")
        return "generate"

### Conditional edge

def grade_generation_v_documents_and_question(state):
    """
    Determines whether the generation is grounded in the document and answers question.

    Args:
        state (dict): The current graph state

    Returns:
        str: Decision for next node to call
    """

    print("---CHECK HALLUCINATIONS---")
    question = state["question"]
    documents = state["documents"]
    generation = state["generation"]

    score = hallucination_grader.invoke({"documents": documents, "generation": generation})
    grade = score['score']

    # Check hallucination
    if grade == "yes":
        print("---DECISION: GENERATION IS GROUNDED IN DOCUMENTS---")
        # Check question-answering
        print("---GRADE GENERATION vs QUESTION---")
        score = answer_grader.invoke({"question": question,"generation": generation})
        grade = score['score']
        if grade == "yes":
            print("---DECISION: GENERATION ADDRESSES QUESTION---")
            return "useful"
        else:
            print("---DECISION: GENERATION DOES NOT ADDRESS QUESTION---")
            return "not useful"
    else:
        pprint("---DECISION: GENERATION IS NOT GROUNDED IN DOCUMENTS, RE-TRY---")
        return "not supported"

from langgraph.graph import END, StateGraph
workflow = StateGraph(GraphState)

# Define the nodes
workflow.add_node("websearch", web_search) # web search
workflow.add_node("retrieve", retrieve) # retrieve
workflow.add_node("grade_documents", grade_documents) # grade documents
workflow.add_node("generate", generate) # generatae

With nodes registered, the edges declare the legal paths through the graph. The entry point is conditional — the router inspects the question before any node runs. After grading, documents that all fail relevance trigger a web search fallback. After generation, the hallucination grader can route back to generate for a retry, or forward to END if everything checks out.

# Build graph
workflow.set_conditional_entry_point(
    route_question,
    {
        "websearch": "websearch",
        "vectorstore": "retrieve",
    },
)

workflow.add_edge("retrieve", "grade_documents")
workflow.add_conditional_edges(
    "grade_documents",
    decide_to_generate,
    {
        "websearch": "websearch",
        "generate": "generate",
    },
)
workflow.add_edge("websearch", "generate")
workflow.add_conditional_edges(
    "generate",
    grade_generation_v_documents_and_question,
    {
        "not supported": "generate",
        "useful": END,
        "not useful": "websearch",
    },
)

Step 5: Compile and run the agent

Compiling converts the StateGraph into an executable app. The three test questions below are deliberately chosen to exercise different paths: the first is in-domain for the vector store, the second is an out-of-domain current-events question (NFL draft), and the third is a recent news question that would not exist in any static knowledge base.

# Compile
app = workflow.compile()

# Test
from pprint import pprint
inputs = {"question": "What are the types of agent memory?"}
for output in app.stream(inputs):
    for key, value in output.items():
        pprint(f"Finished running: {key}:")
pprint(value["generation"])

# Compile
app = workflow.compile()

# Test
from pprint import pprint
inputs = {"question": "Who are the Bears expected to draft first in the NFL draft?"}
for output in app.stream(inputs):
    for key, value in output.items():
        pprint(f"Finished running: {key}:")
pprint(value["generation"])

# Test
from pprint import pprint
inputs = {"question": "Did Emmanuel Macron visit Germany recently?"}
for output in app.stream(inputs):
    for key, value in output.items():
        pprint(f"Finished running: {key}:")
pprint(value["generation"])

app.stream yields one dictionary per node execution as the graph progresses, so you can watch routing decisions happen in real time rather than waiting for a final answer. Each dictionary key is the node name that just finished; value holds the full graph state at that point.

What to watch out for

Infinite hallucination loops. The "not supported" branch in grade_generation_v_documents_and_question routes straight back to generate with no iteration limit. If the model consistently produces answers the grader considers ungrounded — which happens when the model is at the edge of its context window or when retrieved chunks are too sparse — the graph will loop indefinitely. LangGraph supports a recursion_limit parameter on app.stream; set it explicitly in production, for example app.stream(inputs, {"recursion_limit": 10}).

JSON format failures from the LLM. All three graders depend on Ollama returning valid JSON with format="json". Smaller quantised models sometimes emit a JSON blob with extra keys, trailing commas, or prose before the opening brace, causing JsonOutputParser to raise. Wrapping grader invocations in a try/except and defaulting to "yes" on parse failure is a pragmatic fallback, though it silently disables grading when it matters most.

Milvus Lite duplicate ingestion. Re-running Step 2 without deleting milvus_rag.db duplicates every chunk. The retriever will still return results, but the grader scores redundant copies and latency rises on every query. Delete the file or add a collection existence check before calling Milvus.from_documents.

Tavily quota exhaustion. The free tier allows 1,000 searches per month. A single complex question that fails the relevance grader on every retrieved document can trigger multiple web searches within one graph execution. Monitor Tavily usage separately from your LLM usage.

Embedding model cold starts. The first call to HuggingFaceEmbeddings() downloads model weights from HuggingFace Hub if they are not cached. Set TRANSFORMERS_CACHE to a fast local SSD path to avoid re-downloading across sessions.

pprint reference error in the hallucination grader. The source notebook calls pprint(...) inside grade_generation_v_documents_and_question before pprint is imported at the module level. If you define this function in a fresh cell before the from pprint import pprint statement that appears in Step 5, Python will raise a NameError. Import pprint at the top of your notebook or at the top of the cell that defines the grader function.

Where to go next

The graph here is deliberately simple — four nodes, two conditional edge functions, one loop. Real deployments typically add a query-rewriting node between routing and retrieval, replace the LLM-based relevance grader with a cross-encoder re-ranker for higher precision, and swap Milvus Lite for a Milvus standalone or cluster deployment for concurrent access. To understand the failure modes of agentic systems at scale — including the risks of agents operating outside their sanctioned boundaries — the AI Mastery analysis of rogue agent incidents is worth reading alongside this practical implementation. The Milvus Bootcamp repository also contains notebook variants that swap Ollama for other local inference backends and extend the graph with multi-query expansion, both of which address the retrieval precision ceiling that a 250-token chunk size imposes.