Module 8: Capstone RAG Project with ChromaDB
Capsule 07: Final Hardening
Capsule description
Before closing the project, you'll apply a hardening pass: resilience to failures, basic security, and maintainability. A RAG system "works" when retrieval and generation respond; a "production-ready" system keeps working when ChromaDB goes down, the LLM times out, or an attacker sends 10,000 requests per second.
In this capsule you'll implement:
- Error handling with retry logic, fallback responses, and controlled degradation
- Caching (embedding cache + result cache) with hit rate metrics
- Security (API keys, rate limiting, input validation)
- Documentation (README with a diagram, API reference, deployment guide)
At the end you'll have a system that defends itself against failures and abuse, and documentation that lets any developer operate it.
Why Hardening Matters
Real failure scenarios
Day 1: All good
Day 2: ChromaDB restarts due to a deploy → /ask returns 500 for 30s
Day 3: OpenAI has a p99 latency of 15s → cascading timeouts
Day 4: A poorly written script makes 500 req/sec → service down for everyone
Day 5: Someone discovers there is no auth → extracts the entire corpus
Without hardening, each of these is an incident. With hardening: retries, fallbacks, rate limit, and auth reduce the impact.
Hardening Areas
| Area | What to do | Priority |
|---|---|---|
| Reliability | Retries, timeouts, fallbacks when an external service fails | High |
| Security | API key, rate limiting, input validation | High |
| Performance | Caching of embeddings and results | Medium |
| Errors | Uniform 4xx/5xx format, don't expose stack traces | High |
| Documentation | README, API reference, operations runbook | Medium |
Robust Error Handling
Retry Logic for ChromaDB and OpenAI
# app/utils/retry.py
import asyncio
import functools
from tenacity import retry, stop_after_attempt, wait_exponential, retry_if_exception_type
def with_retry(max_attempts=3, exceptions=(ConnectionError, TimeoutError)):
def decorator(func):
@functools.wraps(func)
async def wrapper(*args, **kwargs):
last_exc = None
for attempt in range(max_attempts):
try:
return await func(*args, **kwargs)
except exceptions as e:
last_exc = e
if attempt < max_attempts - 1:
await asyncio.sleep(2 ** attempt) # 1s, 2s, 4s
raise last_exc
return wrapper
return decorator
# Usage
@with_retry(max_attempts=3)
async def get_chroma_results(collection, query_embedding, top_k=5):
return collection.query(query_embeddings=[query_embedding], n_results=top_k)
With tenacity (a more complete option):
from tenacity import retry, stop_after_attempt, wait_exponential
@retry(
stop=stop_after_attempt(3),
wait=wait_exponential(multiplier=1, min=1, max=10),
retry=retry_if_exception_type((ConnectionError, TimeoutError))
)
async def call_openai(messages):
# ...
Fallback Responses
When retrieval or generation fail, return a controlled response instead of a generic 500.
# app/api/ask.py
FALLBACK_MESSAGES = {
"retrieval_failed": "I couldn't search the knowledge base. Please try again in a few seconds.",
"generation_failed": "I found relevant information but had a problem generating the answer. Try rephrasing your question.",
"no_evidence": "I didn't find enough information to answer your question with confidence. Can you be more specific?",
}
async def ask_endpoint(payload: dict):
question = payload.get("question", "").strip()[:500] # limit
if not question:
raise HTTPException(400, "question is required")
try:
docs = await retrieve(question, top_k=5)
except Exception as e:
logger.warning(f"Retrieval failed: {e}")
return {
"answer": FALLBACK_MESSAGES["retrieval_failed"],
"sources": [],
"confidence": 0,
"fallback": True,
}
if not docs or not docs.get("documents") or not docs["documents"][0]:
return {
"answer": FALLBACK_MESSAGES["no_evidence"],
"sources": [],
"confidence": 0,
}
try:
answer = await generate_answer(question, docs["documents"][0])
except Exception as e:
logger.warning(f"Generation failed: {e}")
return {
"answer": FALLBACK_MESSAGES["generation_failed"],
"sources": [m for m in docs.get("metadatas", [[]])[0]],
"confidence": 0.5,
"fallback": True,
}
return {"answer": answer, "sources": docs["metadatas"][0], "confidence": 0.85}
Graceful Degradation
If ChromaDB is down but you have a result cache, you can keep answering frequent questions:
async def ask_with_degradation(question: str):
# 1. Try the result cache first
cached = await result_cache.get(question)
if cached:
return {**cached, "from_cache": True}
# 2. Try normal retrieval
try:
docs = await retrieve(question)
except Exception:
# 3. Degrade: answer with the model only (no RAG) or an unavailability message
return {"answer": "The search service is not available. Try again later.", "sources": []}
# ... normal flow
Caching
Embedding Cache
Avoids recomputing embeddings for the same questions or texts.
# app/cache/embedding_cache.py
import hashlib
import json
from typing import Optional
# Using Redis or an in-memory dict for development
class EmbeddingCache:
def __init__(self, redis_url: Optional[str] = None):
self._redis = redis.from_url(redis_url) if redis_url else {}
self._local = {} if not redis_url else None # fallback in-memory
def _key(self, text: str, model: str) -> str:
h = hashlib.sha256(f"{model}:{text}".encode()).hexdigest()
return f"emb:{h}"
async def get(self, text: str, model: str = "text-embedding-3-small") -> Optional[list]:
k = self._key(text, model)
if self._redis:
val = await self._redis.get(k)
return json.loads(val) if val else None
return self._local.get(k)
async def set(self, text: str, embedding: list, model: str = "text-embedding-3-small", ttl: int = 86400):
k = self._key(text, model)
val = json.dumps(embedding)
if self._redis:
await self._redis.setex(k, ttl, val)
else:
self._local[k] = embedding
Result Cache
Caches complete /ask responses for identical questions.
# app/cache/result_cache.py
class ResultCache:
def __init__(self, redis_url: Optional[str] = None, ttl: int = 3600):
self._redis = redis.from_url(redis_url) if redis_url else {}
self._ttl = ttl
def _key(self, question: str) -> str:
return f"ask:{hashlib.sha256(question.strip().lower().encode()).hexdigest()}"
async def get(self, question: str) -> Optional[dict]:
k = self._key(question)
if self._redis:
val = await self._redis.get(k)
return json.loads(val) if val else None
return None
async def set(self, question: str, result: dict):
k = self._key(question)
if self._redis:
await self._redis.setex(k, self._ttl, json.dumps(result))
Hit Rate Metrics
# app/metrics.py
from prometheus_client import Counter
cache_hits = Counter("rag_cache_hits_total", "Cache hits", ["cache_type"]) # embedding, result
cache_misses = Counter("rag_cache_misses_total", "Cache misses", ["cache_type"])
# In the flow
if cached_emb := await embedding_cache.get(text):
cache_hits.labels(cache_type="embedding").inc()
return cached_emb
cache_misses.labels(cache_type="embedding").inc()
emb = await get_embedding(text)
await embedding_cache.set(text, emb)
Security
API Key in Headers
# app/auth.py
from fastapi import Security, HTTPException
from fastapi.security import APIKeyHeader
api_key_header = APIKeyHeader(name="X-API-Key", auto_error=False)
async def verify_api_key(api_key: str = Security(api_key_header)):
expected = os.getenv("API_KEY")
if not expected:
return None # No API key configured, allow (dev only)
if api_key != expected:
raise HTTPException(403, "Invalid API key")
return api_key
@app.post("/ask", dependencies=[Depends(verify_api_key)])
async def ask(payload: dict):
...
Rate Limiting
# app/middleware/rate_limit.py
from slowapi import Limiter
from slowapi.util import get_remote_address
limiter = Limiter(key_func=get_remote_address)
app.state.limiter = limiter
@app.post("/ask")
@limiter.limit("20/minute")
async def ask(request: Request, payload: dict):
...
Simple alternative without external dependencies:
from collections import defaultdict
from time import time
class SimpleRateLimiter:
def __init__(self, requests_per_minute=60):
self.rpm = requests_per_minute
self.requests = defaultdict(list)
def is_allowed(self, client_id: str) -> bool:
now = time()
self.requests[client_id] = [t for t in self.requests[client_id] if now - t < 60]
if len(self.requests[client_id]) >= self.rpm:
return False
self.requests[client_id].append(now)
return True
Input Validation
from pydantic import BaseModel, Field, validator
class AskPayload(BaseModel):
question: str = Field(..., min_length=1, max_length=500)
@validator("question")
def sanitize(cls, v):
v = v.strip()
if not v:
raise ValueError("question cannot be empty")
# Optional: reject dangerous patterns
if "<?php" in v.lower() or "<script" in v.lower():
raise ValueError("Invalid characters in question")
return v
Don't Expose Secrets in Logs
# Bad
logger.info(f"Calling OpenAI with key {api_key[:8]}...") # Avoid
# Good
logger.info("Calling OpenAI", extra={"key_prefix": api_key[:4] + "***" if api_key else "not_set"})
Uniform Error Format
# app/exceptions.py
from fastapi import Request, status
from fastapi.responses import JSONResponse
def error_response(status_code: int, detail: str, trace_id: str = None):
return JSONResponse(
status_code=status_code,
content={
"error": True,
"detail": detail,
"trace_id": trace_id,
},
)
@app.exception_handler(422)
async def validation_exception_handler(request: Request, exc):
return error_response(422, "Invalid input", getattr(request.state, "trace_id", None))
@app.exception_handler(500)
async def server_exception_handler(request: Request, exc):
logger.exception("Unhandled error")
return error_response(500, "Internal server error", getattr(request.state, "trace_id", None))
Documentation
README with Architecture
# RAG API with ChromaDB
Production-ready RAG system: ingestion, retrieval, generation, and REST API.
## Architecture
\`\`\`
Documents → Ingestion → ChromaDB → Retrieval → Generation → API
↑ ↑ ↑
chunking vector store LLM (OpenAI)
\`\`\`
## Requirements
- Python 3.11+
- Docker and Docker Compose
- OpenAI API Key
## Local Usage
\`\`\`bash
cp .env.example .env
docker-compose up -d
curl -X POST http://localhost:8000/ask -H "Content-Type: application/json" -d '{"question":"What is RAG?"}'
\`\`\`
## API Reference
| Endpoint | Method | Description |
|----------|--------|-------------|
| /health | GET | Health check |
| /search | GET | Semantic search (q, top_k) |
| /ask | POST | RAG question (question) |
| /ingest | POST | Document ingestion |
## Deployment
See [DEPLOYMENT.md](./DEPLOYMENT.md).
Automatic API Reference
FastAPI generates Swagger at /docs and ReDoc at /redoc. Make sure to document parameters:
@app.post("/ask", summary="RAG question")
async def ask(
payload: AskPayload,
request: Request,
api_key: str = Depends(verify_api_key)
):
"""
Receives a question, retrieves context from ChromaDB, and generates an answer with sources.
Requires X-API-Key in the header if configured.
"""
...
Hardening Checklist
- No secrets are exposed in logs
- There are fallbacks when ChromaDB or OpenAI fail
- 4xx/5xx errors have a uniform format
- Retry logic in calls to ChromaDB and OpenAI
- Rate limiting on public endpoints
- Optional API key for protection
- Input validation with Pydantic
- Endpoint documentation up to date
- README with architecture, setup, and troubleshooting
- Documented rollback plan
- Caching of embeddings and results (optional but recommended)
- Cache hit rate metrics
Hardening Prioritization
| Order | Area | Effort | Impact |
|---|---|---|---|
| 1 | Reliability: retries, fallbacks | Medium | High |
| 2 | Security: auth, rate limit, validation | Medium | High |
| 3 | Errors: uniform format, no stack traces | Low | Medium |
| 4 | Caching | Medium | High (performance) |
| 5 | Documentation | Medium | Medium (maintainability) |
| 6 | Rollback plan | Low | High (operations) |
Exercises with Detailed Solutions
Exercise 1: Retry with exponential backoff for OpenAI
Goal: Implement retry only for OpenAI 429 and 5xx errors.
Solution:
@retry(
stop=stop_after_attempt(3),
wait=wait_exponential(multiplier=1, min=2, max=60),
retry=retry_if_exception(lambda e: "429" in str(e) or "5" in str(e)[:1])
)
async def call_openai_completion(messages):
# ...
Exercise 2: Fallback to a short answer when the LLM times out
Goal: If OpenAI takes more than 30s, return "Query too complex, try rephrasing".
Solution:
try:
answer = await asyncio.wait_for(generate_answer(question, docs), timeout=30.0)
except asyncio.TimeoutError:
answer = "Query too complex. Try rephrasing or simplifying your question."
Exercise 3: Rate limit per API key instead of IP
Goal: Limit 100 req/min per API key for multi-tenant.
Solution:
def get_client_id(request: Request) -> str:
return request.headers.get("X-API-Key", get_remote_address(request))
limiter = Limiter(key_func=get_client_id)
Exercise 4: Cache hit rate in the /metrics endpoint
Goal: Expose a gauge rag_cache_hit_rate computed as hits/(hits+misses).
Solution:
from prometheus_client import Gauge
hit_rate = Gauge("rag_cache_hit_rate", "Cache hit rate", ["cache_type"])
# Update on each hit/miss
def update_hit_rate(cache_type):
h, m = cache_hits.labels(cache_type=cache_type)._value.get(), cache_misses.labels(cache_type=cache_type)._value.get()
hit_rate.labels(cache_type=cache_type).set(h / (h + m) if (h + m) > 0 else 0)
Exercise 5: README with a Mermaid diagram
Goal: Include a flow diagram in the README.
Solution:
## Architecture
\`\`\`mermaid
flowchart LR
A[Documents] --> B[Chunking]
B --> C[Embeddings]
C --> D[ChromaDB]
D --> E[Retrieval]
E --> F[Generation]
F --> G[API Response]
\`\`\`
Exercise 6: Rollback plan in 3 steps
Goal: Document the rollback in DEPLOYMENT.md.
Solution:
## Rollback
1. Revert to the previous image: `docker pull rag-api:v1.2.2 && docker-compose up -d rag-api`
2. If there is a ChromaDB migration: restore the volume backup
3. Verify health and metrics before closing the incident
60-Minute Hardening Exercise
Split the time like this:
| Time | Focus |
|---|---|
| 20 min | Errors: retries, fallbacks, uniform format |
| 20 min | Security: API key, rate limit, validation |
| 20 min | Docs: README, runbook, rollback |
Delivery: A list of findings and actions with priority (critical / important / nice-to-have).
Final Troubleshooting
"Everything seems fine, but we don't trust the release"
Do smoke tests + an internal canary before delivering. Deploy to staging, run the integration tests, simulate a ChromaDB failure, and verify that the fallbacks work.
"We have technical hardening but poor docs"
Without a runbook or an operational README, the project is not ready. Any developer must be able to bring up the system and understand what each endpoint does. Invest 1-2 hours in documentation.
"Rollback is not tested"
Define and validate a minimal rollback in staging: which command to run, how to restore data if needed, how to verify everything is fine.
"Cache is causing stale answers"
Adjust the TTL. For the embedding cache, 24h is usually fine; for the result cache, 1h or less if the corpus updates frequently. Add a way to invalidate the cache by pattern or a total flush.
"Rate limit is blocking legitimate users"
Increase limits or implement per-tenant limits. Monitor rejection metrics and adjust.
Summary
- You implemented robust error handling with retries, fallbacks, and controlled degradation.
- You added caching of embeddings and results with hit rate metrics.
- You applied basic security: API key, rate limiting, input validation.
- You standardized errors into a uniform format without exposing stack traces.
- You documented with a README (architecture, setup, API reference) and a deployment guide.
- You defined a rollback plan for operations.
The system is ready for serious technical delivery. The next capsule consolidates everything as the final project.
Additional Resources
- OWASP API Security — top 10 vulnerabilities
- Tenacity — retry decorators
- SlowAPI — rate limiting for FastAPI
- Redis Cache Patterns — caching
- Pydantic Validation — data validation
- FastAPI Security — API keys, OAuth
- Keep a Changelog — CHANGELOG format
- Semantic Versioning — API versioning
Estimated time: 55-65 minutes
Next: 08-project-production-ready-rag.md