| """ |
| rag/cache.py |
| |
| Embedding cache to avoid recomputing vectors for repeated queries. |
| |
| Why this matters for production: |
| - Embedding a query takes ~20–50ms on GPU, ~200ms on CPU |
| - Popular questions get asked repeatedly (e.g. "how to reverse a list") |
| - Cache hit rate of 30–60% is typical → average latency drops significantly |
| - Eliminates redundant GPU compute → reduces cost at scale |
| |
| Cache backends supported: |
| 1. InMemoryCache — LRU dict, process-local, lost on restart (dev/test) |
| 2. DiskCache — Persistent JSON on disk (single-instance production) |
| 3. RedisCache — Distributed, shared across multiple API workers (prod) |
| |
| Benchmark (typical): |
| Cache miss: embed query → 35ms |
| Cache hit: dict lookup → 0.1ms → 350× speedup |
| |
| Architecture note: |
| Cache is keyed on (query_text + model_id) hash to avoid collisions |
| when the embedding model changes between deployments. |
| """ |
|
|
| import hashlib |
| import json |
| import logging |
| import time |
| from abc import ABC, abstractmethod |
| from collections import OrderedDict |
| from pathlib import Path |
| from typing import List, Optional |
|
|
| import numpy as np |
|
|
| import sys |
| sys.path.append(str(Path(__file__).parent.parent)) |
| from config import cfg |
|
|
| logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") |
| log = logging.getLogger(__name__) |
|
|
|
|
| |
| def make_cache_key(text: str, model_id: str) -> str: |
| """ |
| Deterministic cache key: SHA256(text + model_id). |
| Including model_id prevents stale cache hits after model upgrades. |
| """ |
| payload = f"{model_id}::{text}" |
| return hashlib.sha256(payload.encode()).hexdigest() |
|
|
|
|
| |
| class EmbeddingCache(ABC): |
| """Abstract interface for all cache backends.""" |
|
|
| @abstractmethod |
| def get(self, key: str) -> Optional[List[float]]: |
| """Return cached embedding or None on miss.""" |
| ... |
|
|
| @abstractmethod |
| def set(self, key: str, embedding: List[float]) -> None: |
| """Store embedding in cache.""" |
| ... |
|
|
| @abstractmethod |
| def clear(self) -> None: |
| """Wipe all cached entries.""" |
| ... |
|
|
| @property |
| @abstractmethod |
| def size(self) -> int: |
| """Number of cached entries.""" |
| ... |
|
|
|
|
| |
| class InMemoryCache(EmbeddingCache): |
| """ |
| Thread-safe LRU cache using OrderedDict. |
| Evicts least-recently-used entries when max_size is reached. |
| """ |
|
|
| def __init__(self, max_size: int = None): |
| self.max_size = max_size or cfg.cache.max_memory_entries |
| self._cache: OrderedDict = OrderedDict() |
| self._hits = 0 |
| self._misses = 0 |
|
|
| def get(self, key: str) -> Optional[List[float]]: |
| if key in self._cache: |
| self._cache.move_to_end(key) |
| self._hits += 1 |
| return self._cache[key] |
| self._misses += 1 |
| return None |
|
|
| def set(self, key: str, embedding: List[float]) -> None: |
| if key in self._cache: |
| self._cache.move_to_end(key) |
| else: |
| if len(self._cache) >= self.max_size: |
| self._cache.popitem(last=False) |
| self._cache[key] = embedding |
|
|
| def clear(self) -> None: |
| self._cache.clear() |
| self._hits = 0 |
| self._misses = 0 |
|
|
| @property |
| def size(self) -> int: |
| return len(self._cache) |
|
|
| @property |
| def hit_rate(self) -> float: |
| total = self._hits + self._misses |
| return self._hits / total if total > 0 else 0.0 |
|
|
| def stats(self) -> dict: |
| return { |
| "backend": "memory", |
| "size": self.size, |
| "max_size": self.max_size, |
| "hits": self._hits, |
| "misses": self._misses, |
| "hit_rate": round(self.hit_rate, 3), |
| } |
|
|
|
|
| |
| class DiskCache(EmbeddingCache): |
| """ |
| Persistent cache stored as numpy .npy files on disk. |
| Survives process restarts — good for single-instance deployments. |
| """ |
|
|
| def __init__(self, cache_dir: str = None): |
| self.cache_dir = Path(cache_dir or cfg.cache.disk_cache_dir) |
| self.cache_dir.mkdir(parents=True, exist_ok=True) |
| self._hits = 0 |
| self._misses = 0 |
|
|
| def _path(self, key: str) -> Path: |
| return self.cache_dir / f"{key}.npy" |
|
|
| def get(self, key: str) -> Optional[List[float]]: |
| path = self._path(key) |
| if path.exists(): |
| self._hits += 1 |
| return np.load(str(path)).tolist() |
| self._misses += 1 |
| return None |
|
|
| def set(self, key: str, embedding: List[float]) -> None: |
| np.save(str(self._path(key)), np.array(embedding)) |
|
|
| def clear(self) -> None: |
| for f in self.cache_dir.glob("*.npy"): |
| f.unlink() |
|
|
| @property |
| def size(self) -> int: |
| return len(list(self.cache_dir.glob("*.npy"))) |
|
|
| def stats(self) -> dict: |
| return { |
| "backend": "disk", |
| "size": self.size, |
| "hits": self._hits, |
| "misses": self._misses, |
| "hit_rate": round(self._hits / max(self._hits + self._misses, 1), 3), |
| "cache_dir": str(self.cache_dir), |
| } |
|
|
|
|
| |
| class RedisCache(EmbeddingCache): |
| """ |
| Distributed cache via Redis. |
| Shared across multiple API worker processes — required for production |
| multi-instance deployments (e.g. Kubernetes, multiple Uvicorn workers). |
| |
| Embeddings are stored as JSON strings with TTL. |
| """ |
|
|
| def __init__(self): |
| try: |
| import redis |
| self._client = redis.Redis( |
| host=cfg.cache.redis_host, |
| port=cfg.cache.redis_port, |
| db=cfg.cache.redis_db, |
| decode_responses=True, |
| ) |
| self._client.ping() |
| self._ttl = cfg.cache.redis_ttl_seconds |
| log.info(f"Redis cache connected: {cfg.cache.redis_host}:{cfg.cache.redis_port}") |
| except Exception as e: |
| raise RuntimeError(f"Redis connection failed: {e}. Is Redis running?") |
|
|
| def get(self, key: str) -> Optional[List[float]]: |
| value = self._client.get(f"embed:{key}") |
| if value: |
| return json.loads(value) |
| return None |
|
|
| def set(self, key: str, embedding: List[float]) -> None: |
| self._client.setex( |
| name=f"embed:{key}", |
| time=self._ttl, |
| value=json.dumps(embedding), |
| ) |
|
|
| def clear(self) -> None: |
| keys = self._client.keys("embed:*") |
| if keys: |
| self._client.delete(*keys) |
|
|
| @property |
| def size(self) -> int: |
| return len(self._client.keys("embed:*")) |
|
|
| def stats(self) -> dict: |
| info = self._client.info("stats") |
| return { |
| "backend": "redis", |
| "size": self.size, |
| "keyspace_hits": info.get("keyspace_hits", 0), |
| "keyspace_misses": info.get("keyspace_misses", 0), |
| } |
|
|
|
|
| |
| class CachedEmbeddingModel: |
| """ |
| Wraps any HuggingFace embedding model with a cache layer. |
| |
| Usage: |
| model = CachedEmbeddingModel(base_model, cache=InMemoryCache()) |
| embedding = model.embed_query("how to sort a list in Python") |
| # First call: computes embedding (~35ms) |
| # Repeated call: cache hit (~0.1ms) |
| """ |
|
|
| def __init__(self, base_model, cache: Optional[EmbeddingCache] = None): |
| self.base_model = base_model |
| self.model_id = getattr(base_model, "model_name", "unknown") |
| self.cache = cache or _get_default_cache() |
| self._total_queries = 0 |
| self._cache_hits = 0 |
| log.info(f"CachedEmbeddingModel initialized (backend={cfg.cache.backend})") |
|
|
| def embed_query(self, text: str) -> List[float]: |
| """Embed a single query string, using cache when available.""" |
| self._total_queries += 1 |
| key = make_cache_key(text, self.model_id) |
|
|
| |
| cached = self.cache.get(key) |
| if cached is not None: |
| self._cache_hits += 1 |
| log.debug(f"Cache HIT | hit_rate={self.hit_rate:.1%}") |
| return cached |
|
|
| |
| t0 = time.perf_counter() |
| embedding = self.base_model.embed_query(text) |
| latency_ms = (time.perf_counter() - t0) * 1000 |
|
|
| self.cache.set(key, embedding) |
| log.debug(f"Cache MISS | embed={latency_ms:.1f}ms | hit_rate={self.hit_rate:.1%}") |
| return embedding |
|
|
| def embed_documents(self, texts: List[str]) -> List[List[float]]: |
| """Batch embed documents, using cache for each individual text.""" |
| results = [] |
| for text in texts: |
| results.append(self.embed_query(text)) |
| return results |
|
|
| @property |
| def hit_rate(self) -> float: |
| return self._cache_hits / max(self._total_queries, 1) |
|
|
| def stats(self) -> dict: |
| return { |
| "total_queries": self._total_queries, |
| "cache_hits": self._cache_hits, |
| "hit_rate": round(self.hit_rate, 3), |
| "cache_stats": self.cache.stats() if hasattr(self.cache, "stats") else {}, |
| } |
|
|
|
|
| |
| def _get_default_cache() -> EmbeddingCache: |
| backend = cfg.cache.backend |
| if backend == "redis": |
| return RedisCache() |
| elif backend == "disk": |
| return DiskCache() |
| else: |
| return InMemoryCache() |
|
|
|
|
| def build_cached_embeddings(base_model=None) -> CachedEmbeddingModel: |
| """Build a cached embedding model from the configured base model.""" |
| if base_model is None: |
| from langchain_huggingface import HuggingFaceEmbeddings |
| base_model = HuggingFaceEmbeddings( |
| model_name=cfg.model.embedding_model_id, |
| encode_kwargs={"normalize_embeddings": True}, |
| ) |
| return CachedEmbeddingModel(base_model) |
|
|
