Cutting RAG Token Bloat: Compression Before Retrieval in Production Pipelines

#database #ai #python #backend

The Token Tax on RAG Pipelines

When you're running a retrieval-augmented generation system at scale, the real bottleneck isn't your embedding model or your vector index—it's the token overhead between retrieval and generation. I learned this the hard way after our agent memory service started hitting 200K context windows on routine queries, with 60% of those tokens coming from redundant JSON metadata and uncompressed tool outputs.

The architecture was standard: pgvector for semantic search, a Python orchestration layer, and Claude Sonnet routed through a multi-provider gateway at https://b-lost.com/v1 for fallback resilience. The problem surfaced during high-concurrency runs—our median token count per request jumped from 45K to 180K within two weeks of production load, purely from accumulated retrieval chunks and agent step logs.

Compression as a Pre-LLM Gate

That's where headroomlabs-ai/headroom became critical. Instead of throwing raw retrieval results and JSON blobs directly into the LLM context, we inserted a compression stage right after vector search and before prompt assembly. The library handles three core reduction patterns:

Semantic deduplication across retrieved chunks (cosine similarity threshold to drop near-identical paragraphs)
JSON structure flattening (strips redundant keys, inlines nested objects where safe)
Log summarization for agent tool outputs (replaces verbose stack traces with root-cause distillations)

Here's the integration point in our retrieval pipeline:

import psycopg
from pgvector.psycopg import register_vector
from headroom import compress_chunks

def hybrid_search_with_compression(query_embedding, query_text, top_k=10):
    conn = psycopg.connect("dbname=agent_memory")
    register_vector(conn)

    # Vector search with HNSW index
    vector_results = conn.execute(
        """
        SELECT id, content, embedding <=> %s AS distance
        FROM memory_chunks
        WHERE metadata->>'session_id' = %s
        ORDER BY embedding <=> %s
        LIMIT %s
        """,
        (query_embedding, session_id, query_embedding, top_k)
    ).fetchall()

    # Keyword search for hybrid fusion
    keyword_results = conn.execute(
        """
        SELECT id, content, ts_rank(search_vector, plainto_tsquery('english', %s)) AS rank
        FROM memory_chunks
        WHERE search_vector @@ plainto_tsquery('english', %s)
        AND metadata->>'session_id' = %s
        ORDER BY rank DESC
        LIMIT %s
        """,
        (query_text, query_text, session_id, top_k)
    ).fetchall()

    # Reciprocal rank fusion
    fused = reciprocal_rank_fusion(vector_results, keyword_results, k=60)

    # Compression gate before LLM ingestion
    compressed = compress_chunks(
        [r['content'] for r in fused],
        target_ratio=0.4,  # Aim for 60% reduction
        preserve_code_blocks=True
    )

    return compressed

def reciprocal_rank_fusion(vector_results, keyword_results, k=60):
    scores = {}
    for rank, row in enumerate(vector_results, 1):
        scores[row['id']] = scores.get(row['id'], 0) + 1 / (k + rank)
    for rank, row in enumerate(keyword_results, 1):
        scores[row['id']] = scores.get(row['id'], 0) + 1 / (k + rank)

    ranked = sorted(scores.items(), key=lambda x: x[1], reverse=True)
    return [{'id': id, 'content': get_content(id), 'score': score} 
            for id, score in ranked[:10]]

The HNSW Lock Trap

One production trap we hit: during bulk embedding insertion (10K+ vectors/minute), PostgreSQL's HNSW index locking caused severe write contention. Our P95 insert latency spiked to 3.2 seconds. The fix was counterintuitive—disable the index during batch writes, then rebuild:

-- Before bulk insert
DROP INDEX IF EXISTS memory_chunks_embedding_idx;

-- Batch insert
COPY memory_chunks (id, content, embedding, metadata) FROM '/tmp/batch.csv';

-- Rebuild index (takes 2-4 min for 1M vectors)
CREATE INDEX memory_chunks_embedding_idx ON memory_chunks 
USING hnsw (embedding vector_cosine_ops) 
WITH (m = 16, ef_construction = 64);

The m=16 and ef_construction=64 parameters are critical. Higher values improve recall but double index build time. We settled on these after A/B testing showed recall@10 plateaued at 0.91 with m=16.

Embedding Normalization Gotcha

Another root-cause failure: our recall rate dropped from 0.89 to 0.62 after switching embedding models. The culprit was unnormalized vectors. pgvector's cosine distance operator expects unit vectors, but our new model (a fine-tuned bge-large-en-v1.5) returned unnormalized embeddings. The fix:

import numpy as np

def normalize_embedding(embedding):
    vec = np.array(embedding)
    norm = np.linalg.norm(vec)
    return (vec / norm).tolist() if norm > 0 else vec.tolist()

After normalization, recall recovered to 0.87—close enough to justify the model swap for its 40% faster inference.

Production Metrics

After integrating headroom's compression layer:

Median token count: 180K → 72K (60% reduction)
P95 LLM latency: 8.2s → 3.1s (faster decode with shorter context)
Cost per 1K requests: $34 → $14 (fewer input tokens)

The compression hit rate varies by data type: JSON tool outputs compress 85-95%, whereas natural language retrieval chunks compress 40-50%. Code blocks are exempted from aggressive compression to preserve syntax.

The Unresolved Dilemma

One open question: when does semantic deduplication hurt answer quality? We've seen cases where two nearly identical chunks contained subtle differences that mattered for the final answer—one mentioned a deprecated API parameter, the other didn't. The compression layer merged them, and the LLM missed the deprecation warning. We're experimenting with a similarity threshold of 0.92 (down from 0.96) to retain more borderline chunks, but it's a precision-recall trade-off with no clean answer. How are others handling near-duplicate retrieval results in production RAG systems without bloating context?

DEV Community