RAG vs CAG: When Cache-Augmented Generation Beats Retrieval

User Query
    ↓
[1. Embed Query]
    ↓
[2. Search Vector DB] → Top-K Documents
    ↓
[3. Construct Prompt]
    Query + Retrieved Docs
    ↓
[4. LLM Generation]
    ↓
Response

from langchain_openai import ChatOpenAI, OpenAIEmbeddings
from langchain_pinecone import PineconeVectorStore
from langchain.chains import RetrievalQA

class RAGSystem:
    def __init__(self, index_name: str):
        self.llm = ChatOpenAI(model="gpt-5.2")
        self.embeddings = OpenAIEmbeddings(model="text-embedding-3-large")
        self.vectorstore = PineconeVectorStore.from_existing_index(
            index_name=index_name,
            embedding=self.embeddings
        )
        self.retriever = self.vectorstore.as_retriever(
            search_type="mmr",
            search_kwargs={"k": 5, "fetch_k": 20}
        )

    def query(self, question: str) -> str:
        """Query with real-time retrieval."""
        # 1. Retrieve relevant documents
        docs = self.retriever.invoke(question)

        # 2. Build context
        context = "\n\n".join([doc.page_content for doc in docs])

        # 3. Generate response
        prompt = f"""Answer the question based on the following context:

Context:
{context}

Question: {question}

Answer:"""

        response = self.llm.invoke(prompt)
        return response.content

Total latency = Embedding (50ms) + Search (100ms) + Generation (500ms)
             = ~650ms per query

[Offline: Knowledge Preloading]
    All Documents → LLM Context → KV Cache (saved to disk)

[Online: Query Time]
    User Query + Cached KV → LLM Generation → Response

# Simplified attention mechanism
# For each layer, the model computes:
K = input @ W_k  # Keys
V = input @ W_v  # Values

# These are cached and reused for subsequent tokens
# Instead of recomputing for the entire context each time

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import pickle

class CAGSystem:
    def __init__(self, model_name: str = "meta-llama/Llama-3.1-70B-Instruct"):
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForCausalLM.from_pretrained(
            model_name,
            torch_dtype=torch.bfloat16,
            device_map="auto"
        )
        self.kv_cache = None

    def preload_knowledge(self, documents: list[str], cache_path: str):
        """Preload all documents and cache KV states."""
        # Combine all documents into context
        knowledge_text = "\n\n---\n\n".join(documents)

        # Create the preloaded prompt
        preload_prompt = f"""You are a helpful assistant with access to the following knowledge base:

{knowledge_text}

Use this knowledge to answer questions accurately. If the answer is not in the knowledge base, say so.

"""
        # Tokenize
        inputs = self.tokenizer(
            preload_prompt,
            return_tensors="pt",
            truncation=True,
            max_length=self.model.config.max_position_embeddings - 1000  # Leave room for query
        ).to(self.model.device)

        # Forward pass to generate KV cache
        with torch.no_grad():
            outputs = self.model(
                **inputs,
                use_cache=True,
                return_dict=True
            )
            self.kv_cache = outputs.past_key_values

        # Save cache to disk
        with open(cache_path, 'wb') as f:
            pickle.dump(self.kv_cache, f)

        print(f"Cached {len(documents)} documents ({inputs['input_ids'].shape[1]} tokens)")

    def load_cache(self, cache_path: str):
        """Load precomputed KV cache from disk."""
        with open(cache_path, 'rb') as f:
            self.kv_cache = pickle.load(f)

    def query(self, question: str) -> str:
        """Query using cached knowledge (no retrieval)."""
        # Format the query
        query_prompt = f"\n\nQuestion: {question}\n\nAnswer:"

        inputs = self.tokenizer(
            query_prompt,
            return_tensors="pt"
        ).to(self.model.device)

        # Generate using cached KV
        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                past_key_values=self.kv_cache,
                max_new_tokens=500,
                do_sample=True,
                temperature=0.7
            )

        response = self.tokenizer.decode(
            outputs[0][inputs['input_ids'].shape[1]:],
            skip_special_tokens=True
        )
        return response

Total latency = Generation only (500ms)
             = ~500ms per query

Speedup vs RAG: ~23% faster (no retrieval overhead)

KV Cache Size = 2 × num_layers × num_kv_heads × head_dim × seq_len × bytes_per_param

