Retrieval-Augmented Generation with LlamaIndex and Python

By · · AI Engineering

Large language models are impressive, but they hallucinate. They confidently generate answers that sound correct but are entirely fabricated, especially when asked about proprietary data, recent events, or niche domains they were not trained on. Retrieval-Augmented Generation, or RAG, addresses this by retrieving relevant documents before generating a response, grounding the LLM's output in actual data. LlamaIndex is one of the best frameworks for building RAG pipelines in Python, and in this post we will build one from scratch.

How RAG works

The RAG pattern has three stages:

  1. Indexing: Documents are loaded, split into chunks, converted to vector embeddings, and stored in a vector index.
  2. Retrieval: When a user asks a question, the query is embedded and the most similar document chunks are retrieved from the index.
  3. Generation: The retrieved chunks are injected into the LLM prompt as context, and the model generates an answer grounded in that context.

This approach lets you build Q&A systems over your own data without fine-tuning a model.

Installation

Install LlamaIndex and the OpenAI integration:

pip install llama-index llama-index-llms-openai llama-index-embeddings-openai

Set your API key:

export OPENAI_API_KEY="sk-your-api-key-here"

Loading documents

LlamaIndex provides document loaders for many formats. The simplest is SimpleDirectoryReader, which reads all files from a directory:

from llama_index.core import SimpleDirectoryReader

documents = SimpleDirectoryReader("./data").load_data()
print(f"Loaded {len(documents)} documents")

for doc in documents[:3]:
    print(f"  - {doc.metadata.get('file_name', 'unknown')}: {len(doc.text)} chars")

Place your text files, PDFs, or markdown files in a ./data directory. LlamaIndex handles the parsing automatically. For this tutorial, imagine we have a collection of internal engineering documents that we want to make searchable.

Configuring the service context

Before building the index, configure the LLM and embedding model:

from llama_index.core import Settings
from llama_index.llms.openai import OpenAI
from llama_index.embeddings.openai import OpenAIEmbedding

Settings.llm = OpenAI(model="gpt-3.5-turbo", temperature=0)
Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small")
Settings.chunk_size = 512
Settings.chunk_overlap = 50

The chunk_size controls how documents are split. Smaller chunks mean more precise retrieval but less context per chunk. An overlap of 50 tokens ensures that information spanning a chunk boundary is not lost.

Building a vector store index

The vector store index is the core of the RAG pipeline. It converts document chunks into embeddings and stores them for similarity search:

from llama_index.core import VectorStoreIndex

index = VectorStoreIndex.from_documents(documents)
print("Index built successfully")

Under the hood, this splits each document into chunks, generates embeddings using the configured model, and stores everything in an in-memory vector store. For production use, you would connect to a persistent store like ChromaDB or Pinecone.

Querying the index

The simplest way to query is with a query engine:

query_engine = index.as_query_engine(similarity_top_k=3)

response = query_engine.query("What is our deployment process for production releases?")
print(response)

The similarity_top_k parameter controls how many chunks are retrieved. More chunks provide more context but also increase token usage and can dilute relevance.

You can also inspect which source documents were used:

for node in response.source_nodes:
    print(f"Score: {node.score:.4f}")
    print(f"Source: {node.metadata.get('file_name', 'unknown')}")
    print(f"Text: {node.text[:200]}...")
    print("---")

This transparency is one of the key advantages of RAG. You can always trace an answer back to its source documents, which is essential for trust and debugging.

Using a persistent vector store

The in-memory store is fine for development, but production systems need persistence. Here is how to use ChromaDB:

pip install llama-index-vector-stores-chroma chromadb

import chromadb
from llama_index.vector_stores.chroma import ChromaVectorStore
from llama_index.core import StorageContext

chroma_client = chromadb.PersistentClient(path="./chroma_db")
chroma_collection = chroma_client.get_or_create_collection("engineering_docs")

vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
storage_context = StorageContext.from_defaults(vector_store=vector_store)

# Build the index with persistent storage
index = VectorStoreIndex.from_documents(
    documents,
    storage_context=storage_context,
)

# On subsequent runs, load the existing index instead of rebuilding
# index = VectorStoreIndex.from_vector_store(vector_store)

With persistent storage, you only need to index documents once. New documents can be added incrementally without rebuilding the entire index.

Chat engine for conversational RAG

A query engine handles one-off questions, but a chat engine maintains conversation history:

chat_engine = index.as_chat_engine(
    chat_mode="context",
    similarity_top_k=3,
    system_prompt="""You are a helpful assistant that answers questions about
    our engineering documentation. Always cite the source document when
    providing information. If you cannot find the answer in the provided
    context, say so clearly.""",
)

response1 = chat_engine.chat("What testing frameworks do we use?")
print(f"Assistant: {response1}")

response2 = chat_engine.chat("How do we run those tests in CI?")
print(f"Assistant: {response2}")

The second question references "those tests" from the first response. The chat engine handles this by maintaining history and performing context-aware retrieval.

Customizing the retrieval pipeline

For more control, you can build a custom retrieval pipeline using retrievers and response synthesizers separately:

from llama_index.core import get_response_synthesizer
from llama_index.core.retrievers import VectorIndexRetriever
from llama_index.core.query_engine import RetrieverQueryEngine

retriever = VectorIndexRetriever(
    index=index,
    similarity_top_k=5,
)

response_synthesizer = get_response_synthesizer(
    response_mode="compact",
)

query_engine = RetrieverQueryEngine(
    retriever=retriever,
    response_synthesizer=response_synthesizer,
)

response = query_engine.query("What are the key architecture decisions?")
print(response)

The response_mode parameter controls how retrieved chunks are combined:

Evaluation: is your RAG system working?

Building a RAG pipeline is easy. Building one that actually gives good answers is harder. You need to evaluate both retrieval quality and generation quality:

from llama_index.core.evaluation import RelevancyEvaluator

evaluator = RelevancyEvaluator()

query = "What is our incident response process?"
response = query_engine.query(query)

eval_result = evaluator.evaluate_response(query=query, response=response)
print(f"Is relevant: {eval_result.passing}")
print(f"Feedback: {eval_result.feedback}")

You should also build a test set of question-answer pairs and systematically evaluate retrieval recall and answer accuracy. A RAG system that retrieves the wrong chunks will generate confidently wrong answers, which is worse than no answer at all.

Production considerations

Before deploying a RAG system, consider these practical issues:

RAG is not a silver bullet, but it is the most practical way to give LLMs access to your private data today. LlamaIndex handles the heavy lifting of document processing, embedding, retrieval, and response synthesis so you can focus on the data and the user experience. Start with the basic pipeline, evaluate rigorously, and iterate on chunk size and retrieval parameters until the answers meet your quality bar.