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

AreaWhat to doPriority
ReliabilityRetries, timeouts, fallbacks when an external service failsHigh
SecurityAPI key, rate limiting, input validationHigh
PerformanceCaching of embeddings and resultsMedium
ErrorsUniform 4xx/5xx format, don't expose stack tracesHigh
DocumentationREADME, API reference, operations runbookMedium

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

OrderAreaEffortImpact
1Reliability: retries, fallbacksMediumHigh
2Security: auth, rate limit, validationMediumHigh
3Errors: uniform format, no stack tracesLowMedium
4CachingMediumHigh (performance)
5DocumentationMediumMedium (maintainability)
6Rollback planLowHigh (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:

TimeFocus
20 minErrors: retries, fallbacks, uniform format
20 minSecurity: API key, rate limit, validation
20 minDocs: 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


Estimated time: 55-65 minutes
Next: 08-project-production-ready-rag.md