Context Compression Techniques — How We Cut $4k/Month in Token Costs Without Losing RAG Accuracy

import numpy as np
from sentence_transformers import SentenceTransformer
from typing import List, Dict
import openai  # v1.0+

class HybridCompressor:
    def __init__(self, extractive_model: str = "all-MiniLM-L6-v2",
                 abstractive_model: str = "gpt-4o-mini",
                 extractive_threshold: float = 0.3,
                 max_tokens: int = 2000):
        self.embedder = SentenceTransformer(extractive_model)
        self.extractive_threshold = extractive_threshold
        self.max_tokens = max_tokens
        self.abstractive_model = abstractive_model

    def _chunk_sentences(self, text: str) -> List[str]:
        # Simple sentence splitter — use spaCy or nltk for production
        import re
        sentences = re.split(r'(?<=[.!?])\s+', text)
        return [s.strip() for s in sentences if s.strip()]

    def _extractive_filter(self, sentences: List[str], query: str) -> List[str]:
        if not sentences:
            return []
        # Encode query and sentences
        query_emb = self.embedder.encode([query])[0]
        sent_embs = self.embedder.encode(sentences)
        # Compute cosine similarity
        scores = np.dot(sent_embs, query_emb) / (
            np.linalg.norm(sent_embs, axis=1) * np.linalg.norm(query_emb)
        )
        # Keep sentences above threshold, but at least 2
        kept = [s for s, score in zip(sentences, scores) if score >= self.extractive_threshold]
        if len(kept) < 2:
            # Fallback: keep top-2 by score
            top_indices = np.argsort(scores)[-2:]
            kept = [sentences[i] for i in sorted(top_indices)]
        return kept

    def _abstractive_compress(self, text: str, query: str) -> str:
        # Use a cheap LLM to compress
        response = openai.chat.completions.create(
            model=self.abstractive_model,
            messages=[
                {"role": "system", "content": "Compress the following text to its essential information. Keep all facts, dates, numbers, and named entities. Do not add new information."},
                {"role": "user", "content": f"Query: {query}\n\nText: {text}"}
            ],
            max_tokens=self.max_tokens,
            temperature=0.0  # deterministic
        )
        return response.choices[0].message.content

    def compress(self, context: str, query: str) -> str:
        # Step 1: sentence split
        sentences = self._chunk_sentences(context)
        # Step 2: extractive filter
        filtered = self._extractive_filter(sentences, query)
        if not filtered:
            return ""  # no relevant content found
        # Step 3: abstractive compression on the survivors
        compressed = self._abstractive_compress(" ".join(filtered), query)
        return compressed

# Usage example
if __name__ == "__main__":
    compressor = HybridCompressor()
    context = """
    Aspirin is a nonsteroidal anti-inflammatory drug (NSAID).
    It is used to reduce pain, fever, and inflammation.
    Aspirin may increase bleeding risk when used with Warfarin.
    Warfarin requires regular INR monitoring.
    """
    query = "Can I take aspirin with Warfarin?"
    result = compressor.compress(context, query)
    print("Compressed:", result)
    # Expected: keeps the bleeding risk sentence

import numpy as np
from sentence_transformers import SentenceTransformer
from typing import List

class QueryAwareExtractiveCompressor:
    def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
        self.model = SentenceTransformer(model_name)

    def compress(self, documents: List[str], query: str, top_k: int = 5) -> List[str]:
        """
        Compress a list of document strings to the top-k sentences most relevant to the query.
        """
        # Split all documents into sentences
        all_sentences = []
        for doc in documents:
            # Simple split — use spaCy for production
            sentences = doc.split('. ')
            all_sentences.extend([s.strip() + '.' for s in sentences if s.strip()])
        
        if not all_sentences:
            return []
        
        # Encode query and sentences
        query_emb = self.model.encode([query])[0]
        sent_embs = self.model.encode(all_sentences)
        
        # Compute cosine similarity
        scores = np.dot(sent_embs, query_emb) / (
            np.linalg.norm(sent_embs, axis=1) * np.linalg.norm(query_emb)
        )
        
        # Get top-k indices
        top_indices = np.argsort(scores)[-top_k:][::-1]
        
        # Return top-k sentences
        return [all_sentences[i] for i in top_indices]