Example for Llama 3.1 70B (80 layers, 8 KV heads, 128 head_dim, bfloat16):
= 2 × 80 × 8 × 128 × seq_len × 2 bytes
= 327,680 × seq_len bytes
= ~312 KB per 1K tokens
= ~31.2 GB for 100K tokens

RAG (per query):
  - Embed query: ~50ms
  - Vector search: ~100ms
  - Prefill retrieved docs (~5K tokens): ~200ms
  - Decode response: ~300ms
  Total: ~650ms

CAG (per query):
  - Load cached KV: ~50ms (from SSD) or ~5ms (from RAM)
  - Decode response: ~300ms
  Total: ~350ms

CAG prefill (one-time):
  - Process 100K tokens: ~30-60 seconds
  - Amortized over 10K queries: ~3-6ms per query

class BatchCAGSystem:
    """Efficient batch inference with shared KV cache."""

    def __init__(self, model, shared_cache):
        self.model = model
        self.shared_cache = shared_cache  # Same knowledge for all queries

    def batch_query(self, questions: list[str]) -> list[str]:
        # All queries share the same preloaded context
        # Only the query tokens differ per batch item
        batch_inputs = self.tokenizer(
            [f"\n\nQuestion: {q}\n\nAnswer:" for q in questions],
            padding=True,
            return_tensors="pt"
        )

        # Replicate KV cache for batch dimension
        batch_cache = self._expand_cache_for_batch(
            self.shared_cache,
            batch_size=len(questions)
        )

        outputs = self.model.generate(
            **batch_inputs,
            past_key_values=batch_cache,
            max_new_tokens=200
        )

        return self.tokenizer.batch_decode(outputs, skip_special_tokens=True)

Is your knowledge base > 500K tokens?
    ├── YES → Use RAG
    └── NO → Continue...

Does your knowledge change frequently?
    ├── YES (daily/hourly) → Use RAG
    └── NO (weekly/monthly) → Continue...

Is latency critical (<200ms)?
    ├── YES → Use CAG
    └── NO → Either works, consider complexity

Do you already have vector DB infrastructure?
    ├── YES → RAG may be easier
    └── NO → CAG avoids new infrastructure

class HybridSystem:
    """Preload common knowledge, retrieve for edge cases."""

    def __init__(self):
        self.cag = CAGSystem()
        self.rag = RAGSystem(index_name="edge-cases")
        self.cag.load_cache("core_knowledge.cache")

    def query(self, question: str) -> str:
        # Try CAG first (fast path)
        cag_response = self.cag.query(question)

        # Check confidence / detect "I don't know"
        if self.needs_retrieval(cag_response):
            # Fall back to RAG for edge cases
            return self.rag.query(question)

        return cag_response

    def needs_retrieval(self, response: str) -> bool:
        """Detect when CAG doesn't have the answer."""
        uncertainty_phrases = [
            "I don't have information",
            "not in my knowledge",
            "I'm not sure",
            "cannot find"
        ]
        return any(phrase in response.lower() for phrase in uncertainty_phrases)

class TieredCAGSystem:
    """Multiple cache tiers for different knowledge domains."""

    def __init__(self):
        self.caches = {
            "products": "product_docs.cache",
            "policies": "company_policies.cache",
            "technical": "tech_manual.cache"
        }
        self.active_cache = None
        self.classifier = self.load_query_classifier()

    def query(self, question: str) -> str:
        # Classify query to select appropriate cache
        domain = self.classifier.predict(question)

        # Load domain-specific cache
        if self.active_cache != domain:
            self.load_cache(self.caches[domain])
            self.active_cache = domain

        return self.generate(question)

class AugmentedCAGSystem:
    """CAG for static knowledge, RAG for real-time data."""

    def __init__(self):
        self.cag = CAGSystem()  # Product catalog, policies
        self.rag = RAGSystem()  # Inventory, prices, promotions

    def query(self, question: str) -> str:
        # Get static context from CAG
        static_context = self.cag.get_cached_context()

        # Get dynamic context from RAG
        dynamic_docs = self.rag.retrieve(question)
        dynamic_context = "\n".join([d.page_content for d in dynamic_docs])

        # Combine for generation
        prompt = f"""Static Knowledge:
{static_context}

Current Information:
{dynamic_context}

Question: {question}
Answer:"""

        return self.llm.invoke(prompt)

Position in Context → Recall Accuracy
Beginning (0-10%):    ~95%
Middle (40-60%):      ~60-70%
End (90-100%):        ~90%

