A clean, modern Retrieval-Augmented Generation (RAG) pipeline designed to be read, forked, and extended. This repository demonstrates a production-style RAG system with clear, readable code optimized for learning and interviews.
- Document Loading: Support for text files and directories
- Smart Chunking: Configurable text chunking with overlap
- Flexible Embeddings: OpenAI, HuggingFace, or local embedding models
- Vector Store: FAISS-based efficient similarity search
- Retrieval: Similarity and MMR (Maximum Marginal Relevance) search
- RAG Chain: Complete question-answering pipeline
- Optional Reranking: Improve retrieval quality with cross-encoders
- Streaming: Real-time streaming of LLM responses
- Tests: Comprehensive unit tests with pytest
demo-rag/
├── src/ # Core RAG components
│ ├── loader.py # Document loading
│ ├── chunker.py # Text chunking
│ ├── embedder.py # Embedding generation
│ ├── vectorstore.py # FAISS vector store
│ ├── retriever.py # Document retrieval
│ ├── rag_chain.py # RAG pipeline
│ └── config.py # Configuration
├── examples/
│ └── demo.py # Main demo script
├── tests/ # Unit tests
├── docs/ # Documentation
├── data/
│ └── sample_docs/ # Sample documents
├── requirements.txt
├── README.md
├── Dockerfile
└── LICENSE
Option 1: Install from source
git clone <your-repo-url>
cd demo-rag
pip install -r requirements.txtOption 2: Install as package (development)
git clone <your-repo-url>
cd demo-rag
pip install -e ".[dev]"Option 3: Using Make
make installSet up your API key (optional - falls back to HuggingFace if not set):
export OPENAI_API_KEY="your-api-key-here"python examples/demo.pyThe demo will:
- Load documents from
data/sample_docs/ - Chunk the documents
- Generate embeddings
- Build a vector store
- Allow you to ask questions interactively
The RAG pipeline is highly configurable. See src/config.py for all options:
from src.config import RAGConfig, ChunkConfig, EmbeddingConfig, RetrievalConfig
config = RAGConfig(
chunk_config=ChunkConfig(
chunk_size=512,
chunk_overlap=50
),
embedding_config=EmbeddingConfig(
provider="openai", # or "huggingface" or "local"
model_name="text-embedding-3-small"
),
retrieval_config=RetrievalConfig(
top_k=5,
search_type="similarity" # or "mmr"
),
llm_model="gpt-4o-mini",
enable_reranking=False
)The RAG pipeline follows this flow:
- Load: Documents are loaded from files or directories
- Chunk: Documents are split into smaller chunks for embedding
- Embed: Chunks are converted to vector embeddings
- Store: Embeddings are stored in a FAISS vector index
- Retrieve: Query embeddings are used to find similar chunks
- Rerank (optional): Retrieved chunks are reranked for better relevance
- Generate: LLM generates an answer using retrieved context
See docs/architecture.md for detailed architecture documentation.
Run tests with pytest:
pytest tests/With coverage:
pytest tests/ --cov=src --cov-report=htmlExample queries to test the system:
- "What is this document about?"
- "Summarize the main topics."
- "What are the key points?"
- "Explain the concepts in detail."
This repository provides a solid foundation. Here are ideas for extending it:
Combine semantic search with keyword-based search (BM25) for better retrieval:
# Add BM25 retriever
from rank_bm25 import BM25Okapi
class HybridRetriever:
def __init__(self, vector_retriever, bm25_retriever):
self.vector_retriever = vector_retriever
self.bm25_retriever = bm25_retriever
def retrieve(self, query, k=5):
vector_results = self.vector_retriever.retrieve(query, k=k*2)
bm25_results = self.bm25_retriever.retrieve(query, k=k*2)
# Combine and rerank
return combined_resultsImprove retrieval quality with cross-encoder reranking:
config = RAGConfig(
enable_reranking=True,
reranker_model="cross-encoder/ms-marco-MiniLM-L-6-v2"
)Generate multiple query variations and combine results:
def multi_query_retrieve(self, query, k=5):
# Generate query variations
variations = self.generate_query_variations(query)
# Retrieve for each variation
all_results = []
for variation in variations:
results = self.retriever.retrieve(variation, k=k)
all_results.extend(results)
# Deduplicate and rerank
return deduplicate_and_rerank(all_results, k)Handle long documents with hierarchical chunking:
class HierarchicalChunker:
def chunk(self, document):
# Create parent chunks
parent_chunks = self.chunk_by_section(document)
# Create child chunks within each parent
all_chunks = []
for parent in parent_chunks:
children = self.chunk_by_paragraph(parent)
all_chunks.extend(children)
return all_chunksAdd agent capabilities for complex reasoning:
class AgenticRAG:
def __init__(self, rag_chain, tools):
self.rag_chain = rag_chain
self.tools = tools
def answer(self, query):
# Decide if RAG is sufficient or if tools are needed
if self.needs_tools(query):
return self.use_tools(query)
else:
return self.rag_chain.answer(query)- Metadata Filtering: Filter by document metadata (date, author, etc.)
- Query Expansion: Expand queries with synonyms and related terms
- Citation Tracking: Track which chunks contributed to the answer
- Confidence Scores: Provide confidence scores for answers
- Multi-modal: Support images, tables, and other document types
- Caching: Cache embeddings and retrieval results
- Async Processing: Process documents asynchronously for better performance
Build and run with Docker:
docker build -t demo-rag .
docker run -e OPENAI_API_KEY=your-key demo-ragSee LICENSE file for details.
Contributions are welcome! Please see CONTRIBUTING.md for guidelines.
This is a demo repository designed for learning. Feel free to fork and extend it!
make test # Run tests
make lint # Check code style
make format # Format code
make run-demo # Run demo scriptBuilt with ❤️ for learning and exploration.