# Usage
if __name__ == "__main__":
    compressor = QueryAwareExtractiveCompressor()
    docs = [
        "France is a country in Europe. Its capital is Paris. The population is 67 million.",
        "Paris is known for the Eiffel Tower. It is a major fashion capital."
    ]
    query = "What is the capital of France?"
    result = compressor.compress(docs, query, top_k=2)
    print("Compressed:", result)
    # Expected: ['Its capital is Paris.', 'Paris is known for the Eiffel Tower.']

import time
from typing import Optional

class CompressionDecisionEngine:
    def __init__(self, min_tokens: int = 1500, max_latency_ms: int = 200):
        self.min_tokens = min_tokens
        self.max_latency_ms = max_latency_ms

    def should_compress(self, context: str, query: str, 
                        required_exact_wording: bool = False,
                        real_time: bool = False) -> tuple:
        """
        Returns (should_compress: bool, reason: Optional[str])
        """
        # Estimate token count (rough: 1 token ~= 4 chars for English)
        token_count = len(context) // 4
        
        if token_count < self.min_tokens:
            return False, f"Context too small ({token_count} tokens < {self.min_tokens})"
        
        if required_exact_wording:
            return False, "Exact wording required — use extractive only, not compression"
        
        if real_time:
            return False, "Real-time latency requirement — skip compression"
        
        return True, None

# Usage
if __name__ == "__main__":
    engine = CompressionDecisionEngine()
    context = "A" * 10000  # ~2500 tokens
    should, reason = engine.should_compress(context, "test", real_time=True)
    print(f"Should compress: {should}, Reason: {reason}")
    # Output: Should compress: False, Reason: Real-time latency requirement — skip compression

import hashlib
import json
import redis
import numpy as np
from sentence_transformers import SentenceTransformer
from typing import List, Dict

class ProductionCompressor:
    def __init__(self, redis_host: str = "localhost", redis_port: int = 6379,
                 ttl_seconds: int = 3600, batch_size: int = 64):
        self.redis = redis.Redis(host=redis_host, port=redis_port, decode_responses=True)
        self.ttl = ttl_seconds
        self.batch_size = batch_size
        self.model = SentenceTransformer("all-MiniLM-L6-v2")

    def _doc_key(self, doc_id: str) -> str:
        return f"compressed:{doc_id}"

    def compress_document(self, doc_id: str, text: str, query: str) -> str:
        # Check cache first
        cached = self.redis.get(self._doc_key(doc_id))
        if cached:
            return cached
        
        # Compress (simplified — use hybrid from earlier)
        sentences = text.split('. ')
        if not sentences:
            return ""
        
        # Batch encode
        query_emb = self.model.encode([query])[0]
        sent_embs = self.model.encode(sentences, batch_size=self.batch_size)
        scores = np.dot(sent_embs, query_emb) / (
            np.linalg.norm(sent_embs, axis=1) * np.linalg.norm(query_emb)
        )
        top_indices = np.argsort(scores)[-5:][::-1]
        compressed = '. '.join([sentences[i] for i in sorted(top_indices)])
        
        # Cache
        self.redis.setex(self._doc_key(doc_id), self.ttl, compressed)
        return compressed

# Usage
if __name__ == "__main__":
    compressor = ProductionCompressor()
    result = compressor.compress_document("doc_123", "France is a country. Paris is its capital.", "capital of France")
    print("Compressed:", result)

import numpy as np
from sentence_transformers import SentenceTransformer
from typing import List

