Optimizing Token Usage in RAG: Implementing Dynamic Context Pruning

# Optimizing Token Usage in RAG: Implementing Dynamic Context Pruning


*By the engineering team at OnlineCloud — a US-based Agentic AI & MCP engineering studio*

—

## 1. Introduction: The Linear Cost of Long-Context Retrieval

Every RAG pipeline I audit has the same hidden line item. Not in the vector database bill, not in the embedding batch job, but in the prompt assembly code: `context = “\n\n”.join(chunks[:top_k])`. That single slice — usually `top_k=10` or `top_k=20` — decides how many tokens the generation model re-reads on every single request. And because a chunk that added retrieval value in week one rarely gets removed in week twelve, the context grows fat, the latency creeps up, and the cost-per-query line on the FinOps dashboard doubles before anyone asks why.

The economics are brutally linear. A knowledge-base query that matches twenty five-hundred-token chunks ships ten thousand tokens of context to the model — before the question, before the instructions, before the answer. At GPT-4-class input pricing, that is roughly $0.10 per request in context alone, and it compounds: users ask follow-ups, agents run multi-turn loops, and every turn re-encodes the same bloated window. Long-context models made this *possible*; they did not make it *cheap*. The 200K-token window is a safety margin, not a budget to spend.

Worse, the waste is not neutral — it is actively harmful. Research on the “lost in the middle” phenomenon shows models systematically under-weight tokens in the middle of their context window. Stuff the window with marginal chunks and you are not just overpaying; you are diluting the signal that the high-value passages carry. Pruning is therefore a **quality** optimization disguised as a **cost** optimization.

This post walks through the four layers where we cut token spend in production RAG — dynamic pruning at retrieval time, semantic compression before generation, resource-level pruning in the Model Context Protocol (MCP) transport layer, and the ROI math that justifies each — then closes with a working Python/Pydantic implementation of a context optimizer you can adapt today.

—

## 2. Dynamic Pruning: Rerankers and Relevance Scoring

The first and highest-leverage cut happens before generation: **stop shipping chunks that do not answer the question.** Vector search retrieves by embedding proximity, which is a fuzzy proxy for relevance. A bi-encoder embeds query and document independently, then measures cosine distance — fast enough to scan millions of vectors, but blind to the lexical and syntactic interactions that actually decide whether a passage answers a question. Top-20 by cosine similarity reliably contains *some* answer-bearing chunks; it rarely contains *only* them.

The fix is a second-stage **cross-encoder reranker**. Unlike a bi-encoder, a cross-encoder takes the query and a candidate passage *together* through the transformer, producing an interaction-aware relevance score in (0, 1). It is too slow to run over the whole corpus — that is what the approximate-nearest-neighbor index is for — but running it over the top 20–50 candidates costs only a few milliseconds per chunk. The standard recipe:

1. **Recall stage:** embed the query, retrieve the top 40 candidates by cosine similarity (or hybrid BM25 + vector fusion via reciprocal rank fusion — dense vectors alone routinely miss exact identifiers like part numbers and policy IDs).
2. **Precision stage:** score all 40 candidates with a cross-encoder reranker (e.g., a Cohere or BGE reranker, or a fine-tuned `cross-encoder/ms-marco-MiniLM` for on-prem deployments).
3. **Prune stage:** drop everything below a relevance floor *and* enforce a hard token budget, whichever binds first.

The reranker gives you something the vector index cannot: a **calibrated cutoff**. With raw similarity you are guessing; with cross-encoder scores you can set `min_score = 0.35`, measure precision on your eval set, and move the threshold until false-positive context stops appearing in prompts. In our production builds this step alone routinely removes 60–75% of retrieved tokens while holding answer accuracy flat — and frequently *improving* it, because the model no longer has to reason past ten irrelevant passages to find the one good answer.

Two engineering cautions. First, rerank *after* metadata filtering and row-level access enforcement, never before — you do not want to spend ranking compute on chunks the caller cannot legally see. Second, log the scores. The distribution of reranker scores across your traffic is the single best early-warning signal for embedding drift and corpus decay; if the mean relevance score of served chunks is sliding down month over month, your ingestion pipeline is rotting even though nothing is erroring.

—

## 3. Semantic Compression: Summarize Before You Generate

Pruning removes what is irrelevant. **Semantic compression** shrinks what remains. The insight is that a retrieved passage usually contains one load-bearing fact surrounded by prose, examples, and boilerplate — and the generation model does not need the boilerplate to answer correctly.

The pattern is map-reduce summarization applied *inside* the retrieval path. First, group the surviving chunks (optionally by source document or by cluster) and run a cheap summarizer over each group with a strict instruction: *”Preserve every fact, number, name, date, and caveat; remove restatement, examples, and filler. Output at most N tokens.”* Then concatenate the compressed groups into the generation context. Compression ratios of 3:1 to 5:1 are routine on enterprise documents, which are optimized for human skimming, not model ingestion.

The design decision that matters is **what gets compressed versus what stays verbatim**. We classify each reranked chunk into tiers:

| Tier | Reranker score | Handling |
|—|—|—|
| High | ≥ 0.7 | **Verbatim.** Quote-bearing, compliance-critical passages go to the prompt untouched. |
| Medium | 0.35 – 0.7 | **Compressed.** Summarize to ~25% length, preserving all atomic facts. |
| Low | < 0.35 | **Dropped.** Not retrieved at all. | This tiering exists because compression is lossy, and loss is not evenly distributed. A summarizer can faithfully reduce a two-hundred-word eligibility explanation to fifty words; it will occasionally mangle a precise dollar amount or a statutory date. High-value chunks earn the right to full fidelity; medium-value chunks trade a bounded fidelity risk for a 4× token reduction. Quantify that risk on your eval set (see Section 5) before you trust it in production — if compression drops factual-recall on your benchmark by more than ~1–2%, lower the verbatim threshold. Compression also has a **latency tax**: each summarization call is a small LLM round-trip. Amortize it by caching compressed forms keyed on the chunk hash (documents change rarely; their summaries should not be recomputed per query), and by batching — one call compressing three chunks beats three calls compressing one. Done right, the summarization cost lands at 10–20% of the token savings it unlocks, which is the trade you want. --- ## 4. Protocol Layer Optimization: Partial Resources in MCP If you are building RAG on top of Model Context Protocol — as we increasingly do at OnlineCloud — there is a fourth layer most teams miss, and it is the most elegant one. **The bloat can be cut before a single token ever leaves your server.** MCP servers expose data to agents as **resources**, addressed by URI and returned as typed contents — the contract layer between LLMs and your business logic. The naive pattern is to expose a document as one resource: the agent requests `mcp://kb/policy-42`, and the server returns the whole document. That is prompt-stuffing wearing a protocol hat, and it reproduces the linear-cost problem at the transport boundary. The fix is what we call **partial resources**: resource templates that let the client request precisely the span it needs, with the server doing the pruning against the *real* retrieval pipeline rather than the model doing it against a wall of text. Concretely: - **Offset/limit templates.** A resource template `mcp://kb/{doc_id}` accepts a query parameter block — `mcp://kb/policy-42?start=1200&limit=600` — and the server returns only that byte/token span, mirroring HTTP Range semantics. The agent's context-manager, not the model, decides the window. - **Projection queries.** For structured knowledge (Oracle, MS SQL, Snowflake), expose `mcp://kb/{doc_id}?fields=eligibility,effective_date` and let the server emit a row-level-filtered, column-projected view. The database does the pruning; the wire and the model never see the columns the task does not need. - **Scored slices.** For semantic resources, `mcp://kb/{doc_id}?query={question}&top=3&min_score=0.35` runs the retrieval + rerank pipeline *server-side* and returns only the pruned, ranked passages — the exact pipeline from Sections 2 and 3, enforced at the protocol seam. Why this matters architecturally: MCP's design puts the *application* in control of what gets loaded, when, and with what caching. That control is wasted if the only granularity your server offers is "the whole document." Partial resources move context optimization from an emergent property of prompts — unpredictable, unobservable, different for every model — into a deterministic, typed, auditable contract. Every resource read is a log line with the span requested, the spans dropped, and the score threshold applied. Your CFO gets a token-accounting trail; your eval suite gets a stable interface to test against; and your agents stop paying for the middle of documents they never needed. --- ## 5. Cost ROI: Token Savings vs. Accuracy Trade-offs None of this is worth doing if it breaks answers. The discipline is to treat context optimization as an experiment with two measured axes — **cost per resolved query** and **answer quality on a frozen eval set** — and to make the trade-off explicit before shipping. Here is a representative before/after from a production knowledge-base deployment (200–400 token chunks, GPT-4-class model, ~$10/1M input tokens): | Stage | Tokens sent per query | Cost per 1k queries | Latency added | |---|---|---|---| | Baseline (top-20 chunks, no rerank) | ~8,000 | ~$80 | — | | + Rerank & prune (top-5, score floor) | ~2,000 | ~$20 | +40 ms | | + Semantic compression (tiered) | ~800 | ~$8 | +250 ms (cached: ~+30 ms) | | + MCP partial resources | pays only for the span served | ~$6–8 | −network/parse overhead | The 75% token reduction from pruning is nearly free accuracy-wise: cross-encoder rerankers consistently beat bi-encoder top-k on passage-ranking benchmarks, so your *precision* goes up as your *volume* goes down. The honest trade-off lives in the compression tier, where lossy summarization can shave factual recall. The measurement protocol we use: 1. Freeze a labeled eval set of ~500 real queries with expected answer facts. 2. Measure baseline factual-recall and exact-match on served-chunk text. 3. Turn on pruning; confirm recall does not regress beyond noise (±1%). 4. Turn on compression; if factual-recall drops more than 1–2%, raise the verbatim threshold or drop the compression ratio. 5. Compute **net ROI = (tokens saved × blended price) − (rerank + summarization compute)** and track it on a dashboard, per corpus, per quarter. In our deployments the arithmetic lands almost monotonically positive: a system answering 100k queries/month saves $5–8k/month in context tokens against a few hundred dollars of reranker inference — before counting the secondary wins of lower latency, higher throughput per rate-limit, and smaller KV-cache footprints on self-hosted GPUs. The remaining accuracy tax, when it exists at all, sits in the 0.5–1.5% band and is directly tunable. That is a trade you can take to a budget meeting with a straight face. --- ## 6. Implementation: A Python / Pydantic Context Optimizer Here is the optimizer we actually run — compact, typed, and eval-friendly. It implements the full pipeline: rerank-and-prune against a token budget, tiered compression, and deterministic prompt assembly. Pydantic gives us validated boundaries between stages, so a malformed chunk or an out-of-range budget fails loudly at the seam instead of corrupting a prompt. ```python """Dynamic context pruning for RAG: rerank -> budget -> compress -> assemble.”””
from __future__ import annotations

