DEV Community

Syeed Talha
Syeed Talha

Posted on

2-Step RAG in Langchain

If you've ever asked ChatGPT something about your own documents and got a wrong answer, you already understand the problem RAG solves.

LLMs are powerful, but they only know what they were trained on. They have no idea about your internal docs, your company wiki, or anything that happened after their training cutoff. RAG fixes that.

In this post, we'll build the most basic version of RAG — a 2-step RAG — from scratch using LangChain and FAISS. No fluff, just the core idea working in real code.


What is RAG?

RAG stands for Retrieval-Augmented Generation. The name sounds fancy but the idea is simple:

  1. Retrieve — find the most relevant documents for a given question
  2. Generate — feed those documents to the LLM and let it answer

Instead of expecting the LLM to magically know everything, you hand it the right information first, then ask the question. That's it.


Why "2-Step"?

Because the flow is literally two steps:

Query → retrieve relevant docs → LLM sees docs + query → answer
Enter fullscreen mode Exit fullscreen mode

No agents, no complex chains, no fancy orchestration. Just retrieve, then generate. This is the foundation that every advanced RAG system is built on top of.


What We're Building

We have a small knowledge base of 5 documents about LangChain concepts. When a user asks a question, we:

  1. Find the 2 most relevant documents from that knowledge base
  2. Pass them to an LLM as context
  3. Get a grounded answer back

Let's walk through the code.


Step 1: The Knowledge Base

from langchain_core.documents import Document

docs = [
    Document(page_content="LangChain is a framework for building LLM applications. It provides tools for chaining, agents, and retrieval."),
    Document(page_content="Agents in LangChain use an LLM to decide which tools to call. They can reason step-by-step."),
    Document(page_content="RAG stands for Retrieval-Augmented Generation. It combines retrieval with LLM generation for grounded answers."),
    Document(page_content="Vector stores like FAISS and Chroma store embeddings and enable semantic search over documents."),
    Document(page_content="LangGraph is a library for building stateful, multi-actor applications with LLMs."),
]
Enter fullscreen mode Exit fullscreen mode

In a real project, these documents would come from a PDF, a database, a website, or any data source you care about. Here we're keeping it simple with hardcoded strings.


Step 2: Building the Vector Store

This is where the "magic" actually lives, so let's break it down carefully.

from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS

embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
vectorstore = FAISS.from_documents(docs, embeddings)
retriever = vectorstore.as_retriever(search_kwargs={"k": 2})
Enter fullscreen mode Exit fullscreen mode

What's happening here?

First, HuggingFaceEmbeddings loads a model that converts text into vectors — lists of numbers that represent the meaning of a sentence. Similar sentences produce vectors that are numerically close to each other.

"RAG stands for Retrieval-Augmented Generation..."
                    ↓
        [0.87, 0.02, 0.11, ...]  ← 384 numbers
Enter fullscreen mode Exit fullscreen mode

Then FAISS.from_documents() takes all 5 documents, converts each one into a vector, and stores them in a FAISS index — an in-memory database built specifically for fast similarity search.

Finally, .as_retriever(search_kwargs={"k": 2}) wraps everything into a retriever object. The k=2 means: when someone asks a question, return the 2 most similar documents.


Step 3: The 2-Step RAG Function

def answer(query: str) -> str:
    # Step 1: Retrieve
    retrieved = retriever.invoke(query)
    context = "\n\n".join(d.page_content for d in retrieved)

    # Step 2: Generate
    prompt = f"""Answer the question using ONLY the context below.

Context:
{context}

Question: {query}

Answer:"""
    result = free_llm.invoke(prompt)
    return result.content
Enter fullscreen mode Exit fullscreen mode

Step 1 — Retrieve:

retriever.invoke(query) does the following internally:

  • Embeds the query into a vector using the same HuggingFace model
  • Compares that vector against all stored document vectors
  • Returns the top 2 most similar documents
query = "What is RAG?"
    ↓ embed
[0.85, 0.04, 0.10, ...]
    ↓ compare against FAISS index
"RAG stands for..."   → similarity: 0.96  ✅ top 1
"Vector stores..."    → similarity: 0.71  ✅ top 2
"Agents in..."        → similarity: 0.38  ✗
Enter fullscreen mode Exit fullscreen mode

Step 2 — Generate:

We stitch the retrieved documents into a single context string and build a prompt. Notice the instruction: "Answer using ONLY the context below." This is important — it prevents the LLM from going off-script and making things up outside of what we retrieved.


Full Code

import os
from dotenv import load_dotenv

from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_openai import ChatOpenAI
from langchain_core.documents import Document

load_dotenv()

# 1. Knowledge base
docs = [
    Document(page_content="LangChain is a framework for building LLM applications. It provides tools for chaining, agents, and retrieval."),
    Document(page_content="Agents in LangChain use an LLM to decide which tools to call. They can reason step-by-step."),
    Document(page_content="RAG stands for Retrieval-Augmented Generation. It combines retrieval with LLM generation for grounded answers."),
    Document(page_content="Vector stores like FAISS and Chroma store embeddings and enable semantic search over documents."),
    Document(page_content="LangGraph is a library for building stateful, multi-actor applications with LLMs."),
]

# 2. Build vector store
embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
vectorstore = FAISS.from_documents(docs, embeddings)
retriever = vectorstore.as_retriever(search_kwargs={"k": 2})

# 3. LLM setup
free_llm = ChatOpenAI(
    model="nvidia/nemotron-nano-9b-v2:free",
    base_url="https://openrouter.ai/api/v1",
    api_key=os.environ.get("OPENROUTER_API_KEY"),
)

# 4. RAG function
def answer(query: str) -> str:
    retrieved = retriever.invoke(query)
    context = "\n\n".join(d.page_content for d in retrieved)

    prompt = f"""Answer the question using ONLY the context below.

Context:
{context}

Question: {query}

Answer:"""
    result = free_llm.invoke(prompt)
    return result.content

# 5. Run
queries = [
    "What is RAG?",
    "How do agents work in LangChain?",
]

for q in queries:
    print("\n" + "=" * 60)
    print(f"Q: {q}")
    print(f"A: {answer(q)}")
Enter fullscreen mode Exit fullscreen mode

Sample Output

============================================================
Q: What is RAG?
A: RAG stands for Retrieval-Augmented Generation. It combines
   retrieval with LLM generation to produce grounded answers.

============================================================
Q: How do agents work in LangChain?
A: Agents in LangChain use an LLM to decide which tools to
   call. They can reason step-by-step.
Enter fullscreen mode Exit fullscreen mode

The LLM isn't guessing — it's reading the retrieved documents and answering from them directly.


What to Keep in Mind

The quality of retrieval determines the quality of the answer. If the wrong documents get retrieved, the LLM has no chance of answering correctly. Garbage in, garbage out.

k is a tradeoff. A higher k gives the LLM more context to work with, but also introduces more noise and uses more tokens. For most beginner projects, k=2 or k=3 is a reasonable starting point.

This is the foundation. Every advanced RAG technique — reranking, hybrid search, query expansion, agentic RAG — is built on top of this same core loop. Once you understand retrieve-then-generate, everything else is just an improvement on one of those two steps.


What's Next

Once you're comfortable with 2-step RAG, here's where to go next:

  • Add metadata filtering — retrieve documents by category or date, not just similarity
  • Try reranking — use a second model to reorder retrieved docs by relevance
  • Explore agentic RAG — let the LLM decide when to retrieve and what to search for

But honestly, for most straightforward use cases, this basic 2-step approach works surprisingly well. Start here, measure where it fails, then add complexity only where you need it.

Top comments (0)