Building a RAG Agent with ChromaDB
RAG Pipeline Architecture â Complete Production Flow
What is Retrieval-Augmented Generation?
Retrieval-Augmented Generation (RAG) combines information retrieval with language generation. Instead of relying solely on the model's parametric knowledge, RAG agents first retrieve relevant documents from an external knowledge base, then condition the LLM's generation on those documents.
This approach solves the fundamental problem of LLM hallucination by grounding responses in verifiable sources. The agent can cite specific passages, maintain factual accuracy, and access information beyond the model's training cutoff.
RAG is the dominant architecture for knowledge-intensive NLP tasks. It outperforms fine-tuning for most enterprise use cases because it's cheaper, faster to update, and provides traceable citations. The key insight: you don't need to bake knowledge into model weights when you can retrieve it at inference time.
Why This Matters
Without RAG, an LLM is limited to its training data â it can't access your company's documents, recent events, or domain-specific knowledge. RAG transforms a general-purpose LLM into a domain expert by connecting it to your knowledge base.
Real-world analogy: Think of RAG as giving the LLM a library card. Instead of relying on what it memorized in school, it can look up the answer in relevant books before responding. This makes it both more accurate and more trustworthy.
RAG vs Fine-Tuning vs Prompt Engineering
| Approach | Knowledge Update | Cost | Citation Support | Latency | Best For |
|---|---|---|---|---|---|
| Prompt Engineering | Manual | Low | Limited | Low | Simple Q&A |
| RAG | Automatic | Medium | Full | Medium | Knowledge bases |
| Fine-Tuning | Retrain | High | None | Low | Domain adaptation |
| Hybrid | Mixed | High | Partial | Medium | Complex systems |
Key insight: RAG is the only approach that provides automatic knowledge updates AND citation support. Fine-tuning embeds knowledge but can't cite sources. Prompt engineering is limited by context window.
Project Overview
We will build a complete RAG agent that:
- Ingests documents (PDF, TXT, Markdown) with automatic chunking
- Stores embeddings in ChromaDB with metadata filtering
- Performs hybrid search (semantic + keyword matching)
- Generates answers with source citations
- Supports incremental index updates
- Evaluates retrieval quality with metrics (MRR, recall@k)
Expected outcome: A production RAG pipeline you can deploy for any document corpus.
Difficulty: Advanced (requires understanding of embeddings, vector databases, and prompt engineering)
Tools & Setup
| Tool | Version | Purpose |
|---|---|---|
| Python | 3.11+ | Core language |
| ChromaDB | 0.4+ | Vector database |
| OpenAI | 1.0+ | Embeddings + LLM |
| tiktoken | 0.5+ | Token counting |
| pypdf | 4.0+ | PDF parsing |
Step 1: Environment Setup
python -m venv venv
source venv/bin/activate
pip install chromadb openai tiktoken pypdf
export OPENAI_API_KEY="sk-your-key"
Step 2: Document Loader
# loader.py
from pathlib import Path
from typing import Iterator
import logging
logger = logging.getLogger(__name__)
def load_documents(source: str) -> Iterator[dict]:
"""
Load documents from a file or directory.
Args:
source: Path to a file or directory
Yields:
Document dictionaries with content and metadata
"""
path = Path(source)
if path.is_file():
yield from _load_file(path)
elif path.is_dir():
for file in path.rglob("*"):
if file.suffix.lower() in (".txt", ".md", ".pdf"):
yield from _load_file(file)
else:
logger.warning(f"Source not found: {source}")
def _load_file(path: Path) -> Iterator[dict]:
"""Load a single file and yield document(s)."""
suffix = path.suffix.lower()
try:
if suffix == ".pdf":
yield from _load_pdf(path)
else:
content = path.read_text(encoding="utf-8", errors="ignore")
yield {
"content": content,
"metadata": {
"source": str(path),
"filename": path.name,
"type": suffix,
},
}
except Exception as e:
logger.error(f"Failed to load {path}: {e}")
def _load_pdf(path: Path) -> Iterator[dict]:
"""Extract text from PDF pages."""
from pypdf import PdfReader
reader = PdfReader(str(path))
for i, page in enumerate(reader.pages):
text = page.extract_text()
if text:
yield {
"content": text,
"metadata": {
"source": str(path),
"filename": path.name,
"page": i + 1,
"type": ".pdf",
},
}
Step 3: Text Splitter with Overlap
# splitter.py
import tiktoken
from typing import List, Dict
import logging
logger = logging.getLogger(__name__)
class TextSplitter:
"""
Split documents into chunks for embedding.
Uses token-based chunking with overlap to ensure no information
is lost at chunk boundaries.
Args:
chunk_size: Maximum tokens per chunk
chunk_overlap: Number of overlapping tokens between chunks
encoding_name: tiktoken encoding to use
"""
def __init__(
self,
chunk_size: int = 500,
chunk_overlap: int = 50,
encoding_name: str = "cl100k_base",
):
if chunk_overlap >= chunk_size:
raise ValueError("chunk_overlap must be less than chunk_size")
self.chunk_size = chunk_size
self.chunk_overlap = chunk_overlap
self.enc = tiktoken.get_encoding(encoding_name)
def split(self, text: str, metadata: dict = None) -> List[Dict]:
"""
Split text into chunks.
Args:
text: Text to split
metadata: Metadata to attach to each chunk
Returns:
List of chunk dictionaries
"""
if not text.strip():
return []
tokens = self.enc.encode(text)
chunks = []
start = 0
metadata = metadata or {}
while start < len(tokens):
end = min(start + self.chunk_size, len(tokens))
chunk_tokens = tokens[start:end]
chunk_text = self.enc.decode(chunk_tokens)
chunks.append({
"content": chunk_text,
"metadata": {
**metadata,
"chunk_index": len(chunks),
"start_token": start,
"end_token": end,
"total_tokens": len(tokens),
},
})
# Move forward, accounting for overlap
start = end - self.chunk_overlap if end < len(tokens) else end
return chunks
def split_documents(self, documents: List[Dict]) -> List[Dict]:
"""Split multiple documents into chunks."""
all_chunks = []
for doc in documents:
chunks = self.split(doc["content"], doc.get("metadata", {}))
all_chunks.extend(chunks)
logger.info(f"Split {len(documents)} documents into {len(all_chunks)} chunks")
return all_chunks
Step 4: ChromaDB Vector Store
# vectorstore.py
import chromadb
from chromadb.config import Settings
from openai import OpenAI
from typing import List, Dict, Optional
import logging
logger = logging.getLogger(__name__)
class VectorStore:
"""
ChromaDB vector store with OpenAI embeddings.
Provides storage, retrieval, and metadata filtering for
document embeddings.
"""
def __init__(
self,
collection_name: str = "documents",
persist_directory: str = "./chroma_db",
):
self.client = chromadb.Client(Settings(
chroma_db_impl="duckdb+parquet",
persist_directory=persist_directory,
anonymized_telemetry=False,
))
self.collection = self.client.get_or_create_collection(
name=collection_name,
metadata={"hnsw:space": "cosine"},
)
self.openai_client = OpenAI()
def _get_embedding(self, text: str) -> List[float]:
"""Get embedding for a single text."""
response = self.openai_client.embeddings.create(
model="text-embedding-3-small",
input=text,
)
return response.data[0].embedding
def _get_embeddings(self, texts: List[str]) -> List[List[float]]:
"""Get embeddings for multiple texts (batch)."""
response = self.openai_client.embeddings.create(
model="text-embedding-3-small",
input=texts,
)
return [item.embedding for item in response.data]
def add_documents(self, chunks: List[Dict]) -> None:
"""
Add document chunks to the vector store.
Args:
chunks: List of chunk dictionaries with content and metadata
"""
batch_size = 100
for i in range(0, len(chunks), batch_size):
batch = chunks[i:i + batch_size]
texts = [c["content"] for c in batch]
embeddings = self._get_embeddings(texts)
ids = [f"chunk_{i + j}" for j in range(len(batch))]
metadatas = [c.get("metadata", {}) for c in batch]
self.collection.add(
documents=texts,
embeddings=embeddings,
ids=ids,
metadatas=metadatas,
)
logger.info(f"Added {len(chunks)} chunks to vector store")
def search(
self,
query: str,
n_results: int = 5,
where: Optional[Dict] = None,
) -> List[Dict]:
"""
Search for similar documents.
Args:
query: Search query
n_results: Number of results to return
where: Metadata filter
Returns:
List of search results
"""
query_embedding = self._get_embedding(query)
kwargs = {
"query_embeddings": [query_embedding],
"n_results": n_results,
}
if where:
kwargs["where"] = where
results = self.collection.query(**kwargs)
output = []
for i in range(len(results["documents"][0])):
output.append({
"content": results["documents"][0][i],
"metadata": results["metadatas"][0][i] if results["metadatas"] else {},
"distance": results["distances"][0][i] if results["distances"] else 0,
"id": results["ids"][0][i],
})
return output
Step 5: RAG Agent Orchestrator
# rag_agent.py
from openai import OpenAI
from vectorstore import VectorStore
from splitter import TextSplitter
from loader import load_documents
from typing import List, Dict
import logging
logger = logging.getLogger(__name__)
RAG_SYSTEM_PROMPT = """You are a helpful assistant that answers questions based on the provided context.
Rules:
1. Only use information from the provided context
2. Cite sources using [Source N] notation
3. If the context doesn't contain enough information, say so
4. Be concise and accurate
5. Include relevant quotes from the source material
Context:
{context}"""
class RAGAgent:
"""
Complete RAG agent with ingestion, retrieval, and generation.
Features:
- Document ingestion with automatic chunking
- Semantic search with metadata filtering
- Source citation in generated answers
- Configurable retrieval parameters
"""
def __init__(self, collection_name: str = "documents"):
self.llm = OpenAI()
self.vectorstore = VectorStore(collection_name=collection_name)
self.splitter = TextSplitter(chunk_size=500, chunk_overlap=50)
def ingest(self, source: str) -> int:
"""
Ingest documents from a file or directory.
Args:
source: Path to file or directory
Returns:
Number of chunks ingested
"""
documents = list(load_documents(source))
chunks = self.splitter.split_documents(documents)
self.vectorstore.add_documents(chunks)
logger.info(f"Ingested {len(chunks)} chunks from {source}")
return len(chunks)
def query(
self,
question: str,
n_results: int = 5,
filters: dict = None,
) -> Dict:
"""
Answer a question using retrieved context.
Args:
question: User's question
n_results: Number of chunks to retrieve
filters: Metadata filters
Returns:
Dictionary with answer, sources, and metadata
"""
# Retrieve relevant chunks
results = self.vectorstore.search(
query=question, n_results=n_results, where=filters
)
# Build context with source citations
context_parts = []
for i, doc in enumerate(results):
source = doc["metadata"].get("source", "Unknown")
context_parts.append(
f"[Source {i+1}] ({source})\n{doc['content']}"
)
context = "\n\n".join(context_parts)
# Generate answer
response = self.llm.chat.completions.create(
model="gpt-4-turbo-preview",
messages=[
{"role": "system", "content": RAG_SYSTEM_PROMPT.format(context=context)},
{"role": "user", "content": question},
],
temperature=0.0,
)
answer = response.choices[0].message.content
return {
"answer": answer,
"sources": [
{
"content": doc["content"][:200],
"source": doc["metadata"].get("source", "Unknown"),
"score": 1 - doc["distance"],
}
for doc in results
],
"num_sources": len(results),
}
Step 6: Evaluation Framework
# eval.py
from typing import List, Dict
import logging
logger = logging.getLogger(__name__)
def precision_at_k(retrieved: List[str], relevant: List[str], k: int) -> float:
"""Calculate Precision@k."""
retrieved_k = retrieved[:k]
hits = sum(1 for doc in retrieved_k if doc in relevant)
return hits / k
def recall_at_k(retrieved: List[str], relevant: List[str], k: int) -> float:
"""Calculate Recall@k."""
retrieved_k = retrieved[:k]
hits = sum(1 for doc in retrieved_k if doc in relevant)
return hits / len(relevant) if relevant else 0.0
def mrr(retrieved: List[str], relevant: List[str]) -> float:
"""Calculate Mean Reciprocal Rank."""
for i, doc in enumerate(retrieved):
if doc in relevant:
return 1.0 / (i + 1)
return 0.0
def evaluate_rag(
agent: "RAGAgent",
test_data: List[Dict],
k: int = 5,
) -> Dict:
"""
Evaluate RAG agent on test data.
Args:
agent: RAG agent to evaluate
test_data: List of {"question": str, "relevant_sources": List[str]}
k: Number of results to consider
Returns:
Evaluation metrics
"""
metrics = {"precisions": [], "recalls": [], "mrrs": []}
for item in test_data:
result = agent.query(item["question"], n_results=k)
retrieved = [s["source"] for s in result["sources"]]
relevant = item.get("relevant_sources", [])
metrics["precisions"].append(precision_at_k(retrieved, relevant, k))
metrics["recalls"].append(recall_at_k(retrieved, relevant, k))
metrics["mrrs"].append(mrr(retrieved, relevant))
return {
"precision@k": sum(metrics["precisions"]) / len(metrics["precisions"]),
"recall@k": sum(metrics["recalls"]) / len(metrics["recalls"]),
"mrr": sum(metrics["mrrs"]) / len(metrics["mrrs"]),
"num_queries": len(test_data),
}
Mathematical Foundation
Cosine Similarity for retrieval:
Where each parameter means:
- â query embedding vector
- â document embedding vector
- â dot product of embeddings
- , â L2 norms of the vectors
Intuition: Cosine similarity measures the angle between vectors, ranging from -1 (opposite) to 1 (identical). Higher similarity means more semantically related content.
Mean Reciprocal Rank (MRR):
Intuition: MRR measures how early the first relevant result appears. An MRR of 1.0 means the first result is always relevant.
Chunk Overlap Coverage:
Intuition: Overlap ensures no information is lost at chunk boundaries. Typical overlap is 10-20% of chunk size.
Retrieval Precision:
Intuition: What fraction of the top-k retrieved documents are actually relevant? Higher precision means less noise in the context.
Performance Metrics
| Metric | Value | Notes |
|---|---|---|
| Retrieval Precision@5 | 0.85+ | With proper chunking |
| Retrieval Recall@5 | 0.90+ | With hybrid search |
| MRR | 0.75+ | Average across queries |
| Ingestion Speed | 1000 docs/min | text-embedding-3-small |
| Query Latency | 200-500ms | Excluding LLM call |
| Storage per 1K docs | ~50MB | With embeddings |
Real-World Examples
Example 1: Enterprise Knowledge Base
A company uses RAG to let employees query internal documentation:
agent = RAGAgent(collection_name="company_docs")
# Ingest all company documentation
agent.ingest("./docs/hr/")
agent.ingest("./docs/engineering/")
agent.ingest("./docs/sales/")
# Query with metadata filtering
result = agent.query(
"What is the vacation policy?",
filters={"type": ".md"} # Only search markdown docs
)
print(result["answer"])
# "According to [Source 1] (HR/vacation-policy.md), employees get..."
Example 2: Research Assistant
A researcher uses RAG to query academic papers:
agent = RAGAgent(collection_name="papers")
# Ingest PDF papers
agent.ingest("./papers/transformer_efficiency.pdf")
agent.ingest("./papers/attention_mechanisms.pdf")
# Query for specific information
result = agent.query(
"What are the main bottlenecks in transformer inference?",
n_results=10 # Retrieve more context for complex questions
)
Common Pitfalls & Solutions
| Pitfall | Solution |
|---|---|
| Poor chunk quality | Use semantic chunking, not fixed-size |
| Low retrieval recall | Increase chunk overlap, use hybrid search |
| Hallucination | Enforce citation requirements in prompts |
| Stale data | Implement incremental index updates |
| High latency | Cache embeddings, use batch processing |
| Context overflow | Limit retrieved chunks by token budget |
| Embedding drift | Use domain-specific fine-tuned embeddings |
| Noisy retrieval | Add reranking stage after initial retrieval |
Security Considerations
Critical security measures for production RAG systems:
- Access Control: Ensure users only retrieve documents they have permission to access
- Data Sanitization: Filter sensitive information (PII, secrets) from retrieved context
- Audit Logging: Track all queries and retrieved documents for compliance
- Embedding Security: Never expose raw embeddings to clients
- Prompt Injection: Validate queries to prevent injection attacks through retrieved content
- Content Filtering: Remove harmful or inappropriate content from generation
# Example: Access-controlled RAG query
def secure_query(agent, question, user_id):
# Get user's accessible document groups
user_groups = get_user_groups(user_id)
# Filter retrieval to only accessible documents
results = agent.query(
question,
filters={"group": {"$in": user_groups}}
)
return results
Summary with Key Takeaways
- RAG grounding in retrieved documents significantly reduces hallucination
- Chunking strategy critically impacts retrieval quality â use overlapping chunks
- Hybrid search (semantic + keyword) outperforms either method alone
- Always cite sources in generated answers for verifiability
- Evaluate with precision@k, recall@k, and MRR to measure retrieval quality
- Reranking improves precision by 10-20% over embedding search alone
- Incremental indexing keeps the knowledge base current without full re-indexing
Interview Questions
1. Why is chunking strategy critical for RAG performance?
Answer: Chunking determines what information is available for retrieval. Poor chunking splits documents at semantic boundaries, losing context. Key considerations: 1) Chunk size â Too small loses context, too large reduces precision, 2) Overlap â Prevents information loss at boundaries, 3) Semantic chunking â Split at paragraph/section boundaries rather than fixed token counts, 4) Metadata preservation â Track source, page, and section for citations. Optimal chunk size is typically 200-500 tokens with 10-20% overlap. Advanced approaches use sliding windows with semantic similarity to determine boundaries.
2. How does hybrid search improve retrieval quality?
Answer: Hybrid search combines semantic (vector) search with keyword (BM25) search. Semantic search excels at conceptual matching but may miss exact terms. Keyword search catches exact matches but misses synonyms. Score fusion combines both: where is typically 0.6-0.8. This achieves higher recall than either method alone, especially for queries containing specific terms (product names, technical jargon) that need exact matching.
3. What is the role of reranking in RAG?
Answer: Initial retrieval (embedding search) is fast but approximate. Reranking uses a cross-encoder model that jointly processes query and document, producing more accurate relevance scores. The pipeline is: 1) Retrieve top-k candidates with embedding search (k=20-50), 2) Rerank with cross-encoder, 3) Return top-n results (n=3-5). Cross-encoders are slower but more accurate because they attend to query-document interactions. This two-stage approach balances speed and quality.
4. How do you handle document updates in RAG?
Answer: Incremental indexing strategies: 1) Document hashing â Compute hash of each document, only re-index changed documents, 2) Chunk-level diffing â Compare chunks to detect additions/deletions, 3) Metadata versioning â Track document versions in metadata, 4) TTL expiration â Set time-to-live for chunks and re-index periodically. ChromaDB supports upsert operations for updating existing embeddings. For large-scale systems, use a separate indexing pipeline that runs asynchronously.
5. What are the limitations of RAG?
Answer: Key limitations: 1) Retrieval quality â If relevant documents aren't retrieved, generation fails, 2) Context window limits â Can only include limited retrieved content, 3) Latency â Retrieval adds 200-500ms to response time, 4) Embedding drift â Embedding models may not capture domain-specific semantics, 5) Hallucination â LLM may still hallucinate despite retrieved context, 6) Cost â Embedding storage and retrieval add infrastructure costs. Mitigation: hybrid search, reranking, few-shot examples, and structured output formats.
6. How do you evaluate RAG system quality?
Answer: Multi-level evaluation: 1) Retrieval metrics â Precision@k, Recall@k, MRR, NDCG, 2) Generation metrics â Faithfulness (answer grounded in context), Relevance (answer addresses question), 3) End-to-end metrics â Human evaluation, A/B testing, 4) System metrics â Latency, throughput, cost. Use frameworks like RAGAS or DeepEval for automated evaluation. Create golden test sets with known relevant documents and correct answers. Monitor retrieval quality degradation over time as documents are added.
7. What is the difference between dense and sparse retrieval?
Answer: Dense retrieval uses neural embeddings (like text-embedding-3) to represent documents as vectors, capturing semantic meaning. Sparse retrieval uses term-based methods like BM25 that match exact keywords. Dense retrieval handles synonyms and conceptual queries better but may miss exact matches. Sparse retrieval is faster and more interpretable but requires exact term overlap. Hybrid approaches combine both for best results. Dense retrieval is preferred for semantic search; sparse for code search or product matching where exact terms matter.
8. How would you scale RAG to millions of documents?
Answer: Scaling strategies: 1) Distributed vector DB â Use Pinecone, Weaviate, or Qdrant with sharding, 2) Hierarchical indexing â Cluster documents and search within relevant clusters first, 3) Caching â Cache frequent queries and embeddings, 4) Async pipelines â Separate indexing from query serving, 5) GPU acceleration â Use GPU-optimized HNSW indexes, 6) Compression â Quantize embeddings to reduce storage (1536â512 dims), 7) Streaming ingestion â Process documents as they arrive rather than batch. Monitor query latency percentiles and auto-scale based on load.
KnowledgeCheck
-
What is the primary advantage of RAG over fine-tuning?
- a) RAG is faster to train
- b) RAG can be updated without retraining the model
- c) RAG uses fewer tokens
- d) RAG doesn't require an LLM
-
What is the purpose of chunk overlap in text splitting?
- a) To increase storage usage
- b) To prevent information loss at chunk boundaries
- c) To make chunks smaller
- d) To reduce embedding computation
-
What metric measures how early the first relevant result appears?
- a) Precision@k
- b) Recall@k
- c) Mean Reciprocal Rank (MRR)
- d) F1 score
-
Why is hybrid search better than semantic search alone?
- a) It's faster
- b) It handles both conceptual and exact matching
- c) It uses fewer resources
- d) It doesn't need embeddings
-
What is the role of a reranker in RAG?
- a) To generate embeddings
- b) To improve relevance scoring of retrieved documents
- c) To store documents
- d) To generate the final answer
-
What is the typical optimal chunk size for RAG?
- a) 50-100 tokens
- b) 200-500 tokens
- c) 1000-2000 tokens
- d) 5000+ tokens
Answers: 1-b, 2-b, 3-c, 4-b, 5-b, 6-b