Coverage for agentos/cache/llm_cache.py: 100%

1""" 

2AgentOS v0.40 LLM Cache — 语义缓存减少API调用成本。 

3支持：精确匹配缓存、语义相似度缓存、LRU淘汰、TTL过期。 

4""" 

5 

6from __future__ import annotations 

7 

8import hashlib 

9import json 

10import time 

11from collections import OrderedDict 

12from dataclasses import dataclass, field 

13from typing import Any 

14 

15 

16@dataclass 

17class CacheEntry: 

18 """A cached LLM response with metadata. 

19 

20 Attributes: 

21 key: Cache lookup key (typically hash of prompt + model). 

22 value: Cached response content. 

23 tokens_saved: Tokens saved by serving from cache. 

24 cost_saved: Estimated cost saved. 

25 created_at: Unix timestamp of cache insertion. 

26 ttl: Time-to-live in seconds. 

27 hit_count: Number of cache hits. 

28 tags: Optional tags for cache invalidation. 

29 """ 

30 

31 key: str 

32 value: Any 

33 tokens_saved: int = 0 

34 cost_saved: float = 0.0 

35 created_at: float = field(default_factory=time.time) 

36 ttl: float = 3600 # 默认1小时 

37 hit_count: int = 0 

38 tags: list[str] = field(default_factory=list) 

39 

40 @property 

41 def expired(self) -> bool: 

42 return time.time() > self.created_at + self.ttl 

43 

44 

45class LRUCache: 

46 """LRU淘汰的内存缓存。""" 

47 

48 def __init__(self, max_size: int = 500): 

49 self._cache: OrderedDict[str, CacheEntry] = OrderedDict() 

50 self.max_size = max_size 

51 

52 def get(self, key: str) -> CacheEntry | None: 

53 entry = self._cache.get(key) 

54 if entry: 

55 if entry.expired: 

56 del self._cache[key] 

57 return None 

58 self._cache.move_to_end(key) 

59 entry.hit_count += 1 

60 return entry 

61 return None 

62 

63 def put(self, key: str, entry: CacheEntry): 

64 if len(self._cache) >= self.max_size and key not in self._cache: 

65 self._cache.popitem(last=False) # 淘汰最久未用 

66 self._cache[key] = entry 

67 self._cache.move_to_end(key) 

68 

69 def invalidate(self, key: str | None = None, tag: str | None = None): 

70 if key: 

71 self._cache.pop(key, None) 

72 elif tag: 

73 to_delete = [k for k, v in self._cache.items() if tag in v.tags] 

74 for k in to_delete: 

75 del self._cache[k] 

76 

77 def size(self) -> int: 

78 return len(self._cache) 

79 

80 def clear(self): 

81 self._cache.clear() 

82 

83 

84class SemanticCache: 

85 """语义缓存 — 基于embedding相似度的缓存匹配。""" 

86 

87 def __init__(self, similarity_threshold: float = 0.92, embedder: Any = None): 

88 self.threshold = similarity_threshold 

89 self._entries: list[tuple[list[float], CacheEntry]] = [] 

90 self._embedder = embedder # 外部注入的embedding函数 

91 self.max_entries = 200 

92 

93 def _embed(self, text: str) -> list[float]: 

94 if self._embedder: 

95 return self._embedder(text) 

96 # 默认回退：简易TF-IDF风格hash 

97 tokens = text.lower().split() 

98 tf = {} 

99 for t in tokens: 

100 tf[t] = tf.get(t, 0) + 1 

101 vec = [hash(w) % 100 / 100.0 * tf.get(w, 0) for w in sorted(set(tokens))[:128]] 

102 return vec[:64] if len(vec) > 64 else vec + [0.0] * (64 - len(vec)) 

103 

104 @staticmethod 

105 def cosine_sim(a: list[float], b: list[float]) -> float: 

106 if not a or not b: 

107 return 0.0 

108 dot = sum(x * y for x, y in zip(a, b)) 

109 norm_a = sum(x**2 for x in a) ** 0.5 

110 norm_b = sum(x**2 for x in b) ** 0.5 

111 if norm_a == 0 or norm_b == 0: 

112 return 0.0 

113 return dot / (norm_a * norm_b) 

114 

115 def search(self, query: str) -> CacheEntry | None: 

116 query_vec = self._embed(query) 

117 best_sim = 0.0 

118 best_entry = None 

119 for cached_vec, entry in self._entries: 

120 if entry.expired: 

121 continue 

