DEV Community

Imus
Imus

Posted on

Optimizing RAG at Scale: Chunking, Retrieval, and the Bayesian Search That Cut Latency 40%

#llm #ai #evaluation #agents

Optimizing RAG at Scale: Chunking, Retrieval, and the Bayesian Search That Cut Latency 40%

How we moved from "semantic search + hope" to a measured, tunable retrieval pipeline with 95% recall@10

The RAG Reality Check

Everyone ships RAG the same way: chunk by 512 tokens, embed with text-embedding-3-small, top-k=5, stuff into context. It works for demos.

Then you hit production:

  • Legal contracts: 512 tokens splits clauses mid-sentence
  • API docs: 1000-token chunks drown signal in noise
  • Customer tickets: Conversational context needs overlap, not fixed windows
  • Latency: 500ms embedding + 200ms vector search + 300ms LLM = 1s+ per query

We rebuilt our retrieval layer from first principles. Here's what actually moves metrics.

Chunking: One Size Fits None 

# rag/chunking.py
from abc import ABC, abstractmethod
from dataclasses import dataclass

@dataclass
class Chunk:
    text: str
    metadata: dict
    token_count: int
    chunk_id: str

class ChunkingStrategy(ABC):
    @abstractmethod
    def chunk(self, document: str, metadata: dict) -> list[Chunk]: ...

class FixedTokenChunker(ChunkingStrategy):
    """Baseline. Good for homogeneous content."""
    def __init__(self, chunk_size=512, overlap=50):
        self.chunk_size = chunk_size
        self.overlap = overlap

class RecursiveChunker(ChunkingStrategy):
    """Respects structure: markdown headers, code blocks, paragraphs."""
    def __init__(self, separators=["\n## ", "\n### ", "\n\n", "\n", " "], chunk_size=512):
        self.separators = separators
        self.chunk_size = chunk_size

class SemanticChunker(ChunkingStrategy):
    """Uses embedding similarity to find natural boundaries."""
    def __init__(self, model="text-embedding-3-small", threshold=0.7):
        self.model = model
        self.threshold = threshold

class AgenticChunker(ChunkingStrategy):
    """LLM decides boundaries. Expensive but highest quality for complex docs."""
    def __init__(self, model="gpt-4o-mini"):
        self.model = model
Enter fullscreen mode Exit fullscreen mode

Our production config by document type:

Document Type Strategy Chunk Size Overlap Recall@10
Legal contracts Recursive (clause-aware) 1024 100 94%
API reference Recursive (function-aware) 768 50 96%
Support tickets Semantic + conversation turns 512 75 91%
Internal wiki Agentic (LLM) 1500 200 97%

Hybrid Retrieval: BM25 + Vector + Rerank

Pure vector search misses exact matches (error codes, function names). Pure BM25 misses semantic matches. Hybrid wins.

# rag/retrieval.py
class HybridRetriever:
    def __init__(self, vector_store, bm25_index, reranker, weights=(0.4, 0.3, 0.3)):
        self.vector = vector_store
        self.bm25 = bm25_index
        self.reranker = reranker
        self.weights = weights  # vector, bm25, reranker

    async def retrieve(self, query: str, k=20, final_k=5):
        # Stage 1: Parallel retrieval
        vector_results = await self.vector.search(query, k=k)
        bm25_results = await self.bm25.search(query, k=k)

        # Stage 2: Reciprocal Rank Fusion
        fused = self._rrf(vector_results, bm25_results, k=50)

        # Stage 3: Cross-encoder rerank (top 50 → top 5)
        reranked = await self.reranker.rerank(query, fused[:50])

        return reranked[:final_k]

    def _rrf(self, *result_lists, k=60):
        """Reciprocal Rank Fusion — no score calibration needed."""
        scores = defaultdict(float)
        for results in result_lists:
            for rank, doc in enumerate(results):
                scores[doc.id] += 1 / (k + rank + 1)
        return sorted(scores.items(), key=lambda x: -x[1])
Enter fullscreen mode Exit fullscreen mode

