Optimizing Token Usage in RAG: Implementing Dynamic Context Pruning

# Optimizing Token Usage in RAG: Implementing Dynamic Context Pruning


*By the engineering team at [OnlineCloud](https://www.onlinecloud.dev/) — a US-based Agentic AI & MCP engineering studio.*

—

## 1. Introduction: The Linear Cost Gap Between Long-Context Retrieval and Agentic ROI

Every RAG pipeline starts the same way. You chunk the corpus, embed it, stand up a vector store, retrieve the top-k, and stuff the results into the prompt. It works in a demo. Then it hits production and the invoice arrives.

Here is the uncomfortable arithmetic. Retrieval cost is roughly constant per query — an ANN search over HNSW or IVF-PQ is sub-linear in corpus size and costs fractions of a cent. **Generation cost, by contrast, scales linearly with the number of tokens you inject into the prompt.** Every additional chunk you retrieve is not just more context — it is more input tokens billed on every single request, plus it inflates the attention quadratic term, plus it raises the probability that the model latches onto an irrelevant passage and hallucinates around it.

The industry’s response over the last two years has been to lean into long context. “Just retrieve 50 chunks, the model has a 200K window.” That is a prototype optimization masquerading as an architecture. A 200K-token prompt on a frontier model at, say, \$3 per million input tokens costs roughly \$0.60 in input alone per request. At 5,000 requests a day, that is \$3,000 a day, or about **\$90K a month**, before you have made a single useful call to a tool. Latency also degrades: a 200K prefill is seconds of time-to-first-token, which kills any interactive or agentic UX.

The agentic ROI problem is sharper still. Agents do not make one call — they make a plan/act/reflect loop with five, ten, sometimes twenty model invocations per task. If each invocation carries a bloated 30K-token context, the cost multiplies through the loop. **The ROI of an agent is bounded above by its per-request token budget.** You cannot out-prompt a bad retrieval strategy.

The fix is not to retrieve less indiscriminately — it is to retrieve broadly and then *prune aggressively* before generation. This post walks through the four-stage pattern we implement in production: reranker-driven pruning, semantic compression, per-request context budgets, and the accuracy instrumentation that keeps you honest. Then we do the ROI math that shows a 50% monthly bill reduction is not marketing — it is a straightforward consequence of dropping tokens that were never contributing to the answer.

—

## 2. Context Pruning: Using Rerankers to Drop Low-Relevance Chunks Before Generation

The core insight is that **recall and precision are two different jobs and deserve two different stages**. Bi-encoder embedding retrieval (your vector store) is optimized for recall: it is fast and it casts a wide net. But cosine similarity between a query embedding and a chunk embedding is a coarse, symmetric signal. It will happily return a chunk that shares vocabulary with the query but has nothing to do with the answer.

A **cross-encoder reranker** fixes this. Instead of comparing two independently-computed vectors, the cross-encoder feeds the (query, chunk) pair jointly through a transformer and produces a single relevance score. That joint attention is dramatically more accurate at judging relevance — but it is O(k) forward passes, so you only run it on a small candidate set. The pipeline is therefore:

1. **Retrieve wide.** Pull top-50 (or top-100) candidates from the vector store. Cheap, high recall.
2. **Rerank narrow.** Score all candidates with a cross-encoder (Cohere Rerank, `bge-reranker-v2-m3`, or a hosted Bedrock reranker). Sort descending.
3. **Prune by threshold, not just top-k.** Keep chunks above a relevance score cutoff *and* cap the count. Score-threshold pruning is the part teams skip, and it is the part that saves money — on an easy query, it may keep 3 chunks; on a hard multi-hop query, it keeps 12.
4. **Deduplicate.** Near-duplicate chunks (overlapping windows, boilerplate headers) burn tokens for zero information gain. A MinHash or embedding-similarity dedupe pass before generation routinely removes 10–20% of surviving chunks.

The economics are immediate. Reranking 50 candidates costs a fraction of a cent and adds maybe 100–200ms. In exchange, you cut the generation context from 50 chunks to 5–12. If each chunk averages 400 tokens, you go from ~20,000 tokens to ~4,000 tokens — an **80% reduction in prompt size** for a relevance *improvement*. You are simultaneously spending less and answering better, which is the rare quadrant engineering teams should chase relentlessly.

One production caveat: rerankers need monitoring too. Track the score distribution over time. If your retrieval query is drifting (new terminology in the corpus, changed user intent), the reranker will confidently prune good chunks. Log the pruned set alongside the retained set so you can audit what you threw away.

—

## 3. Semantic Compression: Summarizing Retrieved Context to Reduce Prompt Tokens

Pruning removes *whole chunks*. Compression shrinks the chunks you keep. These compose: prune first to get the right 8 chunks, then compress those 8 chunks to fit the token budget.

There are three tiers of semantic compression, in increasing aggressiveness:

**Tier 1 — Extractive compression.** Sentence-level selection. Embed each sentence in a retrieved chunk against the query and keep only sentences above a similarity threshold. This is lossless-ish and cheap; you keep the model’s own words and drop the surrounding narrative. Typical reduction: 40–60%.

**Tier 2 — Abstractive (LLM) compression.** Run a small, cheap model (Haiku-class, GPT-4o-mini-class, or a distilled local model) over the retrieved chunk set with a prompt like: *”Extract only the facts relevant to: {query}. Discard preamble, examples, and redundant restatement. Preserve all numbers, names, dates, and code identifiers verbatim.”* You are trading a cheap-model call for a large reduction in expensive-model input tokens. Typical reduction: 60–80%.

**Tier 3 — Structured compression.** For documents with predictable structure (tables, specs, logs, API references), serialize into a compact schema — key/value pairs, TSV rows, or a typed JSON object — rather than prose. This is the highest-leverage form because it eliminates the language modeling overhead entirely. A 600-token prose paragraph about a product’s specs becomes a 90-token key/value block.

The critical engineering discipline for compression is **preservation of a “must-keep” channel**. Numbers, dates, IDs, error codes, and named entities are exactly what compression can silently corrupt, and exactly what your accuracy eval will catch. We implement this by extracting entities with a deterministic NER/regex pass, compressing the prose, then re-attaching the verbatim facts. Never let an LLM paraphrase a dollar amount.

A second discipline: **compress once, cache forever**. Compression output for a given (chunk, query-cluster) pair is deterministic enough to cache. If 30% of your traffic hits the same query intent (very common in support and e-commerce), you compress each chunk once and reuse the compressed form. This turns a per-request cost into a one-time cost — and the compression model call itself becomes the largest line item you eliminate, not add.

—

## 4. Logic: Implementing ‘Context Budgets’ per-Request to Bound LLM Costs

Compression and pruning are *techniques*. A **context budget** is the *policy* that makes them deterministic and puts a ceiling on spend. Without a budget, token usage is an emergent property of your retrieval — nobody owns it, and it drifts upward as the corpus grows.

A context budget is a per-request allocation, expressed in tokens, that the pipeline must satisfy before it calls the generation model:

“`
ContextBudget:
total_prompt_tokens: 6000 # hard ceiling for the whole prompt
system_prompt: 500 # reserved
conversation_history: 1200 # reserved, recency-weighted
retrieved_context: 3500 # the budget we actually fight over
output_reserve: 800 # leaves room for answer length
per_chunk_max: 600 # no single chunk dominates
min_chunks: 2 # never prune below a floor
“`

The allocator runs a greedy fill that respects these constraints:

“`python
def build_context(query, candidates, budget):
ranked = rerank(query, candidates) # cross-encoder scores
ranked = [c for c in ranked if c.score > MIN_RELEVANCE]
ranked = dedupe(ranked) # MinHash near-dup removal

selected, used = [], 0
for chunk in ranked:
cost = estimate_tokens(chunk)
if cost > budget.per_chunk_max:
chunk = compress(chunk, query, target=budget.per_chunk_max)
cost = estimate_tokens(chunk)
if used + cost > budget.retrieved_context:
break
selected.append(chunk)
used += cost

if len(selected) < budget.min_chunks: selected = selected or top_n(ranked, budget.min_chunks) # floor guard return selected, used ``` Three properties make this work in production. **Hard ceilings, soft preferences.** The budget is enforced *before* the model call, not estimated after. You can pre-compute token counts with the exact tokenizer for the target model (`tiktoken` for OpenAI, the Anthropic tokenizer for Claude, or Bedrock's per-model counts) so there is no surprise overage. Never let a heuristic character-count estimate gate a hard budget. **Budget as a tunable knob.** Because the budget is a config value, you can dial it per route: a low-stakes FAQ endpoint gets a 2,000-token budget, a compliance-heavy endpoint gets 12,000. You can also make it *dynamic* — small budgets for high-volume tier-1 traffic, larger budgets when a confidence signal (reranker score variance) suggests the query is genuinely hard. This is where "dynamic context pruning" earns its name: the budget adapts to the query instead of being a fixed top-k. **Budget observability.** Emit `tokens_requested`, `tokens_after_prune`, `tokens_after_compress`, `tokens_billed`, and `prune_ratio` as metrics on every request. Your dashboard should show prompt tokens as a *controlled variable*, not a mystery. When the average creeps up, you will see exactly which stage failed to hold the line. --- ## 5. Performance: Measuring the Trade-off Between Pruning and Factual Accuracy You cannot prune aggressively in production without an eval harness, because pruning is a recall-destroying operation by definition. The question is never "does it lose information?" — it always does. The question is "does the lost information change the answer?" Build the measurement around three axes: **Accuracy (did the answer stay right?).** Maintain a golden set of 100–300 query/answer pairs with ground-truth citations. Run it on every pipeline change. Track exact-match where possible and judge-scored correctness (an LLM-as-judge with a rubric, validated against human labels on a sample) for open-ended answers. Segment by query type — single-hop lookups should be *unaffected* by aggressive pruning; multi-hop and summarization queries are where you will see regressions. **Attribution (was the answer actually grounded?).** Measure citation precision/recall: of the chunks in the final context, how many were actually referenced, and were the referenced ones the right ones? A drop in citation precision after pruning means you kept the wrong evidence. A drop in citation recall means you pruned a chunk the model needed. **Cost and latency (what did you buy?).** Track `$/request` and p50/p95 time-to-first-token before and after. Pruning plus reranking trades ~150ms of rerank latency for a large prefill reduction — net latency usually *improves* because prefill dominates on long prompts. The empirically observed curve from our deployments looks like this: going from top-50 (unpruned) to a reranked, budget-bounded context of ~8 chunks typically produces **flat-to-positive accuracy** (the reranker is removing distractors that caused errors) while cutting tokens by 70–85%. Push further — top-3, aggressive compression with a small model, sub-1K context budgets — and accuracy starts to bend downward, first on multi-hop queries, then on anything requiring synthesis across documents. The knee of the curve is where the product should live, and it is query-class-specific. Two guardrails keep the trade-off honest. **Ablation on every release:** run the eval with pruning off, with pruning on, and with compression on, and report the delta. **A canary on the golden set before rollout:** if citation recall drops more than a couple of points on any query class, block the deploy. Treat the context budget the way you would treat a database index change — measured, versioned, and reversible. --- ## 6. ROI: How Context Pruning Reduces Monthly Bedrock/OpenAI Bills by 50% Let's do the arithmetic on a realistic mid-size deployment: **5,000 RAG requests/day**, frontier model at \$3.00 per million input tokens, output at \$15.00 per million, average 300 output tokens. **Before — naive top-k retrieval.** - Context: 40 chunks × 450 tokens = **18,000 input tokens/request** - Input cost: 5,000 × 18,000 × \$3.00/1M = **\$270/day** - Output cost: 5,000 × 300 × \$15.00/1M = \$22.50/day - **Total: ~\$292/day → ~\$8,770/month** **After — reranked, compressed, budget-bounded.** - Rerank 50 candidates → keep 8 above threshold → 3,600 tokens - Extractive + structured compression → **2,700 input tokens/request** - Input cost: 5,000 × 2,700 × \$3.00/1M = **\$40.50/day** - Reranker cost: 5,000 × 50 pairs ≈ **\$2.50/day** - Compression (cheap model, 30% cache hit): ≈ **\$3.50/day** - Output cost: unchanged at \$22.50/day - **Total: ~\$69/day → ~\$2,070/month** That is a **76% reduction in input tokens**, and a **~55% reduction in the total monthly bill** — input was 92% of spend and pruning attacked it directly. The reranker and compression model together cost \$6/day and save \$230/day. That is a ~38× return on the pruning infrastructure. The 50% figure is conservative, and it is conservative in two places. First, we assumed output cost is untouched — but shorter, less-distracted contexts often produce *shorter, more direct* answers, shaving output tokens too. Second, we assumed no caching. On repetitive support and e-commerce traffic, prompt caching on the stable system prompt plus cached compression output pushes savings past 60%. The agentic multiplier is where it compounds. An agent that fires 8 model calls per task multiplies every inefficiency by 8 — and multiplies every saving by 8. A pipeline with a 2,700-token context that runs 8 steps costs ~21,600 input tokens per task. The unpruned version costs 144,000. Across a month of agentic traffic, context pruning is the difference between an AI feature that pays for itself and one that quietly eats the operating budget. The strategic point for data engineers: **token spend is a pipeline property, not a model property.** You control it at the retrieval and context-assembly layer, with rerankers, compression, and budgets — long before the request reaches Bedrock or OpenAI. Teams that treat context as a first-class resource with an allocated budget ship agentic systems that stay economically viable as they scale. Teams that treat it as "just pass everything to the model" discover the ceiling the hard way, on the invoice. --- ## Building This in Production? Let's Talk. At **[OnlineCloud](https://www.onlinecloud.dev/)**, we build exactly this class of system — reranker-driven retrieval pipelines, budget-bounded context assembly, per-request cost ceilings, and the eval harnesses that prove pruning did not cost you accuracy. It is the same discipline we apply to every MCP server and agentic workload: instrumented, observable, and measured in dollars per request. We are a US-based studio of senior engineers, we scope it, build it, and operate it, and we handle the AWS/Bedrock FinOps work so your monthly bill stays predictable. If your RAG bill is growing faster than your usage, we can help you find the tokens you are wasting. **→ [Book a free initial consultation at onlinecloud.dev](https://www.onlinecloud.dev/)** — response within 24 hours. Let's chat about how our team can help bring your AI vision to life.