import hashlib
from dataclasses import dataclass, field
from typing import Callable, Protocol, Sequence

from pydantic import BaseModel, Field, field_validator

# ————————————————————————— #
# Stage 1: typed data contracts
# ————————————————————————— #

class RetrievedChunk(BaseModel):
“””A candidate passage returned by the vector recall stage.”””
doc_id: str
text: str
tokens: int = Field(ge=0)

@field_validator(“text”)
@classmethod
def reject_empty(cls, v: str) -> str:
if not v.strip():
raise ValueError(“chunk text must not be empty”)
return v

def cache_key(self) -> str:
return hashlib.sha256(self.text.encode()).hexdigest()

class ScoredChunk(RetrievedChunk):
“””A chunk after cross-encoder relevance scoring.”””
relevance: float = Field(ge=0.0, le=1.0)

class CompressedChunk(BaseModel):
“””A chunk after semantic compression, preserving atomic facts.”””
doc_id: str
text: str
source_tokens: int
tokens: int
verbatim: bool = False

class ContextPlan(BaseModel):
“””The final, budgeted context handed to the generation model.”””
chunks: list[CompressedChunk]
total_tokens: int = 0
dropped_count: int = 0

@field_validator(“total_tokens”, mode=”before”)
@classmethod
def auto_sum(cls, _: int, info) -> int: # pragma: no cover – defensive
return sum(c.tokens for c in info.data.get(“chunks”, []))

# ————————————————————————— #
# Stage 2: pluggable scorers and summarizers (swap per deployment)
# ————————————————————————— #

class RelevanceScorer(Protocol):
def score(self, query: str, chunk: RetrievedChunk) -> float: …

class Summarizer(Protocol):
def summarize(self, chunk: ScoredChunk, max_tokens: int) -> str: …

@dataclass
class BudgetedContextOptimizer:
“””Rerank candidates, enforce a token budget, compress, and assemble.”””

scorer: RelevanceScorer
summarizer: Summarizer
max_context_tokens: int = 3_000
min_score: float = 0.35
verbatim_threshold: float = 0.70
compression_ratio: float = 0.25 # medium-tier chunks shrink to 25%

compression_cache: dict[str, str] = field(default_factory=dict)

# — public API ——————————————————– #

def optimize(self, query: str, candidates: Sequence[RetrievedChunk]) -> ContextPlan:
“””Full pipeline: score -> prune to budget -> compress -> assemble.”””
scored = self._score(query, candidates)
pruned = self._prune(query, scored)
kept, dropped = self._fit_budget(pruned)
compressed = [self._compress(c) for c in kept]
return ContextPlan(chunks=compressed, dropped_count=dropped)

# — internals ——————————————————— #

def _score(self, query: str, candidates: Sequence[RetrievedChunk]) -> list[ScoredChunk]:
“””Run the cross-encoder over recall-stage candidates only.”””
return [
ScoredChunk(**c.model_dump(), relevance=self.scorer.score(query, c))
for c in candidates
]