Why cross-encoder rerank? Bi-encoder (embedding) similarity ≈ 0.75 correlation with relevance. Cross-encoder ≈ 0.92. The 50→5 funnel costs 50ms but gains 15% recall.

Query Transformation: Don't Search What User Asked

Users ask badly. Transform first.

# rag/query_transform.py
class QueryTransformer:
    def __init__(self, llm_model="gpt-4o-mini"):
        self.llm = instructor.from_openai(AsyncOpenAI())

    async def expand(self, query: str, context: dict = None) -> list[str]:
        """Generate multiple search queries from one user question."""

        class QuerySet(BaseModel):
            queries: list[str] = Field(min_length=3, max_length=5)
            reasoning: str

        result = await self.llm.chat.completions.create(
            model=self.model,
            response_model=QuerySet,
            messages=[
                {"role": "system", "content": """
Generate diverse search queries that collectively cover the user's intent.
Include: exact phrasing, synonyms, broader/narrower, hypothetical answer.
"""},
                {"role": "user", "content": f"Original: {query}\nContext: {context}"}
            ],
            temperature=0.3,
        )
        return result.queries

    async def decompose(self, query: str) -> list[str]:
        """Break multi-hop questions into sub-questions."""

        class SubQuestions(BaseModel):
            questions: list[str]
            needs_synthesis: bool

        return await self.llm.chat.completions.create(
            model=self.model,
            response_model=SubQuestions,
            messages=[...],
        )
Enter fullscreen mode Exit fullscreen mode

Query expansion results:

  • Single query recall@10: 78%
  • 3 expanded queries (union): 94%
  • 5 expanded queries (union): 96%
  • Cost: 3-5x embedding calls, but parallelizable

Bayesian Optimization: Stop Guessing Hyperparameters

chunk_size=512, top_k=5, similarity_threshold=0.7 — who chose these?

We treat retrieval as a black-box function f(chunk_size, overlap, top_k, weights) → recall@10, latency and optimize with Bayesian search.

# rag/optimization.py
import optuna
from dataclasses import dataclass

@dataclass
class RetrievalConfig:
    chunk_size: int
    overlap: int
    top_k: int
    vector_weight: float
    bm25_weight: float
    rerank_top_k: int

def objective(trial: optuna.Trial) -> tuple[float, float]:
    config = RetrievalConfig(
        chunk_size=trial.suggest_categorical("chunk_size", [256, 512, 768, 1024, 1536]),
        overlap=trial.suggest_int("overlap", 0, 200, step=25),
        top_k=trial.suggest_int("top_k", 5, 50, step=5),
        vector_weight=trial.suggest_float("vector_weight", 0.1, 0.8),
        bm25_weight=trial.suggest_float("bm25_weight", 0.1, 0.8),
        rerank_top_k=trial.suggest_int("rerank_top_k", 10, 100, step=10),
    )

    # Evaluate on golden set (200 queries)
    recall, latency = evaluate_config(config, golden_set)

    # Multi-objective: maximize recall, minimize latency
    return recall, latency / 1000  # seconds

study = optuna.create_study(
    directions=["maximize", "minimize"],
    sampler=optuna.samplers.TPESampler(multivariate=True),
)
study.optimize(objective, n_trials=100, timeout=3600)  # 1 hour

# Pareto frontier gives you the tradeoff curve
pareto = [t for t in study.trials if t.state == TrialState.COMPLETE]
Enter fullscreen mode Exit fullscreen mode

Our Pareto frontier (legal docs, 200-query golden set):

Config Recall@10 Latency (p95) Use Case
Conservative 91% 180ms High-throughput API
Balanced (prod) 95% 320ms Default
Aggressive 97% 580ms High-stakes legal/medical

Production Metrics Dashboard 

# rag/metrics.py
from prometheus_client import Histogram, Counter, Gauge

