Implementing Context-Aware Prompt Compression: Reducing Token Costs Without Losing Retrieval Fidelity in High-Volume RAG Pipelines
Introduction
Picture this: your RAG pipeline is working beautifully. Retrieval quality is solid, answers are accurate, your team is happy. Then the finance team sends over the monthly API bill and everyone goes quiet. You're pushing 50 million tokens per day through GPT-4o at $5 per million input tokens, and roughly 60% of those tokens are retrieved context chunks that the model could have summarized adequately in a third of the space. The math is uncomfortable.
This is the unglamorous reality of production RAG at scale. The community tends to obsess over retrieval precision — better embeddings, smarter chunking, hybrid search — while treating the prompt construction step as a fixed cost. It isn't. The text you stuff into context windows is compressible, and compressing it intelligently — meaning based on the actual query, not just word count — can cut your input token costs by 40–70% without meaningfully degrading answer quality. The key word is intelligently. Naive truncation destroys retrieval fidelity. Context-aware compression preserves it.
By the end of this lesson, you'll understand the full architecture of a context-aware compression layer: how to implement it as middleware in your existing RAG pipeline, how to select and tune compression strategies per query type, how to measure fidelity so you know when you're cutting too deep, and how to instrument everything so costs and quality stay visible in production.
What you'll learn:
- The theoretical foundations of why context-aware compression works (and why naive approaches fail)
- How to implement extractive, abstractive, and selective compression strategies with working Python code
- How to route queries to different compression levels based on complexity and cost constraints
- How to measure compression fidelity using embedding similarity and LLM-as-judge approaches
- How to instrument a compression layer for production observability, with latency budgets and fallback logic
Prerequisites
This is an expert-level lesson. You should be comfortable with:
You'll need: openai, anthropic, tiktoken, sentence-transformers, numpy, and redis (optional, for caching).
Why Naive Compression Destroys Fidelity
Before writing any code, we need to understand why compression is hard. The intuition is simple: remove words that don't help answer the question. The implementation is where people go wrong.
The most common naive approach is pure length-based truncation: take the top-k retrieved chunks and cut anything beyond a token budget. This sounds reasonable until you realize that retrieval ranking is imperfect. The most syntactically dense but semantically critical sentence in a chunk is just as likely to appear in positions 200–250 as positions 0–50. Truncating by position is essentially random with respect to relevance.
The second naive approach is keyword filtering: keep sentences that contain terms from the query. This is better, but it breaks on paraphrase and synonym relationships. A user asking about "revenue recognition policies" may need context that says "we record income at the point of delivery" — no keyword overlap, maximum relevance.
The third naive approach is uniform summarization: compress every chunk by the same ratio using the same prompt. This fails because different chunk types have radically different information density. A legal clause is load-bearing at every word. A narrative paragraph in a company blog post can lose 80% of its tokens with no information loss. Treating them identically either over-compresses the legal clause (destroys fidelity) or under-compresses the blog post (wastes budget).
Key insight: Context-aware compression means making compression decisions relative to the query, not relative to absolute token counts. The same chunk should be compressed differently depending on what question is being asked about it.
The goal is a compression function compress(chunk, query) → compressed_chunk where the output retains exactly and only the information in chunk that is relevant to answering query. Everything else is expendable.
The Architecture: A Compression Middleware Layer
Think of your compression layer as sitting between retrieval and generation, transforming the raw retrieved context before it hits the main LLM call. Here's the data flow:
User Query
↓
[Retriever] → raw_chunks (potentially large)
↓
[Compression Router] → decides strategy per chunk
↓
[Compression Executor] → applies strategy, respects token budget
↓
[Fidelity Checker] → validates compressed chunks haven't lost key info
↓
[LLM Generator] → works with compressed, high-density context
The compression router is the intelligence here. It looks at the query type, the chunk content type, your token budget, and the cost/latency constraints of the request to decide which compression strategy to use. Let's build each component.
First, the scaffolding:
import tiktoken
import numpy as np
from dataclasses import dataclass, field
from typing import Optional
from enum import Enum
enc = tiktoken.get_encoding("cl100k_base")
def count_tokens(text: str) -> int:
return len(enc.encode(text))
class CompressionStrategy(Enum):
PASSTHROUGH = "passthrough" # No compression — chunk fits budget as-is
EXTRACTIVE = "extractive" # Sentence-level scoring and filtering
ABSTRACTIVE = "abstractive" # LLM-based rewriting/summarization
SELECTIVE = "selective" # Keep only query-relevant spans verbatim
@dataclass
class RetrievedChunk:
content: str
source_id: str
retrieval_score: float
chunk_type: str = "generic" # "legal", "technical", "narrative", "tabular"
token_count: int = field(init=False)
def __post_init__(self):
self.token_count = count_tokens(self.content)
@dataclass
class CompressedChunk:
content: str
source_id: str
original_token_count: int
compressed_token_count: int
strategy_used: CompressionStrategy
fidelity_score: Optional[float] = None
@property
def compression_ratio(self) -> float:
if self.original_token_count == 0:
return 1.0
return self.compressed_token_count / self.original_token_count
Extractive compression keeps whole sentences but filters out those with low relevance to the query. It's fast (no LLM call required), deterministic, and safe — you're never inventing text, only removing it.
The scoring function is the key decision. We'll use a hybrid of:
- Embedding cosine similarity between sentence and query
- Position weighting (first and last sentences of a chunk often contain topic sentences and conclusions)
- Keyword overlap as a cheap signal for obvious relevance
from sentence_transformers import SentenceTransformer
import re
_embedder = SentenceTransformer("all-MiniLM-L6-v2") # Fast, good enough for scoring
def split_sentences(text: str) -> list[str]:
"""Sentence splitter that handles common edge cases in technical docs."""
# Preserve newlines as sentence boundaries (important for structured docs)
text = re.sub(r'\n{2,}', ' <PARA> ', text)
sentences = re.split(r'(?<=[.!?])\s+', text)
return [s.replace(' <PARA> ', ' ').strip() for s in sentences if s.strip()]
def score_sentences(
sentences: list[str],
query: str,
query_keywords: set[str]
) -> list[float]:
if not sentences:
return []
# Batch embed for efficiency
all_texts = sentences + [query]
embeddings = _embedder.encode(all_texts, normalize_embeddings=True)
sentence_embeddings = embeddings[:-1]
query_embedding = embeddings[-1]
# Cosine similarity (embeddings are already normalized)
semantic_scores = np.dot(sentence_embeddings, query_embedding)
# Keyword overlap score
query_kws = {kw.lower() for kw in query_keywords}
keyword_scores = []
for sent in sentences:
words = set(sent.lower().split())
if query_kws:
overlap = len(words & query_kws) / len(query_kws)
else:
overlap = 0.0
keyword_scores.append(overlap)
# Position weighting — first 15% and last 15% of sentences get a bonus
n = len(sentences)
position_scores = []
for i in range(n):
relative_pos = i / max(n - 1, 1)
# U-shaped: high at 0 and 1, low in middle
pos_score = 1.0 - (2 * abs(relative_pos - 0.5) - 0.0) * 0.3
position_scores.append(pos_score)
# Combine: semantic is dominant, keywords and position are tiebreakers
final_scores = []
for sem, kw, pos in zip(semantic_scores, keyword_scores, position_scores):
combined = 0.70 * float(sem) + 0.20 * kw + 0.10 * pos
final_scores.append(combined)
return final_scores
def extractive_compress(
chunk: RetrievedChunk,
query: str,
query_keywords: set[str],
target_ratio: float = 0.5,
min_sentences: int = 2
) -> str:
sentences = split_sentences(chunk.content)
if len(sentences) <= min_sentences:
return chunk.content # Nothing to compress
scores = score_sentences(sentences, query, query_keywords)
target_count = max(min_sentences, int(len(sentences) * target_ratio))
# Rank by score, but preserve original order of kept sentences
ranked_indices = sorted(range(len(scores)), key=lambda i: scores[i], reverse=True)
kept_indices = sorted(ranked_indices[:target_count])
return " ".join(sentences[i] for i in kept_indices)
Tip: For technical documentation with dense tables or code snippets, skip sentence-level splitting entirely and treat each line or code block as a unit. Regex sentence splitters will mangle structured content badly.
One subtlety here: target_ratio shouldn't be a fixed constant. It should vary by chunk_type. Narrative text can handle 0.35. Regulatory or legal text should never drop below 0.70 without human review. We'll wire this up in the router.
Strategy 2: Abstractive Compression
Abstractive compression uses an LLM to rewrite the chunk in fewer tokens, focusing only on what's relevant to the query. It's more powerful than extractive — it can synthesize information across sentences — but it costs money and adds latency.
The quality of abstractive compression depends almost entirely on the prompt. Here's what works:
from openai import AsyncOpenAI
_client = AsyncOpenAI()
ABSTRACTIVE_COMPRESSION_PROMPT = """\
You are a precision context extractor. Your job is to compress the following document chunk \
into a dense summary that retains ALL information needed to answer the query, and ONLY that information.
Rules:
- Preserve exact numbers, dates, names, and technical terms verbatim
- If a sentence is partially relevant, keep the relevant clause, not the whole sentence
- Do not add interpretation or outside knowledge
- Write in third person, present tense, maintaining the factual register of the source
- Target length: {target_tokens} tokens or fewer
Query: {query}
Document chunk:
{chunk_content}
Compressed context:"""
async def abstractive_compress(
chunk: RetrievedChunk,
query: str,
target_token_count: int,
model: str = "gpt-4o-mini" # Cheap model for compression — save the big model for generation
) -> str:
prompt = ABSTRACTIVE_COMPRESSION_PROMPT.format(
target_tokens=target_token_count,
query=query,
chunk_content=chunk.content
)
response = await _client.chat.completions.create(
model=model,
messages=[{"role": "user", "content": prompt}],
max_tokens=target_token_count + 50, # Small buffer
temperature=0.0 # Determinism matters here
)
compressed = response.choices[0].message.content.strip()
# Safety check: if the model returned more tokens than target, fall back to extractive
if count_tokens(compressed) > target_token_count * 1.15:
return None # Signal to caller to use fallback
return compressed
Warning: Never use your primary generation model for abstractive compression. You're adding an LLM call before the main LLM call, so model choice here has a multiplicative effect on costs. gpt-4o-mini or claude-haiku handles this task well at a fraction of the cost. See Implementing LLM Router Architecture: Dynamically Selecting Models by Cost, Latency, and Task Complexity in Production for patterns on routing to the right model by task.
A critical concern with abstractive compression is hallucination under compression. An LLM rewriting a chunk can subtly alter numbers, change conditional language ("may require" → "requires"), or smooth over important caveats. This is why the fidelity checker we'll build later is non-negotiable for abstractive strategies.
Strategy 3: Selective (Span) Compression
Selective compression is the precision instrument. Instead of summarizing or filtering sentences, it identifies contiguous spans of text — passages, clauses, even individual sentences — that directly address the query, and returns them verbatim. No rewriting, no inference, just surgical extraction.
This is ideal for high-stakes domains (legal, medical, financial) where the exact wording matters, but where chunks contain a lot of surrounding boilerplate.
def selective_compress(
chunk: RetrievedChunk,
query: str,
query_keywords: set[str],
context_window: int = 1, # Sentences before/after each relevant sentence to keep
relevance_threshold: float = 0.45
) -> str:
sentences = split_sentences(chunk.content)
if not sentences:
return chunk.content
scores = score_sentences(sentences, query, query_keywords)
# Find all sentences above threshold
relevant_indices = {
i for i, score in enumerate(scores)
if score >= relevance_threshold
}
# Expand each relevant sentence by context_window
expanded_indices = set()
for idx in relevant_indices:
for offset in range(-context_window, context_window + 1):
neighbor = idx + offset
if 0 <= neighbor < len(sentences):
expanded_indices.add(neighbor)
if not expanded_indices:
# Nothing clears the threshold — fall back to top-3 sentences
top_indices = sorted(
range(len(scores)),
key=lambda i: scores[i],
reverse=True
)[:3]
expanded_indices = set(top_indices)
kept = [sentences[i] for i in sorted(expanded_indices)]
return " ".join(kept)
The context_window parameter is important: isolated sentences often lose meaning without their neighbors. A sentence like "This exception applies only in the case described above" is useless without the preceding sentence. Setting context_window=1 is usually enough, but for dense legal or technical text, context_window=2 is safer.
The Compression Router
Now we assemble the router that decides which strategy to apply to each chunk. The routing logic considers:
- Token budget remaining — how much space do we have left?
- Chunk type — is this legal text? Narrative? Code?
- Retrieval score — high-scoring chunks get less aggressive compression
- Query complexity — a simple factual lookup warrants different treatment than a multi-hop reasoning question
@dataclass
class CompressionConfig:
total_token_budget: int = 4000
abstractive_model: str = "gpt-4o-mini"
enable_abstractive: bool = True
fidelity_threshold: float = 0.80
chunk_type_ratios: dict = field(default_factory=lambda: {
"legal": 0.80,
"financial": 0.75,
"technical": 0.60,
"narrative": 0.40,
"generic": 0.55,
})
class CompressionRouter:
def __init__(self, config: CompressionConfig):
self.config = config
def decide_strategy(
self,
chunk: RetrievedChunk,
query: str,
remaining_budget: int,
query_is_complex: bool = False
) -> tuple[CompressionStrategy, int]:
"""Returns (strategy, target_token_count)."""
# If chunk already fits in the remaining budget, no compression needed
if chunk.token_count <= remaining_budget:
return CompressionStrategy.PASSTHROUGH, chunk.token_count
# Determine target token count based on chunk type
ratio = self.config.chunk_type_ratios.get(chunk.chunk_type, 0.55)
# High retrieval score = be more conservative (less compression)
if chunk.retrieval_score > 0.85:
ratio = min(ratio * 1.2, 0.90)
target_tokens = max(
int(chunk.token_count * ratio),
min(remaining_budget, 150) # Never compress below 150 tokens
)
# Route to strategy based on chunk type and query complexity
if chunk.chunk_type in ("legal", "financial"):
# High-stakes: use selective (verbatim span extraction)
return CompressionStrategy.SELECTIVE, target_tokens
if query_is_complex and self.config.enable_abstractive:
# Complex queries benefit from synthesis
return CompressionStrategy.ABSTRACTIVE, target_tokens
if chunk.token_count > 800:
# Long generic chunks: extractive is efficient
return CompressionStrategy.EXTRACTIVE, target_tokens
# Short chunks that still don't fit: extractive
return CompressionStrategy.EXTRACTIVE, target_tokens
Note: The query_is_complex flag can be set by a cheap pre-classification step — a regex check for question words indicating multi-hop reasoning ("how does X compare to Y", "what changed between X and Y", "explain the relationship between..."). You don't need another LLM call for this; a simple heuristic classifier works fine and keeps latency low.
Fidelity Measurement: Knowing When You've Cut Too Deep
Compression without measurement is just hope. You need a systematic way to verify that a compressed chunk still contains the information needed to answer the query. We use two complementary approaches:
Approach 1: Embedding Similarity
Compare the embedding of the compressed chunk to the embedding of the original chunk. A cosine similarity below ~0.75 usually indicates significant information loss.
from sentence_transformers import SentenceTransformer
_fidelity_embedder = SentenceTransformer("all-MiniLM-L6-v2")
def compute_embedding_fidelity(
original: str,
compressed: str
) -> float:
"""Returns cosine similarity between original and compressed embeddings."""
embeddings = _fidelity_embedder.encode(
[original, compressed],
normalize_embeddings=True
)
return float(np.dot(embeddings[0], embeddings[1]))
Approach 2: LLM-as-Judge (for High-Stakes Chunks)
For chunks where you're using abstractive compression on important content, a lightweight LLM judge check adds a meaningful safety net:
FIDELITY_CHECK_PROMPT = """\
You are evaluating whether a compressed document preserves the key facts relevant to a query.
Query: {query}
Original chunk:
{original}
Compressed version:
{compressed}
Does the compressed version preserve all factual information from the original that is needed \
to answer the query? Reply with only a JSON object: {{"preserved": true/false, "missing": ["list of missing facts"]}}
"""
async def llm_fidelity_check(
original: str,
compressed: str,
query: str,
model: str = "gpt-4o-mini"
) -> dict:
prompt = FIDELITY_CHECK_PROMPT.format(
query=query,
original=original,
compressed=compressed
)
response = await _client.chat.completions.create(
model=model,
messages=[{"role": "user", "content": prompt}],
temperature=0.0,
response_format={"type": "json_object"}
)
import json
return json.loads(response.choices[0].message.content)
Warning: LLM-as-judge fidelity checks add latency and cost. Don't run them on every chunk in every request. Reserve them for: (1) abstractively compressed chunks, (2) chunks from high-stakes source types (legal, medical, financial), (3) chunks where embedding fidelity dropped below your threshold. Gate them behind the embedding check as a cheaper first filter.
Assembling the Full Compression Pipeline
Now let's wire everything together into a production-ready pipeline:
import asyncio
from typing import Optional
async def compress_chunk(
chunk: RetrievedChunk,
query: str,
query_keywords: set[str],
remaining_budget: int,
router: CompressionRouter,
query_is_complex: bool = False,
run_fidelity_check: bool = True
) -> CompressedChunk:
strategy, target_tokens = router.decide_strategy(
chunk, query, remaining_budget, query_is_complex
)
original_content = chunk.content
compressed_content = None
if strategy == CompressionStrategy.PASSTHROUGH:
compressed_content = original_content
elif strategy == CompressionStrategy.EXTRACTIVE:
ratio = target_tokens / chunk.token_count
compressed_content = extractive_compress(
chunk, query, query_keywords, target_ratio=ratio
)
elif strategy == CompressionStrategy.SELECTIVE:
compressed_content = selective_compress(
chunk, query, query_keywords
)
# If selective overshot the budget, trim further with extractive
if count_tokens(compressed_content) > target_tokens * 1.1:
temp_chunk = RetrievedChunk(
content=compressed_content,
source_id=chunk.source_id,
retrieval_score=chunk.retrieval_score,
chunk_type=chunk.chunk_type
)
ratio = target_tokens / temp_chunk.token_count
compressed_content = extractive_compress(
temp_chunk, query, query_keywords, target_ratio=ratio
)
elif strategy == CompressionStrategy.ABSTRACTIVE:
compressed_content = await abstractive_compress(
chunk, query, target_tokens
)
# Fallback to extractive if abstractive fails or overshoots
if compressed_content is None:
ratio = target_tokens / chunk.token_count
compressed_content = extractive_compress(
chunk, query, query_keywords, target_ratio=ratio
)
strategy = CompressionStrategy.EXTRACTIVE
# Fidelity check
fidelity_score = None
if run_fidelity_check and strategy != CompressionStrategy.PASSTHROUGH:
fidelity_score = compute_embedding_fidelity(original_content, compressed_content)
# If fidelity is too low, fall back to passthrough (budget overrun is preferable
# to losing information in high-value chunks)
if fidelity_score < router.config.fidelity_threshold and chunk.retrieval_score > 0.80:
compressed_content = original_content
strategy = CompressionStrategy.PASSTHROUGH
fidelity_score = 1.0
final_token_count = count_tokens(compressed_content)
return CompressedChunk(
content=compressed_content,
source_id=chunk.source_id,
original_token_count=chunk.token_count,
compressed_token_count=final_token_count,
strategy_used=strategy,
fidelity_score=fidelity_score
)
async def compress_context(
chunks: list[RetrievedChunk],
query: str,
query_keywords: set[str],
config: CompressionConfig,
query_is_complex: bool = False
) -> tuple[list[CompressedChunk], dict]:
"""
Compress a list of retrieved chunks to fit within the total token budget.
Returns compressed chunks and a stats dict for observability.
"""
router = CompressionRouter(config)
results = []
tokens_used = 0
# Sort by retrieval score descending so highest-value chunks get first pick of budget
sorted_chunks = sorted(chunks, key=lambda c: c.retrieval_score, reverse=True)
# Run compressions concurrently for efficiency
tasks = []
for chunk in sorted_chunks:
remaining = config.total_token_budget - tokens_used
if remaining <= 0:
break
tasks.append(
compress_chunk(
chunk=chunk,
query=query,
query_keywords=query_keywords,
remaining_budget=remaining,
router=router,
query_is_complex=query_is_complex,
run_fidelity_check=True
)
)
# Estimate token usage assuming compression succeeds
estimated_ratio = config.chunk_type_ratios.get(chunk.chunk_type, 0.55)
tokens_used += int(chunk.token_count * estimated_ratio)
compressed_chunks = await asyncio.gather(*tasks)
# Compute stats
total_original = sum(c.original_token_count for c in compressed_chunks)
total_compressed = sum(c.compressed_token_count for c in compressed_chunks)
strategy_breakdown = {}
for cc in compressed_chunks:
strategy_breakdown[cc.strategy_used.value] = \
strategy_breakdown.get(cc.strategy_used.value, 0) + 1
stats = {
"total_original_tokens": total_original,
"total_compressed_tokens": total_compressed,
"overall_compression_ratio": total_compressed / max(total_original, 1),
"token_savings": total_original - total_compressed,
"strategy_breakdown": strategy_breakdown,
"avg_fidelity": np.mean([
c.fidelity_score for c in compressed_chunks
if c.fidelity_score is not None
]).item() if any(c.fidelity_score is not None for c in compressed_chunks) else None
}
return list(compressed_chunks), stats
Caching Compressed Chunks
One often-overlooked optimization: if the same chunk gets retrieved for similar queries, you don't need to recompress it every time. You can cache compressed versions keyed on (source_id, query_embedding_hash).
This pairs well with semantic deduplication — if two queries are embedding-similar enough that they'd retrieve the same chunks, there's a good chance the compressed representation is reusable. For the full caching architecture, see Implementing LLM Response Caching with Redis: Semantic Deduplication, TTL Strategies, and Cache Invalidation Patterns.
A lightweight version without Redis:
import hashlib
from functools import lru_cache
def _make_cache_key(source_id: str, query: str, strategy: str) -> str:
content = f"{source_id}::{query}::{strategy}"
return hashlib.sha256(content.encode()).hexdigest()
# In-process cache — use Redis for multi-instance deployments
_compression_cache: dict[str, CompressedChunk] = {}
def get_cached_compression(
source_id: str,
query: str,
strategy: str
) -> Optional[CompressedChunk]:
key = _make_cache_key(source_id, query, strategy)
return _compression_cache.get(key)
def cache_compression(
source_id: str,
query: str,
strategy: str,
result: CompressedChunk
) -> None:
key = _make_cache_key(source_id, query, strategy)
_compression_cache[key] = result
Tip: For semantic query matching in the cache key, use a rounded embedding vector (quantized to 2 decimal places) rather than the raw query string. Two queries that are phrased differently but mean the same thing will then hit the same cache entry. This requires a quick embedding call upfront but avoids redundant compression work across thousands of near-duplicate queries.
Production Observability: Instrumenting Your Compression Layer
A compression layer that isn't monitored is a liability. You need to track three things in production:
- Cost savings — are you actually saving money?
- Fidelity drift — is answer quality degrading over time?
- Latency impact — is the compression step adding unacceptable latency?
import time
import logging
from contextlib import asynccontextmanager
logger = logging.getLogger("compression_pipeline")
@asynccontextmanager
async def compression_span(query_id: str, query: str):
"""Context manager that wraps a compression operation with timing and logging."""
start = time.perf_counter()
span_data = {"query_id": query_id, "query_preview": query[:80]}
try:
yield span_data
finally:
elapsed = time.perf_counter() - start
span_data["compression_latency_ms"] = round(elapsed * 1000, 2)
logger.info("compression_complete", extra=span_data)
async def compress_context_instrumented(
chunks: list[RetrievedChunk],
query: str,
query_keywords: set[str],
config: CompressionConfig,
query_id: str,
query_is_complex: bool = False
) -> tuple[list[CompressedChunk], dict]:
async with compression_span(query_id, query) as span:
compressed, stats = await compress_context(
chunks=chunks,
query=query,
query_keywords=query_keywords,
config=config,
query_is_complex=query_is_complex
)
# Enrich span with compression metrics
span.update(stats)
# Alert if fidelity drops below threshold
avg_fidelity = stats.get("avg_fidelity")
if avg_fidelity is not None and avg_fidelity < config.fidelity_threshold:
logger.warning(
"low_fidelity_compression",
extra={
"query_id": query_id,
"avg_fidelity": avg_fidelity,
"threshold": config.fidelity_threshold
}
)
# Alert if compression ratio is unexpectedly high (possible over-compression)
ratio = stats["overall_compression_ratio"]
if ratio < 0.30:
logger.warning(
"aggressive_compression_detected",
extra={"query_id": query_id, "ratio": ratio}
)
return compressed, stats
For full tracing infrastructure that integrates with LLM observability platforms, see Implementing LLM Observability: Tracing, Logging, and Monitoring Requests in Production.
Token Budget Allocation: The Strategic Layer
The token budget isn't just a single number — it's a resource allocation problem. In a full RAG pipeline, your context window budget must cover:
- System prompt
- Conversation history (if multi-turn)
- Retrieved and compressed chunks
- Generation headroom for the model's output
Here's a practical budget allocator:
@dataclass
class TokenBudgetAllocation:
system_prompt_tokens: int
conversation_history_tokens: int
context_token_budget: int
generation_headroom: int
total_model_context: int
def allocate_token_budget(
system_prompt: str,
conversation_history: list[dict],
model_context_limit: int = 128_000,
generation_headroom: int = 2048,
context_allocation_ratio: float = 0.65 # Give 65% of remaining to context
) -> TokenBudgetAllocation:
system_tokens = count_tokens(system_prompt)
history_tokens = sum(
count_tokens(msg.get("content", ""))
for msg in conversation_history
)
fixed_overhead = system_tokens + history_tokens + generation_headroom
if fixed_overhead >= model_context_limit:
raise ValueError(
f"Fixed overhead ({fixed_overhead} tokens) exceeds model limit "
f"({model_context_limit}). Trim conversation history."
)
available = model_context_limit - fixed_overhead
context_budget = int(available * context_allocation_ratio)
return TokenBudgetAllocation(
system_prompt_tokens=system_tokens,
conversation_history_tokens=history_tokens,
context_token_budget=context_budget,
generation_headroom=generation_headroom,
total_model_context=model_context_limit
)
Key insight: In multi-turn RAG applications, conversation history grows with every turn and competes with the context budget. If you're not compressing both history and retrieved chunks, you may find that by turn 10 your context budget for retrieved chunks has dropped to nearly nothing. For strategies on managing growing conversation state, see Implementing Conversational Memory: Managing Context Windows and Chat History at Scale.
Evaluating End-to-End Fidelity in Practice
Embedding similarity tells you whether compression preserved the semantic content. But for RAG, the question you ultimately care about is: did the answer quality degrade? That requires end-to-end evaluation.
Here's a framework for running offline evaluations of your compression settings:
async def evaluate_compression_impact(
test_cases: list[dict], # Each: {"query": str, "chunks": list, "ground_truth_answer": str}
config: CompressionConfig,
generation_model: str = "gpt-4o"
) -> dict:
"""
test_cases: list of dicts with query, chunks (pre-retrieved), and ground truth answer
Returns evaluation metrics comparing uncompressed vs compressed pipelines.
"""
results = []
for case in test_cases:
query = case["query"]
chunks = case["chunks"]
ground_truth = case["ground_truth_answer"]
# Get keywords (simple approach: non-stopword query tokens)
stop_words = {"the", "a", "an", "is", "are", "was", "were", "in", "of", "and", "or"}
keywords = {
w.lower() for w in query.split()
if w.lower() not in stop_words and len(w) > 2
}
# Run compression
compressed_chunks, stats = await compress_context(
chunks=chunks,
query=query,
query_keywords=keywords,
config=config,
)
# Build context strings
uncompressed_context = "\n\n".join(c.content for c in chunks)
compressed_context = "\n\n".join(c.content for c in compressed_chunks)
# Generate answers from both
async def get_answer(context: str) -> str:
resp = await _client.chat.completions.create(
model=generation_model,
messages=[
{"role": "system", "content": "Answer the query using only the provided context."},
{"role": "user", "content": f"Context:\n{context}\n\nQuery: {query}"}
],
temperature=0.0
)
return resp.choices[0].message.content
uncompressed_answer, compressed_answer = await asyncio.gather(
get_answer(uncompressed_context),
get_answer(compressed_context)
)
# Score both answers against ground truth using embedding similarity
answer_embeddings = _fidelity_embedder.encode(
[ground_truth, uncompressed_answer, compressed_answer],
normalize_embeddings=True
)
uncompressed_score = float(np.dot(answer_embeddings[0], answer_embeddings[1]))
compressed_score = float(np.dot(answer_embeddings[0], answer_embeddings[2]))
results.append({
"query": query,
"uncompressed_answer_score": uncompressed_score,
"compressed_answer_score": compressed_score,
"fidelity_delta": compressed_score - uncompressed_score,
"compression_ratio": stats["overall_compression_ratio"],
"token_savings": stats["token_savings"]
})
avg_fidelity_delta = np.mean([r["fidelity_delta"] for r in results])
avg_token_savings = np.mean([r["token_savings"] for r in results])
pct_within_5pct = np.mean([abs(r["fidelity_delta"]) < 0.05 for r in results])
return {
"avg_fidelity_delta": float(avg_fidelity_delta),
"avg_token_savings": float(avg_token_savings),
"pct_answers_within_5pct_quality": float(pct_within_5pct),
"per_case_results": results
}
A healthy compression configuration will show avg_fidelity_delta between -0.03 and +0.02 (sometimes compression actually helps by removing noise) and pct_answers_within_5pct_quality above 0.85. If you're below those thresholds, your compression ratios are too aggressive for your content types.
Hands-On Exercise
Build a complete compression pipeline for a corpus of legal contracts and test it against a benchmark query set. Here's the structured exercise:
Setup:
Download 10–20 public contracts from SEC EDGAR (search for "exhibit 10" filings). Split them into chunks of approximately 500 tokens using semantic chunking — the Chunking Strategies for RAG: How to Split Documents by Size, Sentence, and Semantic Meaning article covers the mechanics. Label each chunk with chunk_type="legal".
Task 1: Baseline measurement
Build a small query set (20 questions) that require answering from the contracts: payment terms, termination clauses, liability limits, governing law. For each, manually retrieve 5 relevant chunks and record the total token counts. This is your baseline cost.
Task 2: Implement and tune compression
Run the compression pipeline on your chunks with:
CompressionConfig(total_token_budget=2000, enable_abstractive=False)CompressionConfig(total_token_budget=2000, enable_abstractive=True)CompressionConfig(total_token_budget=3000, enable_abstractive=True)
Record token savings and fidelity scores for each configuration.
Task 3: End-to-end answer quality evaluation
For each configuration, generate answers to your 20 questions using GPT-4o-mini. Score the answers against ground truth (which you provide manually for these 20 questions). Plot: compression ratio vs. answer quality degradation.
Expected finding: You should see that legal chunks with selective strategy and moderate targets (0.75–0.80 ratio) maintain answer quality within 3–5% of uncompressed, while saving 20–25% of tokens. Abstractive compression without sufficient fidelity checking will show quality drops on clauses with conditional language.
Task 4: Add the reranking integration
Before compression, add a reranking step using the patterns from Building a Reranking Layer for RAG: Improving Retrieval Precision with Cross-Encoders and LLM-Based Scoring. Measure whether reranking before compression improves or degrades fidelity scores. (Hypothesis: it should improve them, because reranking surfaces truly relevant chunks that compression then preserves more conservatively.)
Common Mistakes & Troubleshooting
Mistake 1: Compressing all chunks to the same ratio
Symptom: Consistently low fidelity scores on a specific source type even though aggregate metrics look fine.
Fix: Always segment by chunk_type in your router config. Run per-type fidelity analysis and tune chunk_type_ratios separately for each type.
Mistake 2: Running abstractive compression synchronously on every chunk
Symptom: Compression layer adds 800ms–1.2s of latency per request, making the total pipeline unacceptably slow.
Fix: Use asyncio.gather() to parallelize compression across chunks. For very high-throughput applications, consider moving abstractive compression to a background worker with a cache, so frequently-retrieved chunks arrive pre-compressed. See Orchestrating Parallel LLM Calls: Batching, Concurrency, and Async Patterns for High-Throughput Production Pipelines for the async patterns.
Mistake 3: Ignoring the interaction between compression and citation systems
Symptom: Your citation attribution system (linking claims back to source documents) starts pointing to passages that don't exist verbatim in the compressed context.
Fix: Always preserve source_id mapping between compressed and original chunks. When using extractive or selective strategies, also track exact character offsets of kept sentences so the citation system can map back to the original document. See Building a Citation and Source Attribution System for RAG for the full attribution architecture.
Mistake 4: Setting the fidelity threshold too high globally
Symptom: The fallback-to-passthrough logic triggers for nearly every chunk, resulting in near-zero actual token savings.
Fix: Fidelity thresholds should be per-chunk-type, not global. A fidelity_threshold=0.80 is appropriate for legal chunks. For narrative marketing content, 0.65 is perfectly acceptable. The current implementation uses a single threshold — extend the CompressionConfig to hold a dict of per-type thresholds.
Mistake 5: Not accounting for compressor latency in overall SLA budgets
Symptom: P99 latency SLAs are being violated even though the generation step is fast.
Fix: Track compression latency separately in your observability layer. Profile per-strategy latencies: extractive is typically <5ms, selective <10ms, abstractive 200–600ms per chunk. For requests where you have less than 300ms total budget, disable abstractive compression via a config flag and fall back to extractive.
Mistake 6: Using a single compression config across query types
Symptom: Simple factual lookups ("what is the contract effective date?") are being treated identically to complex analytical queries ("compare the liability structures across all three contracts"). The former needs almost no context; the latter needs rich, well-preserved context.
Fix: Classify query complexity before routing to the compressor. A simple heuristic: queries with comparison words, temporal scope ("between X and Y"), or multi-entity references (count of named entities > 2) are complex. Route complex queries to a higher total_token_budget and more conservative compression ratios.
Summary & Next Steps
You've built a complete context-aware compression layer from the ground up: scoring functions, three distinct compression strategies matched to content types, a routing engine that integrates query complexity and retrieval scores, fidelity measurement at both the embedding and LLM-as-judge level, and production instrumentation.
The key principles to carry forward:
- Compression is a function of both content and query. The same chunk should be compressed differently depending on what's being asked.
- Different content types have different tolerance for compression. Legal and financial text is load-bearing at the word level; narrative text is not.
- Fidelity measurement is not optional. Without it, you don't know whether your token savings are coming at the cost of answer quality.
- Abstractive compression is powerful but needs guardrails. Use cheap models, run fidelity checks, and always have an extractive fallback.
- The compression layer compounds with other optimizations. Combining better chunking, reranking before compression, and semantic caching can push token savings to 60–70% with minimal quality loss.
Where to go from here:
