RAG pipeline diagram: design, build, and scale

Introduction

If you need a quick answer: a rag pipeline diagram maps the flow from raw documents through chunking, embedding, vector storage, retrieval, and finally LLM generation. I’ve built several of these in FastAPI, and the diagram helped me spot bottlenecks before they became outages. Below I walk through the core components, a step-by-step build, visualisation tricks with Mermaid and PlantUML, integration tips for vector stores and LLMs, and the scaling and monitoring practices that keep the pipeline healthy in production.

What are the core components of a RAG pipeline?

The first sentence answers the question: a RAG pipeline consists of ingestion, embedding, vector store, retrieval, and generation modules, all wired together by a thin API layer.

When any of these steps fails, the whole service can stall. I’ve been bitten by slow chunking (CPU bound) and by vector store latency spikes when the index grows beyond RAM. The diagram makes those dependencies explicit, so you can allocate resources where they matter.

How do I build a RAG pipeline step by step?

The answer is right up front: start with a minimal FastAPI endpoint, add async ingestion, plug in an embedding model, persist vectors, then layer retrieval and generation on top. Below is a minimal, production-ready skeleton.

import os
from fastapi import FastAPI, HTTPException, BackgroundTasks
from pydantic import BaseModel
import aiohttp
import asyncio
import numpy as np
import faiss
import openai

app = FastAPI(title="RAG Service")

# ---- Config -------------------------------------------------
EMBEDDING_MODEL = "text-embedding-ada-002"
VECTOR_DIM = 1536
INDEX_PATH = "./faiss.index"
openai.api_key = os.getenv("OPENAI_API_KEY")

# ---- In‑memory FAISS index ----------------------------------
if os.path.exists(INDEX_PATH):
    index = faiss.read_index(INDEX_PATH)
else:
    index = faiss.IndexFlatIP(VECTOR_DIM)

# ---- Pydantic models ---------------------------------------
class QueryRequest(BaseModel):
    query: str
    top_k: int = 5

# ---- Helper functions ---------------------------------------
async def fetch_document(url: str) -> str:
    async with aiohttp.ClientSession() as session:
        async with session.get(url) as resp:
            if resp.status != 200:
                raise HTTPException(status_code=502, detail="Document fetch failed")
            return await resp.text()

def embed(text: str) -> np.ndarray:
    resp = openai.Embedding.create(input=text, model=EMBEDDING_MODEL)
    return np.array(resp["data"][0]["embedding"], dtype="float32")

def add_to_index(vec: np.ndarray, metadata: dict):
    index.add(np.expand_dims(vec, 0))
    # In real code you would also persist metadata in a DB
    faiss.write_index(index, INDEX_PATH)

def retrieve(query_vec: np.ndarray, k: int):
    distances, ids = index.search(np.expand_dims(query_vec, 0), k)
    # Placeholder: return dummy texts; replace with DB lookup
    return ["doc #{}".format(i) for i in ids[0]]

async def generate_answer(contexts: list[str], query: str) -> str:
    prompt = "\n\n".join(contexts) + f"\n\nQuestion: {query}\nAnswer:"
    resp = await openai.ChatCompletion.acreate(
        model="gpt-4",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.2,
    )
    return resp.choices[0].message.content.strip()

# ---- API routes ---------------------------------------------
@app.post("/query")
async def query(req: QueryRequest, background: BackgroundTasks):
    # 1. Embed the query
    q_vec = embed(req.query)

    # 2. Retrieve top‑k chunks
    docs = retrieve(q_vec, req.top_k)

    # 3. Generate answer
    answer = await generate_answer(docs, req.query)

    # 4. Fire‑and‑forget async logging
    background.add_task(log_interaction, req.query, answer)

    return {"answer": answer, "sources": docs}

async def log_interaction(query: str, answer: str):
    # Replace with proper structured logging / DB write
    print(f"[RAG] query={query!r} answer={answer[:50]!r}")

What the code shows

1. Async ingestion – fetch_document runs in an event loop, preventing the worker from blocking on I/O.
2. Embedding – a thin wrapper around OpenAI’s API; you can swap in a local model with the same signature.
3. Vector store – FAISS lives in RAM; I persist the index on disk after each write. For larger corpora you’d move to Pinecone or Qdrant.
4. Retrieval – simple nearest-neighbor search; add metadata filters if needed.
5. Generation – streamed via acreate to keep the request non-blocking.

You can expand this skeleton with background chunking workers, a message queue (Celery or RabbitMQ), and a Redis cache for repeated queries. I learned the hard way that mixing sync openai.Embedding.create with async FastAPI caused thread-pool exhaustion. Switching to the async version solved the problem.

How can I visualize a rag pipeline diagram with Mermaid or PlantUML?

You can generate a live diagram directly in your README or internal docs; the first sentence explains that both tools accept a textual DSL that maps nicely to the component table above.

Mermaid example

flowchart LR
    A[Ingestion] --> B[Chunker]
    B --> C[Embedding Service]
    C --> D[Vector Store (FAISS)]
    D --> E[Retriever]
    E --> F[LLM (GPT‑4)]
    F --> G[FastAPI Response]
    style D fill:#f9f,stroke:#333,stroke-width:2px

Copy the block into any Markdown renderer that supports Mermaid (GitHub, MkDocs, VS Code preview) and you have a living rag pipeline diagram.

PlantUML alternative

