Spaces:
Running on Zero
Running on Zero
File size: 5,910 Bytes
712c393 dda0d18 712c393 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 | from __future__ import annotations
import hashlib
import json
import logging
import threading
import time
from collections import OrderedDict
from typing import Any, Dict, List, Optional, Tuple
logger = logging.getLogger(__name__)
class LRUTTLCache:
"""Thread-safe LRU cache with Time-To-Live (TTL) expiration."""
def __init__(self, max_size: int = 1000, ttl_seconds: int = 3600) -> None:
self.max_size = max(1, max_size)
self.ttl_seconds = ttl_seconds
self._cache: OrderedDict[str, Tuple[float, Any]] = OrderedDict()
self._lock = threading.RLock()
self.hits = 0
self.misses = 0
def get(self, key: str) -> Optional[Any]:
with self._lock:
if key not in self._cache:
self.misses += 1
return None
ts, val = self._cache[key]
if time.time() - ts > self.ttl_seconds:
del self._cache[key]
self.misses += 1
return None
self._cache.move_to_end(key)
self.hits += 1
return val
def set(self, key: str, value: Any) -> None:
with self._lock:
if key in self._cache:
del self._cache[key]
elif len(self._cache) >= self.max_size:
self._cache.popitem(last=False)
self._cache[key] = (time.time(), value)
def clear(self) -> None:
with self._lock:
self._cache.clear()
def stats(self) -> Dict[str, Any]:
with self._lock:
total = self.hits + self.misses
hit_rate = (self.hits / total) if total > 0 else 0.0
return {
"size": len(self._cache),
"max_size": self.max_size,
"hits": self.hits,
"misses": self.misses,
"hit_rate": round(hit_rate, 4),
}
class CacheManager:
"""
Centralized caching manager supporting:
- Query embedding caching
- Retrieval result caching (invalidated on index update)
- Reranker pair scoring cache
"""
_instance: Optional[CacheManager] = None
_singleton_lock = threading.RLock()
def __init__(
self,
enabled: bool = True,
max_size: int = 1000,
ttl_seconds: int = 3600,
) -> None:
self.enabled = enabled
self.index_version = 1
self._lock = threading.RLock()
self.embedding_cache = LRUTTLCache(max_size=max_size * 2, ttl_seconds=ttl_seconds)
self.retrieval_cache = LRUTTLCache(max_size=max_size, ttl_seconds=ttl_seconds)
self.reranker_cache = LRUTTLCache(max_size=max_size * 4, ttl_seconds=ttl_seconds)
@classmethod
def get_instance(cls) -> CacheManager:
if cls._instance is None:
with cls._singleton_lock:
if cls._instance is None:
cls._instance = cls()
return cls._instance
@staticmethod
def _hash_str(val: str) -> str:
return hashlib.sha256(val.encode("utf-8")).hexdigest()
def get_embedding(self, model_name: str, text: str) -> Optional[List[float]]:
if not self.enabled:
return None
key = self._hash_str(f"emb:{model_name}:{text}")
return self.embedding_cache.get(key)
def set_embedding(self, model_name: str, text: str, embedding: List[float]) -> None:
if not self.enabled:
return
key = self._hash_str(f"emb:{model_name}:{text}")
self.embedding_cache.set(key, embedding)
def get_retrieval(
self,
query: str,
filter_repr: Any,
top_k: int,
mode: str,
) -> Optional[Any]:
if not self.enabled:
return None
filter_str = str(filter_repr)
key = self._hash_str(f"ret:{self.index_version}:{mode}:{top_k}:{query}:{filter_str}")
return self.retrieval_cache.get(key)
def set_retrieval(
self,
query: str,
filter_repr: Any,
top_k: int,
mode: str,
results: Any,
) -> None:
if not self.enabled:
return
filter_str = str(filter_repr)
key = self._hash_str(f"ret:{self.index_version}:{mode}:{top_k}:{query}:{filter_str}")
self.retrieval_cache.set(key, results)
def get_rerank_score(self, model_name: str, query: str, text: str) -> Optional[float]:
if not self.enabled:
return None
key = self._hash_str(f"rrk:{model_name}:{query}:{text[:200]}")
return self.reranker_cache.get(key)
def set_rerank_score(self, model_name: str, query: str, text: str, score: float) -> None:
if not self.enabled:
return
key = self._hash_str(f"rrk:{model_name}:{query}:{text[:200]}")
self.reranker_cache.set(key, float(score))
def bump_index_version(self) -> int:
"""Invalidate all cached retrieval queries after an index modification."""
with self._lock:
self.index_version += 1
self.retrieval_cache.clear()
logger.info(f"Bumped index version to {self.index_version}; cleared retrieval cache.")
return self.index_version
def clear(self) -> None:
"""Clear all in-memory caches."""
with self._lock:
self.embedding_cache.clear()
self.retrieval_cache.clear()
self.reranker_cache.clear()
def get_stats(self) -> Dict[str, Any]:
return {
"enabled": self.enabled,
"index_version": self.index_version,
"embedding": self.embedding_cache.stats(),
"retrieval": self.retrieval_cache.stats(),
"reranker": self.reranker_cache.stats(),
}
|