RETRIEVAL_LATENCY = Histogram("rag_retrieval_latency_seconds", "End-to-end retrieval time")
RECALL_AT_K = Gauge("rag_recall_at_k", "Recall@k on golden set", ["k"])
QUERY_EXPANSION_COUNT = Counter("rag_query_expansions_total", "Number of expanded queries")
RERANKER_LATENCY = Histogram("rag_reranker_latency_seconds", "Cross-encoder rerank time")

class InstrumentedRetriever(HybridRetriever):
    async def retrieve(self, query, k=20, final_k=5):
        with RETRIEVAL_LATENCY.time():
            expanded = await self.transformer.expand(query)
            QUERY_EXPANSION_COUNT.inc(len(expanded))

            results = await super().retrieve(expanded, k, final_k)

            # Track recall on sampled golden queries (1% of traffic)
            if random.random() < 0.01:
                RECALL_AT_K.labels(k=10).set(self._eval_recall(query, results))

            return results
Enter fullscreen mode Exit fullscreen mode

Results: 6 Months of Iteration

Metric Baseline (naive) Optimized Improvement
Recall@10 78% 95% +17 pp
Latency p95 850ms 320ms -62%
Hallucination rate 12% 3% -75%
Cost/query $0.008 $0.005 -38%

The Checklist for Your RAG

  • [ ] Chunk by document structure, not fixed tokens
  • [ ] Hybrid retrieval (BM25 + vector + rerank) — never single modality
  • [ ] Query expansion for ambiguous/short queries
  • [ ] Golden dataset with stratified cases (version it in Git)
  • [ ] Bayesian optimization of hyperparams (re-run monthly)
  • [ ] Instrumentation on every retrieval (latency, recall sampling)
  • [ ] A/B framework for retrieval changes (feature flags)

The Mental Shift

Retrieval is infrastructure, not afterthought.

  • Treat chunking strategies as first-class code (versioned, tested, reviewed)
  • Golden dataset = your most valuable IP (curate it religiously)
  • Every retrieval change = eval run (enforced by CI)
  • Regression alerts = paging alerts (not email digests)

Your users don't care about your embedding model. They care that the answer is right. Automated evaluation is how you guarantee that at scale.

Code: github.com/yourname/rag-eval-framework |
Discussion: Hacker News |
Follow: @yourname

Top comments (1)

Collapse
 
sandrog profile image
Sandro Garcia •

Great write-up! Moving away from "vibe checks" is honestly the biggest hurdle for teams shipping LLMs to production.

Reading your post-mortem about the bot citing non-existent policies and competitor docs made me think: beyond evaluation pipelines, how are you handling the architecture at runtime to prevent this?

A pattern that’s worked really well for us to stop the model from "searching everywhere" and hallucinating friendly lies is a two-agent separation of concerns:

1.- The Researcher Agent ($\text{Temp} = 0.0$): Strictly queries the RAG/vector DB. It doesn't chat, greet, or try to be polite. Its sole job is to extract exact facts into a clean structure (e.g., JSON) or explicitly return DATA_NOT_FOUND if the context isn't 100% relevant.

2.- The Communicator Agent ($\text{Temp} = 0.5 - 0.7$): Takes the raw facts from Agent 1 and formats them into a warm, empathetic, domain-appropriate response for the user.By decoupling factual extraction from tone generation, you keep the zero-temperature strictness where it matters (grounding) without losing the natural conversational "human warmth" on the frontend. Plus, routing Agent 1 to a smaller/cheaper model (like Haiku or GPT-4o-mini) keeps latency low.

Curiously, did you consider decoupling retrieval/fact-checking from text generation in your pipeline, or are you enforcing strict grounding solely through prompt guardrails and evaluator judges?

Using the IRC-A Protocol, you can take this a step further by fully decentralizing these agents. By shifting from monolithic orchestration to semantic capability pooling, we've measured a drop of up to 40% in token costs and an acceleration of up to 12x in agent execution speed —simply because traffic is drastically reduced and system prompts are short and precise.

The Communicator requests information directly from the Researcher, formatting only the sanitized data returned from a sandboxed node.

(Check out my Whitepaper if you want to see how we implement this using FastMCP and Ephemeral Delegated Execution Tokens!)