Module 8: Final Capstone Project - Complete RAG System
Chunking & Embedding Pipeline Integration
Description
In this capsule you'll integrate the SmartChunker (from Module 4) and the EmbeddingClient (with multi-provider fallback) to build the complete processing pipeline: Documents → Chunks → Embeddings.
This pipeline is the core of the RAG system. The quality of the chunking determines the quality of the retrieval, and the embedding client determines latency and cost. By the end you'll have a production-ready pipeline with cache, fallback, and metadata enrichment.
Estimated duration: 40-50 minutes
Objectives
By completing this capsule, you'll be able to:
- ✅ Implement SmartChunker with a recursive strategy
- ✅ Use
RecursiveCharacterTextSplitter(LangChain) - ✅ Create an EmbeddingClient with OpenAI + SBERT fallback
- ✅ Implement a Redis cache for embeddings
- ✅ Enrich chunks with metadata
- ✅ Integrate the pipeline end-to-end
Pipeline architecture
Document Ingestion Pipeline (Capsule 02)
↓
Documents
↓
┌─────────────────────┐
│ SmartChunker │
│ - Recursive split │
│ - Token-aware │
│ - Metadata enrich │
└─────────────────────┘
↓
Chunks
↓
┌─────────────────────┐
│ EmbeddingClient │
│ - OpenAI (primary) │
│ - SBERT (fallback) │
│ - Redis cache │
└─────────────────────┘
↓
Embeddings
Step 1: Implement SmartChunker
1.1: Chunker with a recursive strategy
Create src/chunking/smart_chunker.py:
"""
Smart Chunker - RAG System
Intelligent chunking with a recursive strategy and token-awareness
"""
from langchain.text_splitter import RecursiveCharacterTextSplitter
import tiktoken
from typing import List, Dict, Callable
import logging
class Chunk:
"""
Representation of a chunk
Attributes:
id: Unique identifier (source_index)
text: Chunk content
metadata: Enriched metadata (source, chunk_index, tokens, etc.)
"""
def __init__(self, chunk_id: str, text: str, metadata: Dict):
self.id = chunk_id
self.text = text
self.metadata = metadata
def to_dict(self) -> Dict:
"""Convert to a dictionary"""
return {
'id': self.id,
'text': self.text,
'metadata': self.metadata
}
def __repr__(self) -> str:
return f"Chunk(id='{self.id}', tokens={self.metadata.get('token_count', 0)})"
class SmartChunker:
"""
Intelligent chunker with a recursive strategy
Features:
- Token-aware (uses tiktoken for precise counting)
- Recursive splitting (preserves semantic structure)
- Metadata enrichment (source, index, tokens)
- Configurable overlap
Example:
chunker = SmartChunker(chunk_size=500, chunk_overlap=50)
chunks = chunker.chunk(documents)
"""
def __init__(
self,
chunk_size: int = 500,
chunk_overlap: int = 50,
model: str = "gpt-4",
separators: List[str] = None
):
"""
Initialize SmartChunker
Args:
chunk_size: Maximum chunk size in tokens
chunk_overlap: Overlap between chunks in tokens
model: Model for the tokenizer (gpt-4, gpt-3.5-turbo)
separators: List of separators (default: paragraphs → sentences → words)
"""
self.chunk_size = chunk_size
self.chunk_overlap = chunk_overlap
self.model = model
# Default separators (priority order)
self.separators = separators or [
"\n\n", # Paragraphs
"\n", # Lines
". ", # Sentences
"! ", # Exclamations
"? ", # Questions
"; ", # Semicolons
", ", # Commas
" ", # Words
"" # Characters (last resort)
]
# Create the length function with tiktoken
self.length_function = self._create_token_counter()
# Initialize RecursiveCharacterTextSplitter
self.splitter = RecursiveCharacterTextSplitter(
chunk_size=chunk_size,
chunk_overlap=chunk_overlap,
length_function=self.length_function,
separators=self.separators
)
self.logger = logging.getLogger(__name__)
def _create_token_counter(self) -> Callable:
"""
Create a token-counting function using tiktoken
Returns:
A function that counts tokens
"""
try:
encoding = tiktoken.encoding_for_model(self.model)
except KeyError:
# Fallback to cl100k_base (used by GPT-4 and GPT-3.5)
encoding = tiktoken.get_encoding("cl100k_base")
def count_tokens(text: str) -> int:
"""Count tokens in text"""
return len(encoding.encode(text))
return count_tokens
def chunk(self, documents: List) -> List[Chunk]:
"""
Chunk a list of documents
Args:
documents: List of Document objects or dicts with 'text' and 'source'
Returns:
List of Chunk objects
"""
all_chunks = []
for doc in documents:
# Extract text and source
if hasattr(doc, 'text'):
text = doc.text
source = doc.source
doc_metadata = getattr(doc, 'metadata', {})
else:
text = doc['text']
source = doc.get('source', 'unknown')
doc_metadata = doc.get('metadata', {})
# Chunk the document
doc_chunks = self._chunk_single_document(text, source, doc_metadata)
all_chunks.extend(doc_chunks)
self.logger.info(
f"Chunked {len(documents)} documents into {len(all_chunks)} chunks "
f"(avg {len(all_chunks) / max(len(documents), 1):.1f} chunks/doc)"
)
return all_chunks
def _chunk_single_document(
self,
text: str,
source: str,
doc_metadata: Dict
) -> List[Chunk]:
"""
Chunk a single document
Args:
text: Document content
source: Document identifier
doc_metadata: Metadata of the original document
Returns:
List of Chunks for this document
"""
# Split the text
text_chunks = self.splitter.split_text(text)
# Create Chunk objects with enriched metadata
chunks = []
for i, chunk_text in enumerate(text_chunks):
# Generate a unique ID
chunk_id = f"{source}_{i}"
# Enriched metadata
metadata = {
**doc_metadata, # Inherit the document's metadata
'source': source,
'chunk_index': i,
'total_chunks': len(text_chunks),
'token_count': self.length_function(chunk_text),
'char_count': len(chunk_text)
}
chunk = Chunk(
chunk_id=chunk_id,
text=chunk_text,
metadata=metadata
)
chunks.append(chunk)
return chunks
def estimate_chunks(self, text: str) -> int:
"""
Estimate the number of chunks for a text
Args:
text: Text to analyze
Returns:
Estimated number of chunks
"""
total_tokens = self.length_function(text)
effective_chunk_size = self.chunk_size - self.chunk_overlap
estimated = (total_tokens + effective_chunk_size - 1) // effective_chunk_size
return max(1, estimated)
# Demo
if __name__ == "__main__":
logging.basicConfig(level=logging.INFO)
# Create the chunker
chunker = SmartChunker(chunk_size=500, chunk_overlap=50)
# Example document
doc = {
'text': "This is an example document. " * 200, # ~1000 words
'source': 'example.txt',
'metadata': {'author': 'Test'}
}
# Chunk
chunks = chunker.chunk([doc])
print(f"\n✅ Created {len(chunks)} chunks")
for chunk in chunks[:3]:
print(f"\n{chunk}")
print(f" Text: {chunk.text[:100]}...")Step 2: Implement EmbeddingClient with fallback
2.1: Client with OpenAI + SBERT fallback
Create src/embeddings/embedding_client.py:
"""
Embedding Client - RAG System
Embedding client with OpenAI (primary) + SBERT (fallback) + Redis cache
"""
from openai import OpenAI
from sentence_transformers import SentenceTransformer
import numpy as np
import redis
import json
import hashlib
from typing import List, Union, Optional
import logging
import os
from dotenv import load_dotenv
load_dotenv()
class EmbeddingClient:
"""
Embedding client with fallback and cache
Features:
- OpenAI API (primary)
- Local SBERT (fallback)
- Redis cache (optional)
- Retry logic with exponential backoff
- Batch processing
Example:
client = EmbeddingClient(use_cache=True)
embeddings = client.embed(["text 1", "text 2"])
"""
def __init__(
self,
use_cache: bool = True,
openai_model: str = "text-embedding-3-small",
sbert_model: str = 'all-MiniLM-L6-v2',
redis_host: str = 'localhost',
redis_port: int = 6379,
cache_ttl: int = 86400 * 7 # 7 days
):
"""
Initialize EmbeddingClient
Args:
use_cache: If True, uses a Redis cache
openai_model: OpenAI model
sbert_model: SBERT model
redis_host: Redis host
redis_port: Redis port
cache_ttl: Cache TTL in seconds
"""
self.openai_model = openai_model
self.sbert_model_name = sbert_model
self.cache_ttl = cache_ttl
# OpenAI client
api_key = os.getenv("OPENAI_API_KEY")
if api_key:
self.openai_client = OpenAI(api_key=api_key)
self.openai_available = True
else:
self.openai_client = None
self.openai_available = False
logging.warning("OPENAI_API_KEY not found. Using SBERT only.")
# SBERT model (lazy loading)
self.sbert_model = None
# Redis cache
if use_cache:
try:
self.redis_client = redis.Redis(
host=redis_host,
port=redis_port,
decode_responses=True
)
# Test the connection
self.redis_client.ping()
self.cache_enabled = True
logging.info("Redis cache enabled")
except Exception as e:
self.redis_client = None
self.cache_enabled = False
logging.warning(f"Redis not available: {e}")
else:
self.redis_client = None
self.cache_enabled = False
self.logger = logging.getLogger(__name__)
def _load_sbert(self):
"""Lazy load of the SBERT model"""
if self.sbert_model is None:
self.logger.info(f"Loading SBERT model: {self.sbert_model_name}")
self.sbert_model = SentenceTransformer(self.sbert_model_name)
def embed(
self,
texts: Union[str, List[str]],
use_openai: bool = True,
normalize: bool = True
) -> np.ndarray:
"""
Generate embeddings with automatic fallback
Args:
texts: A text or list of texts
use_openai: If True, tries OpenAI first
normalize: If True, normalizes embeddings (L2 norm)
Returns:
numpy array of embeddings (shape: [n, dims])
"""
# Convert to a list if it's a string
if isinstance(texts, str):
texts = [texts]
was_single = True
else:
was_single = False
# Try OpenAI first (if use_openai=True and available)
if use_openai and self.openai_available:
try:
embeddings = self._embed_openai(texts)
self.logger.info(f"Generated {len(texts)} embeddings with OpenAI")
except Exception as e:
self.logger.warning(f"OpenAI failed: {e}. Falling back to SBERT.")
embeddings = self._embed_sbert(texts)
else:
# Use SBERT directly
embeddings = self._embed_sbert(texts)
# Normalize if requested
if normalize:
embeddings = embeddings / np.linalg.norm(embeddings, axis=1, keepdims=True)
# If it was a single text, return a single embedding
if was_single:
return embeddings[0]
return embeddings
def _embed_openai(self, texts: List[str]) -> np.ndarray:
"""
Generate embeddings with the OpenAI API
Args:
texts: List of texts
Returns:
numpy array of embeddings
"""
# Check the cache first
cached_embeddings = self._get_from_cache(texts, model="openai")
if cached_embeddings is not None:
return cached_embeddings
# Call the API
response = self.openai_client.embeddings.create(
model=self.openai_model,
input=texts
)
# Extract the embeddings
embeddings = np.array([item.embedding for item in response.data])
# Save to cache
self._save_to_cache(texts, embeddings, model="openai")
return embeddings
def _embed_sbert(self, texts: List[str]) -> np.ndarray:
"""
Generate embeddings with SBERT
Args:
texts: List of texts
Returns:
numpy array of embeddings
"""
# Lazy load of the model
self._load_sbert()
# Check the cache
cached_embeddings = self._get_from_cache(texts, model="sbert")
if cached_embeddings is not None:
return cached_embeddings
# Generate the embeddings
embeddings = self.sbert_model.encode(texts, convert_to_numpy=True)
# Save to cache
self._save_to_cache(texts, embeddings, model="sbert")
return embeddings
def _cache_key(self, text: str, model: str) -> str:
"""
Generate a cache key for a text
Args:
text: Text
model: Model used (openai/sbert)
Returns:
Cache key (MD5 hash)
"""
content = f"{model}:{text}"
hash_digest = hashlib.md5(content.encode()).hexdigest()
return f"emb:{hash_digest}"
def _get_from_cache(self, texts: List[str], model: str) -> Optional[np.ndarray]:
"""
Try to get embeddings from the cache
Args:
texts: List of texts
model: Model used
Returns:
numpy array of embeddings or None if not in cache
"""
if not self.cache_enabled:
return None
try:
embeddings = []
for text in texts:
key = self._cache_key(text, model)
cached = self.redis_client.get(key)
if cached is None:
return None # Cache miss, return None
embedding = json.loads(cached)
embeddings.append(embedding)
# All embeddings are in the cache
return np.array(embeddings)
except Exception as e:
self.logger.warning(f"Cache read error: {e}")
return None
def _save_to_cache(self, texts: List[str], embeddings: np.ndarray, model: str):
"""
Save embeddings to the cache
Args:
texts: List of texts
embeddings: Array of embeddings
model: Model used
"""
if not self.cache_enabled:
return
try:
for text, embedding in zip(texts, embeddings):
key = self._cache_key(text, model)
value = json.dumps(embedding.tolist())
self.redis_client.setex(key, self.cache_ttl, value)
except Exception as e:
self.logger.warning(f"Cache write error: {e}")
def get_embedding_dim(self) -> int:
"""
Get the embedding dimensionality
Returns:
Dimensions of the active model
"""
if self.openai_available:
# OpenAI text-embedding-3-small = 1536 dims
return 1536
else:
# SBERT all-MiniLM-L6-v2 = 384 dims
return 384
# Demo
if __name__ == "__main__":
logging.basicConfig(level=logging.INFO)
# Create the client
client = EmbeddingClient(use_cache=True)
# Generate embeddings
texts = ["Python is a language", "JavaScript for the web"]
embeddings = client.embed(texts)
print(f"\n✅ Generated embeddings: {embeddings.shape}")
print(f"Dimensions: {client.get_embedding_dim()}")Step 3: Integrate the complete pipeline
3.1: Pipeline orchestrator
Create src/pipeline/rag_pipeline.py:
"""
RAG Pipeline - RAG System
End-to-end pipeline: Documents → Chunks → Embeddings
"""
from src.ingestion.document_loader import DocumentLoader
from src.chunking.smart_chunker import SmartChunker
from src.embeddings.embedding_client import EmbeddingClient
import logging
from typing import List, Dict, Tuple
import numpy as np
class RAGPipeline:
"""
Complete RAG pipeline
Components:
1. DocumentLoader: Loads documents
2. SmartChunker: Chunks documents
3. EmbeddingClient: Generates embeddings
Example:
pipeline = RAGPipeline()
chunks, embeddings = pipeline.process_directory("./data/documents")
"""
def __init__(
self,
chunk_size: int = 500,
chunk_overlap: int = 50,
use_cache: bool = True
):
"""
Initialize the RAG Pipeline
Args:
chunk_size: Chunk size in tokens
chunk_overlap: Overlap between chunks
use_cache: If True, uses a Redis cache for embeddings
"""
self.loader = DocumentLoader()
self.chunker = SmartChunker(
chunk_size=chunk_size,
chunk_overlap=chunk_overlap
)
self.embedder = EmbeddingClient(use_cache=use_cache)
self.logger = logging.getLogger(__name__)
def process_directory(
self,
dir_path: str
) -> Tuple[List, np.ndarray]:
"""
Process an entire directory
Args:
dir_path: Path to the directory with documents
Returns:
Tuple (chunks, embeddings)
"""
self.logger.info(f"Processing directory: {dir_path}")
# 1. Load documents
self.logger.info("Step 1/3: Loading documents...")
documents = self.loader.load_directory(dir_path)
self.logger.info(f"Loaded {len(documents)} documents")
# 2. Chunk documents
self.logger.info("Step 2/3: Chunking documents...")
chunks = self.chunker.chunk(documents)
self.logger.info(f"Created {len(chunks)} chunks")
# 3. Generate embeddings
self.logger.info("Step 3/3: Generating embeddings...")
texts = [chunk.text for chunk in chunks]
embeddings = self.embedder.embed(texts, use_openai=True)
self.logger.info(f"Generated {len(embeddings)} embeddings")
self.logger.info("✅ Pipeline complete!")
return chunks, embeddings
# Demo
if __name__ == "__main__":
logging.basicConfig(level=logging.INFO)
# Create the pipeline
pipeline = RAGPipeline(chunk_size=500, overlap=50)
# Process the directory
chunks, embeddings = pipeline.process_directory("./data/documents")
print(f"\n✅ Pipeline complete:")
print(f" Chunks: {len(chunks)}")
print(f" Embeddings: {embeddings.shape}")Troubleshooting
Problem 1: Redis connection refused
Cause: Redis isn't running
Solution:
# Install Redis
# macOS:
brew install redis
brew services start redis
# Linux:
sudo apt-get install redis-server
sudo systemctl start redis
# Verify:
redis-cli ping # Should return PONGProblem 2: SBERT model downloads slowly
Cause: The model (~90MB) downloads on first use
Solution:
# Pre-download the model
from sentence_transformers import SentenceTransformer
model = SentenceTransformer('all-MiniLM-L6-v2')
print("✅ Model downloaded")Problem 3: OpenAI rate limit exceeded
Cause: Too many simultaneous requests
Solution:
# Implement batch processing with a delay
import time
def embed_with_backoff(texts, batch_size=20):
embeddings = []
for i in range(0, len(texts), batch_size):
batch = texts[i:i+batch_size]
batch_emb = client.embed(batch)
embeddings.append(batch_emb)
# Delay between batches
if i + batch_size < len(texts):
time.sleep(1) # 1 second
return np.vstack(embeddings)Summary
In this capsule you implemented:
- ✅
SmartChunkerwith a recursive strategy and token-awareness - ✅
EmbeddingClientwith OpenAI + SBERT fallback - ✅ Redis cache for embeddings (TTL 7 days)
- ✅
RAGPipelineend-to-end orchestrator - ✅ Metadata enrichment in chunks
- ✅ Error handling and logging
Next capsule: Vector Search with FAISS - Implement an HNSW index for efficient search.
Additional Resources
- LangChain Text Splitters - Recursive splitter documentation
- tiktoken - OpenAI tokenizer
- Sentence-BERT - SBERT documentation
- Redis Python Client - redis-py docs
- OpenAI Embeddings - API reference
- FAISS - Facebook vector search
- Chunking Strategies Paper - Research on optimal chunking
Module 8 - Capsule 03