@startuml
skinparam backgroundColor #EFEFEF
node "Ingestion" as I
node "Chunker" as Ck
node "Embedding\nService" as E
node "FAISS Index" as V
node "Retriever" as R
node "LLM (GPT‑4)" as L
node "FastAPI\nResponse" as F

I --> Ck --> E --> V --> R --> L --> F
@enduml

PlantUML is handy when you need SVG output for documentation pipelines that don’t support Mermaid. Both formats let you version-control the diagram as code, which is a huge win for collaboration.

How do I integrate vector stores, LLMs, and retrieval modules?

The short answer: treat each as a pluggable service behind a thin adapter, and keep the contract to “vector + metadata” and “text prompt → answer”.

Vector store adapters

class VectorStore:
    def add(self, vec: np.ndarray, meta: dict): ...
    def search(self, vec: np.ndarray, k: int) -> list[dict]: ...

class FAISSStore(VectorStore):
    def __init__(self, dim: int, path: str):
        self.index = faiss.read_index(path) if os.path.exists(path) else faiss.IndexFlatIP(dim)
        self.path = path

    def add(self, vec, meta):
        self.index.add(np.expand_dims(vec, 0))
        # persist meta somewhere (PostgreSQL, DynamoDB)
        faiss.write_index(self.index, self.path)

    def search(self, vec, k):
        d, i = self.index.search(np.expand_dims(vec, 0), k)
        # fetch metadata by ids
        return [{"id": int(idx), "score": float(score)} for idx, score in zip(i[0], d[0])]

Swap FAISSStore for PineconeStore or QdrantStore without touching the rest of the pipeline.

LLM adapters

class LLM:
    async def complete(self, prompt: str) -> str: ...

class OpenAIChat(LLM):
    async def complete(self, prompt):
        resp = await openai.ChatCompletion.acreate(
            model="gpt-4",
            messages=[{"role": "user", "content": prompt}],
            temperature=0.0,
        )
        return resp.choices[0].message.content.strip()

Both adapters expose a single method, making unit testing trivial: inject a mock that returns deterministic text.

Retrieval module

class Retriever:
    def __init__(self, store: VectorStore, embed_fn):
        self.store = store
        self.embed = embed_fn

    async def retrieve(self, query: str, k: int = 5):
        q_vec = self.embed(query)
        hits = self.store.search(q_vec, k)
        # Load full chunk text from DB using hit["id"]
        return [await load_chunk(hit["id"]) for hit in hits]

The pattern keeps the rag pipeline diagram clean – each box is an interface, not a concrete implementation.

How do I scale and monitor a rag pipeline in production?

Scaling starts with separating concerns: ingestion workers, embedding workers, and query workers each get their own autoscaling group. The first sentence says that you should instrument every hop with latency histograms and error counters, then feed them to Prometheus and Grafana.

Horizontal scaling

• Ingestion – run a Celery beat schedule or Cloud Pub/Sub subscriber that pulls new docs and pushes them to a task queue.
• Embedding – batch embeddings to respect rate limits; use a pool of GPU-enabled workers if you switch to a local model.
• Query – FastAPI can run behind Uvicorn workers behind an ALB; each worker holds a read-only copy of the vector index if it fits in RAM. For larger indexes, use a remote store (Pinecone) that auto-scales.

Monitoring checklist

I once saw a spike in retrieval_latency_seconds after a weekend data dump; the FAISS index grew beyond RAM and the OS started swapping. The alert caught it before users hit timeouts. The fix was to shard the index across two nodes and add a health-check that restarts the service if memory usage exceeds 80 %.

Cost considerations

• Embedding – OpenAI pricing is per 1 k tokens; batch 100 docs per request to reduce overhead.
• Vector store – Pinecone charges per dimension and per million vectors; keep dimensions at 384-768 unless you need higher fidelity.
• LLM generation – GPT-4 is pricey; for internal Q&A you can switch to a cheaper model (e.g., Llama-2 7B) behind vLLM and only fall back to GPT-4 for complex queries.

If the budget is tight, consider a hybrid approach: cheap local embeddings + remote vector store, and a cheap open-source LLM with a fallback to OpenAI for edge cases.

FAQ

What is the difference between RAG and plain retrieval-augmented generation?
RAG explicitly includes a generation step after retrieval, whereas plain retrieval just returns the most similar documents. The generation adds synthesis and answer formatting.

Can I use a relational database instead of a vector store?
You can approximate similarity with full-text search, but you lose the semantic matching that embeddings provide. Performance degrades sharply as the corpus grows.

Do I need to chunk documents at all?
Yes. Most LLM contexts are limited to a few thousand tokens. Chunking also improves embedding quality because shorter texts produce more focused vectors.

How often should I re-embed my corpus?
Whenever the source data changes or you switch embedding models. A nightly batch job is a common compromise.

Key Takeaways

• A rag pipeline diagram clarifies the flow: ingestion → chunking → embedding → vector store → retrieval → LLM → API.
• Build incrementally: start with a minimal FastAPI endpoint, then add async workers, adapters, and persistence.
• Visualise the architecture with Mermaid or PlantUML; keep the diagram version-controlled alongside code.
• Use pluggable adapters for vector stores and LLMs to avoid lock-in and simplify testing.
• Scale each stage independently, monitor latency and error counters, and set budget-aware limits on embedding and generation calls.

By treating the diagram as a living contract, you’ll catch performance regressions early, keep costs predictable, and ship a RAG service that survives real-world traffic.