122 sim = self.cosine_sim(query_vec, cached_vec) 

123 if sim > best_sim: 

124 best_sim = sim 

125 best_entry = entry 

126 if best_sim >= self.threshold and best_entry: 

127 best_entry.hit_count += 1 

128 return best_entry 

129 return None 

130 

131 def add(self, query: str, entry: CacheEntry): 

132 vec = self._embed(query) 

133 self._entries.append((vec, entry)) 

134 if len(self._entries) > self.max_entries: 

135 self._entries = self._entries[-self.max_entries :] 

136 

137 def clear(self): 

138 self._entries.clear() 

139 

140 

141@dataclass 

142class CacheStats: 

143 """缓存统计。""" 

144 

145 total_requests: int = 0 

146 hits: int = 0 

147 misses: int = 0 

148 tokens_saved: int = 0 

149 cost_saved: float = 0.0 

150 exact_hits: int = 0 

151 semantic_hits: int = 0 

152 

153 @property 

154 def hit_rate(self) -> float: 

155 if self.total_requests == 0: 

156 return 0.0 

157 return self.hits / self.total_requests 

158 

159 

160class LLMCache: 

161 """ 

162 LLM响应缓存 — 减少API调用成本。 

163 

164 三层策略: 

165 1. 精确匹配缓存 (LRU + TTL) 

166 2. 语义相似度缓存 

167 3. 透传 (无缓存命中) 

168 """ 

169 

170 def __init__( 

171 self, lru_size: int = 500, semantic_threshold: float = 0.92, enable_semantic: bool = True 

172 ): 

173 self.lru = LRUCache(max_size=lru_size) 

174 self.semantic = ( 

175 SemanticCache(similarity_threshold=semantic_threshold) if enable_semantic else None 

176 ) 

177 self.stats = CacheStats() 

178 

179 @staticmethod 

180 def _hash_key(prompt: str, model: str = "", **kwargs) -> str: 

181 payload = prompt + model + json.dumps(kwargs, sort_keys=True) 

182 return hashlib.sha256(payload.encode()).hexdigest()[:32] 

183 

184 def get(self, prompt: str, model: str = "", **kwargs) -> Any | None: 

185 self.stats.total_requests += 1 

186 

187 # 1. 精确匹配 

188 exact_key = self._hash_key(prompt, model, **kwargs) 

189 entry = self.lru.get(exact_key) 

190 if entry: 

191 self.stats.hits += 1 

192 self.stats.exact_hits += 1 

193 self.stats.tokens_saved += entry.tokens_saved 

194 self.stats.cost_saved += entry.cost_saved 

195 return entry.value 

196 

197 # 2. 语义匹配 

198 if self.semantic: 

199 entry = self.semantic.search(prompt) 

200 if entry: 

201 self.stats.hits += 1 

202 self.stats.semantic_hits += 1 

203 self.stats.tokens_saved += entry.tokens_saved 

204 self.stats.cost_saved += entry.cost_saved 

205 return entry.value 

206 

207 self.stats.misses += 1 

208 return None 

209 

210 def set( 

211 self, 

212 prompt: str, 

213 value: Any, 

214 model: str = "", 

215 tokens: int = 0, 

216 cost: float = 0.0, 

217 ttl: float = 3600, 

218 **kwargs, 

219 ): 

220 exact_key = self._hash_key(prompt, model, **kwargs) 

221 entry = CacheEntry( 

222 key=exact_key, value=value, tokens_saved=tokens, cost_saved=cost, ttl=ttl 

223 ) 

224 self.lru.put(exact_key, entry) 

225 

226 if self.semantic: 

227 self.semantic.add(prompt, entry) 

228 

229 def invalidate(self, key: str = "", tag: str = ""): 

230 self.lru.invalidate(key=key or None, tag=tag or None) 

231 

232 def clear(self): 

233 self.lru.clear() 

234 if self.semantic: 

235 self.semantic.clear() 

236 

237 def snapshot(self) -> dict: 

238 return { 

239 "lru_entries": self.lru.size(), 

240 "semantic_entries": len(self.semantic._entries) if self.semantic else 0, 

241 "hit_rate": f"{self.stats.hit_rate:.1%}", 

242 "tokens_saved": self.stats.tokens_saved, 

243 "cost_saved": f"${self.stats.cost_saved:.4f}", 

244 "total_requests": self.stats.total_requests, 

245 "exact_hits": self.stats.exact_hits, 

246 "semantic_hits": self.stats.semantic_hits, 

247 }