Spaces:
Running
Running
File size: 5,084 Bytes
c650d43 | 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 | import math
import threading
import time
from typing import Any, Dict, List, Optional, Tuple
from src.embeddings.embedder import embed_texts
def normalize_query(query: str) -> str:
return " ".join((query or "").strip().lower().split())
class ExactMatchCache:
def __init__(self, ttl_seconds: int = 3600):
self.ttl_seconds = max(1, int(ttl_seconds))
self._store: Dict[str, Dict[str, Any]] = {}
self._lock = threading.Lock()
self._hits = 0
self._misses = 0
def _is_expired(self, created_at: float) -> bool:
return (time.time() - created_at) > self.ttl_seconds
def get(self, query: str) -> Optional[Dict[str, Any]]:
key = normalize_query(query)
with self._lock:
value = self._store.get(key)
if not value:
self._misses += 1
return None
if self._is_expired(value["timestamp"]):
self._store.pop(key, None)
self._misses += 1
return None
self._hits += 1
return dict(value)
def set(self, query: str, answer: str, docs: Optional[List[Dict[str, Any]]] = None):
key = normalize_query(query)
payload = {
"answer": answer,
"docs": list(docs or []),
"timestamp": time.time(),
}
with self._lock:
self._store[key] = payload
def invalidate(self, query: Optional[str] = None):
with self._lock:
if query is None:
self._store.clear()
return
self._store.pop(normalize_query(query), None)
def stats(self) -> Dict[str, Any]:
with self._lock:
total = self._hits + self._misses
ratio = (self._hits / total) if total else 0.0
return {
"entries": len(self._store),
"hits": self._hits,
"misses": self._misses,
"hit_ratio": round(ratio, 4),
"ttl_seconds": self.ttl_seconds,
}
class SemanticCache:
def __init__(self, ttl_seconds: int = 3600, similarity_threshold: float = 0.9):
self.ttl_seconds = max(1, int(ttl_seconds))
self.similarity_threshold = float(similarity_threshold)
self._items: List[Dict[str, Any]] = []
self._lock = threading.Lock()
self._hits = 0
self._misses = 0
def _is_expired(self, created_at: float) -> bool:
return (time.time() - created_at) > self.ttl_seconds
def _cosine(self, a: List[float], b: List[float]) -> float:
if not a or not b or len(a) != len(b):
return 0.0
dot = sum(x * y for x, y in zip(a, b))
norm_a = math.sqrt(sum(x * x for x in a))
norm_b = math.sqrt(sum(y * y for y in b))
if norm_a == 0.0 or norm_b == 0.0:
return 0.0
return dot / (norm_a * norm_b)
def _prune_expired(self):
self._items = [
item for item in self._items if not self._is_expired(item["timestamp"])
]
def get(self, query: str) -> Optional[Dict[str, Any]]:
normalized = normalize_query(query)
if not normalized:
self._misses += 1
return None
query_vector = embed_texts([normalized], batch_size=1)[0]
with self._lock:
self._prune_expired()
best: Optional[Tuple[float, Dict[str, Any]]] = None
for item in self._items:
score = self._cosine(query_vector, item["vector"])
if best is None or score > best[0]:
best = (score, item)
if best is None or best[0] < self.similarity_threshold:
self._misses += 1
return None
self._hits += 1
result = dict(best[1]["value"])
result["semantic_similarity"] = round(best[0], 4)
return result
def set(self, query: str, answer: str, docs: Optional[List[Dict[str, Any]]] = None):
normalized = normalize_query(query)
if not normalized:
return
vector = embed_texts([normalized], batch_size=1)[0]
item = {
"query": normalized,
"vector": vector,
"value": {
"answer": answer,
"docs": list(docs or []),
},
"timestamp": time.time(),
}
with self._lock:
self._prune_expired()
self._items.append(item)
def invalidate(self):
with self._lock:
self._items = []
def stats(self) -> Dict[str, Any]:
with self._lock:
total = self._hits + self._misses
ratio = (self._hits / total) if total else 0.0
return {
"entries": len(self._items),
"hits": self._hits,
"misses": self._misses,
"hit_ratio": round(ratio, 4),
"ttl_seconds": self.ttl_seconds,
"similarity_threshold": self.similarity_threshold,
}
|