def _prune(self, query: str, scored: Sequence[ScoredChunk]) -> list[ScoredChunk]:
“””Drop chunks below the calibrated relevance floor.”””
return [c for c in scored if c.relevance >= self.min_score]

def _fit_budget(self, pruned: list[ScoredChunk]) -> tuple[list[ScoredChunk], int]:
“””Greedy fill: keep the most relevant chunks until the budget binds.”””
ordered = sorted(pruned, key=lambda c: c.relevance, reverse=True)
kept, budget = [], self.max_context_tokens
for chunk in ordered:
if chunk.tokens <= budget: kept.append(chunk) budget -= chunk.tokens return kept, len(pruned) - len(kept) def _compress(self, chunk: ScoredChunk) -> CompressedChunk:
“””Verbatim for high-relevance chunks; cached summary for the rest.”””
if chunk.relevance >= self.verbatim_threshold:
return CompressedChunk(
doc_id=chunk.doc_id, text=chunk.text,
source_tokens=chunk.tokens, tokens=chunk.tokens, verbatim=True,
)
key = chunk.cache_key()
summary = self.compression_cache.get(key)
if summary is None:
budget = max(64, int(chunk.tokens * self.compression_ratio))
summary = self.summarizer.summarize(chunk, max_tokens=budget)
self.compression_cache[key] = summary
return CompressedChunk(
doc_id=chunk.doc_id, text=summary,
source_tokens=chunk.tokens,
tokens=len(summary) // 4, # ~4 chars/token; replace with a tokenizer
)

def assemble(self, plan: ContextPlan, question: str, system: str) -> str:
“””Deterministic prompt assembly — the only place the model sees text.”””
blocks = [
f”[{c.doc_id} | {‘verbatim’ if c.verbatim else ‘summary’}]\n{c.text}”
for c in plan.chunks
]
return f”{system}\n\nCONTEXT:\n” + “\n\n—\n\n”.join(blocks) + \
f”\n\nQUESTION: {question}\nANSWER:”

def count_tokens(text: str) -> int:
“””Bare-bones fallback token estimate (replace with tiktoken in prod).”””
return len(text) // 4
“`

A few implementation notes worth carrying into your codebase:

– **Budget before compression.** Prune to the token budget *first*, then compress what survives. Compressing chunks you are about to drop wastes summarizer spend — the expensive operation in this pipeline.
– **Boundaries are Pydantic-validated.** `Relevance` outside [0,1], negative token counts, empty chunk text — all rejected at construction. When the optimizer is fed from a live retrieval service, this catches schema drift at the seam, not as a garbled prompt three hundred tokens later.
– **Cache by content hash, not doc_id.** The same passage appearing under two document versions compresses once. Invalidation is trivial: corpus re-ingestion bumps the hash.
– **Return the plan, not just the string.** `ContextPlan` carries `dropped_count` and per-chunk provenance so your observability layer can log *why* each prompt looks the way it does — the audit trail your eval and your finance team both need.
– **The scorers and summarizers are protocols.** In production these are a Cohere/BGE reranker over HTTP and a small summarization model; in tests they are stubs returning fixed values. The optimizer itself never changes when the model vendors do.

—

## Conclusion: Prune at Every Layer, Measure at Every Step

Token optimization in RAG is not a single knob — it is a stack of cuts, each with its own cost and its own risk. Prune aggressively at retrieval with a calibrated reranker and you remove the 60–75% of context that was never going to help. Compress what remains, tiered by relevance, and you shrink the final prompt another 60% at a bounded, measurable fidelity cost. Push the pruning into the protocol layer with MCP partial resources and the savings become deterministic and auditable instead of prompt-dependent. Then close the loop with an ROI measurement that tracks tokens saved against factual recall — because an optimization you cannot measure is a regression you cannot see.

The systems that survive contact with production are the ones where every token that reaches the model has earned its place: scored, budgeted, compressed, and accounted for. That discipline is exactly what we build at OnlineCloud — evaluated, observable RAG and agentic systems wired into your real data estate, from Oracle and MS SQL to Postgres and Snowflake, delivered by senior engineers in 6–10 weeks.

**Ready to cut the token waste in your retrieval pipeline?** Talk to the OnlineCloud architects. We will map your data, pick the model class, and write the eval that defines “done” — with the cost and accuracy numbers on the table before you commit. [Contact us](https://www.onlinecloud.dev/) or email [contact@onlinecloud.dev](mailto:contact@onlinecloud.dev) to start the conversation.