class AdaptiveThresholdCompressor:
    def __init__(self, model_name: str = "all-MiniLM-L6-v2"):
        self.model = SentenceTransformer(model_name)

    def _query_specificity(self, query: str) -> float:
        # Simple specificity: number of unique tokens / total tokens
        # Higher = more specific
        tokens = query.lower().split()
        if not tokens:
            return 0.0
        unique_ratio = len(set(tokens)) / len(tokens)
        return unique_ratio

    def compress(self, sentences: List[str], query: str) -> List[str]:
        if not sentences:
            return []
        
        specificity = self._query_specificity(query)
        # Adaptive threshold: 0.2 for vague, 0.5 for specific
        threshold = 0.2 + (specificity * 0.3)  # range: 0.2 to 0.5
        
        query_emb = self.model.encode([query])[0]
        sent_embs = self.model.encode(sentences)
        scores = np.dot(sent_embs, query_emb) / (
            np.linalg.norm(sent_embs, axis=1) * np.linalg.norm(query_emb)
        )
        
        kept = [s for s, score in zip(sentences, scores) if score >= threshold]
        # Always keep at least 2
        if len(kept) < 2:
            top_indices = np.argsort(scores)[-2:]
            kept = [sentences[i] for i in sorted(top_indices)]
        
        return kept

# Usage
if __name__ == "__main__":
    compressor = AdaptiveThresholdCompressor()
    sentences = ["France is in Europe.", "Paris is the capital.", "The GDP is $3 trillion."]
    result = compressor.compress(sentences, "What is the GDP of France?")
    print("Compressed:", result)
    # Expected: keeps the GDP sentence and capital sentence

import time
from sentence_transformers import CrossEncoder

class RerankingCompressor:
    def __init__(self):
        self.reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')

    def rerank_and_compress(self, documents: List[str], query: str, top_k: int = 3) -> List[str]:
        # Rerank documents using cross-encoder
        pairs = [(query, doc) for doc in documents]
        scores = self.reranker.predict(pairs)
        top_indices = np.argsort(scores)[-top_k:][::-1]
        top_docs = [documents[i] for i in top_indices]
        
        # Now compress the top documents (extractive)
        # ... (use the extractive compressor from earlier)
        return top_docs  # simplified

# Benchmark: compare latency
if __name__ == "__main__":
    # Simulate 20 documents, each 500 tokens
    docs = ["A " * 125] * 20
    query = "test"
    
    # Compression only
    start = time.time()
    # compressor.compress(docs, query)
    print(f"Compression only: {time.time() - start:.2f}s")
    
    # Reranking + compression
    start = time.time()
    # rerank_and_compress(docs, query)
    print(f"Reranking + compression: {time.time() - start:.2f}s")

import logging
import json
from datetime import datetime
from typing import Dict

class CompressionMonitor:
    def __init__(self, log_file: str = "compression_metrics.log"):
        self.logger = logging.getLogger(__name__)
        handler = logging.FileHandler(log_file)
        handler.setFormatter(logging.Formatter('%(message)s'))
        self.logger.addHandler(handler)
        self.logger.setLevel(logging.INFO)

    def log_compression(self, query: str, original_tokens: int, compressed_tokens: int,
                        sentences_dropped: int, strategy: str):
        ratio = compressed_tokens / original_tokens if original_tokens > 0 else 0
        record = {
            "timestamp": datetime.utcnow().isoformat(),
            "query": query,
            "original_tokens": original_tokens,
            "compressed_tokens": compressed_tokens,
            "compression_ratio": round(ratio, 2),
            "sentences_dropped": sentences_dropped,
            "strategy": strategy
        }
        self.logger.info(json.dumps(record))
        
        # Alert on anomalies
        if ratio < 0.3:
            logging.warning(f"Aggressive compression: ratio {ratio} for query '{query}'")
        if ratio > 0.9:
            logging.warning(f"Ineffective compression: ratio {ratio} for query '{query}'")

# Usage
if __name__ == "__main__":
    monitor = CompressionMonitor()
    monitor.log_compression("test", 1000, 300, 50, "extractive")