# Simplified attention visualization
def attention_per_token(context_length: int, query_tokens: int = 100):
    """Average attention each context token receives."""
    # Softmax distributes attention across all tokens
    # More tokens = less attention per token
    return query_tokens / context_length

# 10K context: 1% attention per token
# 100K context: 0.1% attention per token
# 1M context: 0.01% attention per token

# Example: Retrieval failure detection
def detect_retrieval_failure(query: str, docs: list, response: str) -> bool:
    """Heuristics for detecting poor retrieval."""
    signals = []

    # Low relevance scores
    if all(doc.score < 0.5 for doc in docs):
        signals.append("low_relevance")

    # Response contradicts retrieved docs
    if contains_contradiction(response, docs):
        signals.append("contradiction")

    # Response uses knowledge not in docs (hallucination)
    if contains_external_knowledge(response, docs):
        signals.append("potential_hallucination")

    return len(signals) > 0

Question: "What is the revenue of the company that acquired the
           startup founded by the person who invented X?"

Required reasoning:
1. Who invented X? → Person A
2. What startup did Person A found? → Startup B
3. Who acquired Startup B? → Company C
4. What is Company C's revenue? → $Y

RAG typically retrieves docs for step 1, missing steps 2-4.

Latency distribution (p50/p95/p99):
- Embedding: 30ms / 80ms / 200ms
- Vector search: 50ms / 150ms / 500ms
- Network: 20ms / 100ms / 300ms
- Total retrieval: 100ms / 330ms / 1000ms

CAG has consistent latency (no retrieval variance)

import hashlib
from datetime import datetime
from dataclasses import dataclass

@dataclass
class CacheMetadata:
    version: str
    created_at: datetime
    document_hashes: dict[str, str]
    model_name: str
    token_count: int

class VersionedCAGCache:
    def __init__(self, cache_dir: str):
        self.cache_dir = cache_dir

    def create_cache(self, documents: dict[str, str], model) -> str:
        """Create versioned cache with metadata."""
        # Hash documents for change detection
        doc_hashes = {
            name: hashlib.sha256(content.encode()).hexdigest()[:16]
            for name, content in documents.items()
        }

        # Version based on content + model
        version_string = f"{sorted(doc_hashes.items())}-{model.name}"
        version = hashlib.sha256(version_string.encode()).hexdigest()[:12]

        # Build cache
        kv_cache = self._build_kv_cache(documents, model)

        # Save with metadata
        metadata = CacheMetadata(
            version=version,
            created_at=datetime.now(),
            document_hashes=doc_hashes,
            model_name=model.name,
            token_count=sum(len(d.split()) for d in documents.values())
        )

        self._save_cache(version, kv_cache, metadata)
        return version

    def needs_rebuild(self, current_docs: dict[str, str]) -> bool:
        """Check if cache needs rebuilding due to document changes."""
        latest = self._load_latest_metadata()
        if not latest:
            return True

        current_hashes = {
            name: hashlib.sha256(content.encode()).hexdigest()[:16]
            for name, content in current_docs.items()
        }

        return current_hashes != latest.document_hashes

from prometheus_client import Counter, Histogram, Gauge

# RAG-specific metrics
rag_retrieval_latency = Histogram(
    'rag_retrieval_latency_seconds',
    'Time spent on retrieval',
    buckets=[0.05, 0.1, 0.25, 0.5, 1.0, 2.5]
)
rag_retrieval_count = Counter(
    'rag_retrieval_total',
    'Total retrieval operations',
    ['status']  # success, failure, timeout
)
rag_relevance_score = Histogram(
    'rag_top_doc_relevance',
    'Relevance score of top retrieved document',
    buckets=[0.1, 0.3, 0.5, 0.7, 0.9]
)

# CAG-specific metrics
cag_cache_load_latency = Histogram(
    'cag_cache_load_latency_seconds',
    'Time to load KV cache'
)
cag_cache_size_bytes = Gauge(
    'cag_cache_size_bytes',
    'Size of loaded KV cache'
)
cag_cache_hit_rate = Gauge(
    'cag_cache_hit_rate',
    'Percentage of queries served from cache'
)

# Shared metrics
generation_latency = Histogram(
    'llm_generation_latency_seconds',
    'Time spent on LLM generation',
    ['method']  # rag, cag, hybrid
)
response_quality_score = Histogram(
    'response_quality_score',
    'Quality score from evaluation model'
)

