Caching › Code Implementation

# We will use sentence-transformers for embeddings
# uv pip install sentence-transformers
from sentence_transformers import SentenceTransformer

# Load a small, fast embedding model suitable for caching
embedding_model = SentenceTransformer("all-MiniLM-L6-v2")

@dataclass
class CacheEntry:
    query: str
    response: str
    embedding: np.ndarray
    created_at: float
    hit_count: int = 0


@dataclass
class SemanticCache:
    threshold: float = 0.90
    ttl_seconds: float = 3600.0
    entries: list = field(default_factory=list)
    hits: int = 0
    misses: int = 0

    def _cosine_similarity(self, a: np.ndarray, b: np.ndarray) -> float:
        """Compute cosine similarity between two unit-normalized vectors."""
        return float(
            np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) + 1e-10)
        )

    def _is_expired(self, entry: CacheEntry) -> bool:
        return (time.time() - entry.created_at) > self.ttl_seconds

    def lookup(self, query: str) -> Optional[str]:
        """Search cache for a semantically similar query. Returns response or None."""
        query_embedding = embedding_model.encode(
            query, normalize_embeddings=True
        )
        now = time.time()

        best_similarity = 0.0
        best_entry = None

        for entry in self.entries:
            if (now - entry.created_at) > self.ttl_seconds:
                continue
            sim = self._cosine_similarity(query_embedding, entry.embedding)
            if sim > best_similarity:
                best_similarity = sim
                best_entry = entry

        if best_entry is not None and best_similarity >= self.threshold:
            best_entry.hit_count += 1
            self.hits += 1
            return best_entry.response
        else:
            self.misses += 1
            return None

    def store(self, query: str, response: str) -> None:
        """Store a new query-response pair in the cache."""
        embedding = embedding_model.encode(query, normalize_embeddings=True)
        entry = CacheEntry(
            query=query,
            response=response,
            embedding=embedding,
            created_at=time.time(),
        )
        self.entries.append(entry)

    @property
    def hit_rate(self) -> float:
        total = self.hits + self.misses
        return self.hits / total if total > 0 else 0.0

# Simulated LLM responses for demonstration
def mock_llm(query: str) -> str:
    """Simulate an LLM call with a small delay."""
    time.sleep(0.05)  # Simulate 50ms inference time
    responses = {
        "What is the capital of France?": "The capital of France is Paris.",
        "What is the capital of Germany?": "The capital of Germany is Berlin.",
        "Explain neural networks briefly.": "Neural networks are computational models inspired by the brain, consisting of layers of interconnected nodes that learn patterns from data.",
        "What are transformers in NLP?": "Transformers are a neural network architecture that uses self-attention to process sequences in parallel, forming the basis of models like BERT and GPT.",
    }
    for key, val in responses.items():
        if any(word in query.lower() for word in key.lower().split()):
            return val
    return f"Response to: {query}"


# Populate the cache with some initial queries
cache = SemanticCache(threshold=0.88, ttl_seconds=3600)

initial_queries = [
    (
        "What is the capital of France?",
        mock_llm("What is the capital of France?"),
    ),
    (
        "Explain neural networks briefly.",
        mock_llm("Explain neural networks briefly."),
    ),
    (
        "What are transformers in NLP?",
        mock_llm("What are transformers in NLP?"),
    ),
]

for query, response in initial_queries:
    cache.store(query, response)

# Test with queries at varying semantic distances from cached content
test_queries = [
    "What is the capital of France?",  # Exact match
    "Which city is the capital of France?",  # Paraphrase (should hit)
    "Where is France's capital located?",  # Rephrased (should hit)
    "Tell me about French neural networks.",  # Mixed topics (may hit or miss)
    "What is the capital of Japan?",  # Different country (should miss)
    "How do transformers work in deep learning?",  # Paraphrase (should hit)
]

results = []
for query in test_queries:
    start = time.time()
    cached_response = cache.lookup(query)
    latency_ms = (time.time() - start) * 1000

    if cached_response:
        source = "CACHE HIT"
        response = cached_response
    else:
        source = "CACHE MISS"
        response = mock_llm(query)
        cache.store(query, response)

    results.append(
        {
            "query": query,
            "source": source,
            "latency_ms": latency_ms,
        }
    )

# Simulate how threshold affects hit rate vs. response quality

# A set of base queries and their known paraphrases
query_groups = [
    {
        "canonical": "What is the capital of France?",
        "variants": [
            ("What is the capital of France?", 1.00),  # Exact
            ("Which city is France's capital?", 0.93),  # Near paraphrase
            ("France capital city name?", 0.87),  # Compressed form
            ("Where is France's government located?", 0.75),  # Related concept
            ("Name the French capital city.", 0.91),  # Rephrased
        ],
    },
    {
        "canonical": "Explain transformer architecture.",
        "variants": [
            ("Explain transformer architecture.", 1.00),
            ("How does the transformer model work?", 0.92),
            ("What is a transformer neural network?", 0.88),
            ("Describe the attention mechanism.", 0.72),
            ("What are key components of transformers?", 0.85),
        ],
    },
]

thresholds = np.round(np.arange(0.60, 1.01, 0.02), 2)
hit_rates = []
correct_match_rates = []

for threshold in thresholds:
    hits = 0
    correct = 0
    total = 0
    for group in query_groups:
        for variant_query, true_similarity in group["variants"]:
            total += 1
            if true_similarity >= threshold:
                hits += 1
                # A "correct" match is one where the similarity is genuinely high
                if true_similarity >= 0.85:
                    correct += 1
    hit_rates.append(hits / total)
    correct_match_rates.append(correct / total)

# Simulate a request workload with a Zipfian query distribution
# (a small number of popular queries, a long tail of rare ones)

n_unique_queries = 500
n_total_requests = 5000
inference_latency_ms = 800
cache_lookup_ms = 15

# Zipf distribution: query i has probability proportional to 1/i
ranks = np.arange(1, n_unique_queries + 1)
zipf_probs = 1.0 / ranks
zipf_probs /= zipf_probs.sum()

# Sample request sequence
query_ids = np.random.choice(
    n_unique_queries, size=n_total_requests, p=zipf_probs
)

# Simulate caching with various cache sizes (max unique entries before eviction)
cache_sizes = [10, 25, 50, 100, 200, 500]
simulated_hit_rates = []
simulated_avg_latencies = []

for max_size in cache_sizes:
    cache_store = {}  # Maps query_id -> True (just track presence)
    cache_order = []  # LRU eviction order
    hits = 0

    for qid in query_ids:
        if qid in cache_store:
            hits += 1
            # Move to end (most recently used)
            cache_order.remove(qid)
            cache_order.append(qid)
        else:
            # Cache miss: add to cache
            if len(cache_order) >= max_size:
                # Evict least recently used
                evicted = cache_order.pop(0)
                del cache_store[evicted]
            cache_store[qid] = True
            cache_order.append(qid)

    hit_rate_sim = hits / n_total_requests
    avg_latency = hit_rate_sim * cache_lookup_ms + (1 - hit_rate_sim) * (
        inference_latency_ms + cache_lookup_ms
    )

    simulated_hit_rates.append(hit_rate_sim)
    simulated_avg_latencies.append(avg_latency)