import random
from typing import Literal

class RAGvsCAGExperiment:
    """A/B test RAG vs CAG in production."""

    def __init__(self, rag_system, cag_system, cag_traffic_pct: float = 0.1):
        self.rag = rag_system
        self.cag = cag_system
        self.cag_traffic_pct = cag_traffic_pct
        self.metrics = MetricsCollector()

    def query(self, question: str, user_id: str) -> tuple[str, Literal["rag", "cag"]]:
        # Deterministic assignment based on user_id
        use_cag = int(hashlib.md5(user_id.encode()).hexdigest(), 16) % 100 < self.cag_traffic_pct * 100

        method = "cag" if use_cag else "rag"
        start_time = time.time()

        try:
            if use_cag:
                response = self.cag.query(question)
            else:
                response = self.rag.query(question)

            latency = time.time() - start_time
            self.metrics.record(method, latency, success=True)

            return response, method

        except Exception as e:
            self.metrics.record(method, time.time() - start_time, success=False)
            # Fallback to RAG on CAG failure
            if use_cag:
                return self.rag.query(question), "rag"
            raise

    def get_results(self) -> dict:
        """Get A/B test results."""
        return {
            "rag": {
                "p50_latency": self.metrics.percentile("rag", 50),
                "p95_latency": self.metrics.percentile("rag", 95),
                "success_rate": self.metrics.success_rate("rag"),
                "sample_size": self.metrics.count("rag")
            },
            "cag": {
                "p50_latency": self.metrics.percentile("cag", 50),
                "p95_latency": self.metrics.percentile("cag", 95),
                "success_rate": self.metrics.success_rate("cag"),
                "sample_size": self.metrics.count("cag")
            }
        }

class MultiTenantCAG:
    """Isolated CAG caches per tenant."""

    def __init__(self, model, cache_dir: str):
        self.model = model
        self.cache_dir = cache_dir
        self.loaded_tenant: str | None = None
        self.current_cache = None

    def query(self, tenant_id: str, question: str) -> str:
        # Load tenant-specific cache if needed
        if self.loaded_tenant != tenant_id:
            cache_path = f"{self.cache_dir}/{tenant_id}/knowledge.cache"

            if not os.path.exists(cache_path):
                raise TenantCacheNotFound(tenant_id)

            self.current_cache = self._load_cache(cache_path)
            self.loaded_tenant = tenant_id

        return self._generate(question, self.current_cache)

    def build_tenant_cache(self, tenant_id: str, documents: list[str]):
        """Build isolated cache for a tenant."""
        cache_path = f"{self.cache_dir}/{tenant_id}/knowledge.cache"
        os.makedirs(os.path.dirname(cache_path), exist_ok=True)

        # Tenant data never mixes with other tenants
        kv_cache = self._build_cache(documents)
        self._save_cache(cache_path, kv_cache)

class ResilientKnowledgeSystem:
    """Production system with fallbacks."""

    def __init__(self, cag: CAGSystem, rag: RAGSystem):
        self.cag = cag
        self.rag = rag
        self.circuit_breaker = CircuitBreaker(
            failure_threshold=5,
            recovery_timeout=60
        )

    async def query(self, question: str) -> str:
        # Try CAG first (fast path)
        if self.circuit_breaker.is_closed("cag"):
            try:
                return await asyncio.wait_for(
                    self.cag.query(question),
                    timeout=5.0
                )
            except (asyncio.TimeoutError, CacheLoadError) as e:
                self.circuit_breaker.record_failure("cag")
                logger.warning(f"CAG failed, falling back to RAG: {e}")

        # Fallback to RAG
        try:
            return await self.rag.query(question)
        except Exception as e:
            logger.error(f"Both CAG and RAG failed: {e}")
            return "I'm having trouble accessing my knowledge base. Please try again."

import json
from dataclasses import dataclass
from typing import Callable

@dataclass
class BenchmarkResult:
    method: str
    accuracy: float
    faithfulness: float
    avg_latency_ms: float
    p95_latency_ms: float
    cost_per_query: float
    total_queries: int

class RAGvsCAGBenchmark:
    """Comprehensive benchmark for comparing RAG and CAG."""

    def __init__(self, rag_system, cag_system, eval_model):
        self.rag = rag_system
        self.cag = cag_system
        self.eval_model = eval_model  # For LLM-as-judge

    def run_benchmark(self, test_cases: list[dict]) -> dict[str, BenchmarkResult]:
        """
        test_cases format:
        [
            {
                "question": "What is the return policy?",
                "expected_answer": "30-day money-back guarantee",
                "source_docs": ["policy.md"],  # For faithfulness check
                "difficulty": "easy"  # easy, medium, hard
            }
        ]
        """
        results = {"rag": [], "cag": []}

        for test in test_cases:
            # Run both systems
            rag_result = self._evaluate_single(self.rag, test, "rag")
            cag_result = self._evaluate_single(self.cag, test, "cag")

            results["rag"].append(rag_result)
            results["cag"].append(cag_result)

        return {
            "rag": self._aggregate_results(results["rag"], "rag"),
            "cag": self._aggregate_results(results["cag"], "cag")
        }

    def _evaluate_single(self, system, test: dict, method: str) -> dict:
        start = time.time()
        response = system.query(test["question"])
        latency = (time.time() - start) * 1000

        # LLM-as-judge evaluation
        eval_prompt = f"""Evaluate this response:

Question: {test["question"]}
Expected: {test["expected_answer"]}
Response: {response}

Rate on a scale of 1-5:
1. Accuracy (correctness):
2. Completeness (covers all aspects):
3. Faithfulness (no hallucinations):

Return JSON: {{"accuracy": X, "completeness": X, "faithfulness": X}}"""

        scores = json.loads(self.eval_model.invoke(eval_prompt))

        return {
            "latency_ms": latency,
            "accuracy": scores["accuracy"] / 5,
            "faithfulness": scores["faithfulness"] / 5,
            "response": response
        }

    def _aggregate_results(self, results: list, method: str) -> BenchmarkResult:
        latencies = [r["latency_ms"] for r in results]
        return BenchmarkResult(
            method=method,
            accuracy=sum(r["accuracy"] for r in results) / len(results),
            faithfulness=sum(r["faithfulness"] for r in results) / len(results),
            avg_latency_ms=sum(latencies) / len(latencies),
            p95_latency_ms=sorted(latencies)[int(len(latencies) * 0.95)],
            cost_per_query=self._calculate_cost(method),
            total_queries=len(results)
        )

def needle_in_haystack_test(system, knowledge_base: str, needles: list[dict]):
    """
    Test retrieval of specific facts buried in large context.

    needles format:
    [
        {
            "fact": "The secret code is ALPHA-7742",
            "position": 0.5,  # Middle of context
            "question": "What is the secret code?"
        }
    ]
    """
    results = []

    for needle in needles:
        # Insert needle at specified position
        insert_pos = int(len(knowledge_base) * needle["position"])
        test_context = (
            knowledge_base[:insert_pos] +
            f"\n{needle['fact']}\n" +
            knowledge_base[insert_pos:]
        )

        # Query the system
        response = system.query(needle["question"], context=test_context)

        # Check if needle was found
        found = needle["fact"].lower() in response.lower()

        results.append({
            "position": needle["position"],
            "found": found,
            "response": response
        })

    # Aggregate by position
    position_accuracy = {}
    for r in results:
        pos_bucket = round(r["position"], 1)
        if pos_bucket not in position_accuracy:
            position_accuracy[pos_bucket] = []
        position_accuracy[pos_bucket].append(r["found"])

    return {
        pos: sum(found) / len(found)
        for pos, found in position_accuracy.items()
    }

## RAG vs CAG Decision Checklist

### Knowledge Base Characteristics
- [ ] Size < 500K tokens? → Favor CAG
- [ ] Size > 1M tokens? → Favor RAG
- [ ] Updates daily or more? → Favor RAG
- [ ] Static or monthly updates? → Favor CAG

### Query Patterns
- [ ] Multi-hop reasoning required? → Favor CAG
- [ ] Queries span many topics? → Favor RAG
- [ ] Queries focused on specific domain? → Favor CAG

### Infrastructure
- [ ] Already have vector DB? → RAG easier
- [ ] Have high-memory GPUs? → CAG viable
- [ ] Serverless deployment? → RAG more flexible

### Requirements
- [ ] Latency < 500ms required? → Favor CAG
- [ ] Real-time data needed? → Favor RAG
- [ ] Multi-tenant isolation needed? → Both work, different tradeoffs

### Score
- More CAG checkmarks: Start with CAG
- More RAG checkmarks: Start with RAG
- Mixed: Consider hybrid approach