# -*- coding: utf-8 -*- """ RAG 三层缓存模块 缓存层次： 1. Query Cache: 完整问答结果缓存 2. Embedding Cache: 向量化结果缓存 3. Rerank Cache: 重排序分数缓存 缓存失效：基于知识库版本号（kb_version）的自动失效机制 """ import hashlib import time import threading from dataclasses import dataclass, field from typing import Dict, List, Optional, Any, Tuple from collections import OrderedDict import logging logger = logging.getLogger(__name__) @dataclass class CacheEntry: """缓存条目""" key: str value: Any created_at: float ttl: float # 秒 hits: int = 0 kb_version: int = 0 def is_expired(self) -> bool: return time.time() - self.created_at > self.ttl @dataclass class CacheStats: """缓存统计""" total_entries: int = 0 hits: int = 0 misses: int = 0 evictions: int = 0 @property def hit_rate(self) -> float: total = self.hits + self.misses return self.hits / total if total > 0 else 0.0 class LRUCache: """线程安全的 LRU 缓存实现""" def __init__(self, max_size: int = 1000, default_ttl: float = 3600): self.max_size = max_size self.default_ttl = default_ttl self._cache: OrderedDict[str, CacheEntry] = OrderedDict() self._lock = threading.RLock() self._stats = CacheStats() def get(self, key: str) -> Optional[Any]: """获取缓存值""" with self._lock: if key not in self._cache: self._stats.misses += 1 return None entry = self._cache[key] # 检查过期 if entry.is_expired(): del self._cache[key] self._stats.misses += 1 self._stats.evictions += 1 return None # LRU 更新 self._cache.move_to_end(key) entry.hits += 1 self._stats.hits += 1 return entry.value def set(self, key: str, value: Any, ttl: float = None, kb_version: int = 0) -> None: """设置缓存值""" with self._lock: if key in self._cache: del self._cache[key] entry = CacheEntry( key=key, value=value, created_at=time.time(), ttl=ttl or self.default_ttl, kb_version=kb_version ) self._cache[key] = entry # LRU 淘汰 while len(self._cache) > self.max_size: oldest_key = next(iter(self._cache)) del self._cache[oldest_key] self._stats.evictions += 1 self._stats.total_entries = len(self._cache) def invalidate_by_version(self, kb_version: int) -> int: """失效指定版本的所有缓存""" count = 0 with self._lock: keys_to_delete = [ k for k, v in self._cache.items() if v.kb_version == kb_version ] for key in keys_to_delete: del self._cache[key] count += 1 self._stats.evictions += count self._stats.total_entries = len(self._cache) return count def clear(self) -> None: """清空缓存""" with self._lock: self._cache.clear() self._stats.total_entries = 0 def get_stats(self) -> CacheStats: """获取统计信息""" with self._lock: return self._stats class RAGCacheManager: """RAG 三层缓存管理器""" # 默认配置（可从 config 覆盖） DEFAULT_QUERY_CACHE_SIZE = 500 DEFAULT_QUERY_CACHE_TTL = 3600 # 1小时 DEFAULT_EMBEDDING_CACHE_SIZE = 2000 DEFAULT_EMBEDDING_CACHE_TTL = 86400 # 24小时 DEFAULT_RERANK_CACHE_SIZE = 1000 DEFAULT_RERANK_CACHE_TTL = 3600 # 1小时 def __init__( self, query_cache_size: int = None, query_cache_ttl: float = None, embedding_cache_size: int = None, embedding_cache_ttl: float = None, rerank_cache_size: int = None, rerank_cache_ttl: float = None, kb_versions: Dict[str, int] = None ): """初始化缓存管理器""" self.query_cache = LRUCache( max_size=query_cache_size or self.DEFAULT_QUERY_CACHE_SIZE, default_ttl=query_cache_ttl or self.DEFAULT_QUERY_CACHE_TTL ) self.embedding_cache = LRUCache( max_size=embedding_cache_size or self.DEFAULT_EMBEDDING_CACHE_SIZE, default_ttl=embedding_cache_ttl or self.DEFAULT_EMBEDDING_CACHE_TTL ) self.rerank_cache = LRUCache( max_size=rerank_cache_size or self.DEFAULT_RERANK_CACHE_SIZE, default_ttl=rerank_cache_ttl or self.DEFAULT_RERANK_CACHE_TTL ) self._kb_versions: Dict[str, int] = kb_versions or {} self._version_lock = threading.Lock() def get_kb_version(self, kb_name: str) -> int: """获取知识库当前版本号""" with self._version_lock: return self._kb_versions.get(kb_name, 0) def increment_kb_version(self, kb_name: str) -> int: """递增知识库版本号（文档更新时调用）""" with self._version_lock: old_version = self._kb_versions.get(kb_name, 0) self._kb_versions[kb_name] = old_version + 1 new_version = self._kb_versions[kb_name] # 失效旧版本缓存 self.query_cache.invalidate_by_version(old_version) self.embedding_cache.invalidate_by_version(old_version) # Rerank 缓存未携带知识库版本；文档变更后必须全量清空， # 否则相同 doc_id 可能继续复用旧文档内容对应的分数。 self.rerank_cache.clear() logger.info(f"知识库 {kb_name} 版本更新: {old_version} -> {new_version}") return new_version # ==================== Query Cache 方法 ==================== @staticmethod def _make_query_cache_key(query: str, kb_name: str, kb_version: int, doc_hash: str = "") -> str: """ 生成查询缓存 key Args: query: 查询文本 kb_name: 知识库名称 kb_version: 知识库版本号 doc_hash: 相关文档版本哈希（细粒度失效） Returns: 缓存 key """ if doc_hash: # 细粒度：只失效相关文档的缓存 return hashlib.md5( f"query:{query}:{kb_name}:{doc_hash}".encode() ).hexdigest() else: # 粗粒度：整个知识库版本变化时失效 return hashlib.md5( f"query:{query}:{kb_name}:{kb_version}".encode() ).hexdigest() def get_query_result(self, query: str, kb_name: str, doc_ids: List[str] = None) -> Optional[Dict]: """ 获取查询缓存结果 始终使用粗粒度 key（基于 kb_version），确保 GET/SET key 一致。 Args: query: 查询文本 kb_name: 知识库名称 doc_ids: 保留参数以兼容调用方签名（当前未使用） """ kb_version = self.get_kb_version(kb_name) # 使用粗粒度 key，与 SET 保持一致 key = self._make_query_cache_key(query, kb_name, kb_version) return self.query_cache.get(key) def set_query_result(self, query: str, kb_name: str, result: Dict, doc_ids: List[str] = None) -> None: """ 设置查询缓存结果 始终使用粗粒度 key（基于 kb_version），确保 GET/SET key 一致。 kb_version 在文档变更时自增，触发整个知识库的缓存失效。 Args: query: 查询文本 kb_name: 知识库名称 result: 缓存结果 doc_ids: 保留参数以兼容调用方签名（当前未使用） """ kb_version = self.get_kb_version(kb_name) # 使用与 GET 相同的粗粒度 key，确保缓存可命中 key = self._make_query_cache_key(query, kb_name, kb_version) self.query_cache.set(key, result, kb_version=kb_version) def _compute_doc_hash(self, kb_name: str, doc_ids: List[str]) -> str: """ 计算文档版本哈希 用于细粒度缓存失效：只失效相关文档变化时的缓存 """ if not doc_ids: return "" # 从文档 ID 中提取 source（文件名） sources = set() for doc_id in doc_ids: # doc_id 格式通常为 "filename_text_0" 或类似 parts = doc_id.split('_') if parts: sources.add(parts[0]) # 生成哈希 sources_str = ','.join(sorted(sources)) return hashlib.md5(f"docs:{sources_str}".encode()).hexdigest() # ==================== Embedding Cache 方法 ==================== @staticmethod def _make_embedding_key(text: str) -> str: return hashlib.md5(f"emb:{text}".encode()).hexdigest() def get_embedding(self, text: str) -> Optional[List[float]]: """获取文本的 Embedding 缓存""" key = self._make_embedding_key(text) return self.embedding_cache.get(key) def set_embedding(self, text: str, embedding: List[float], kb_version: int = 0) -> None: """设置 Embedding 缓存""" key = self._make_embedding_key(text) self.embedding_cache.set(key, embedding, kb_version=kb_version) def get_embeddings_batch(self, texts: List[str]) -> Tuple[List[Optional[List[float]]], List[int]]: """ 批量获取 Embedding Returns: (embeddings, missed_indices): 命中的 embedding 列表（未命中为 None）和未命中的索引列表 """ embeddings: List[Optional[List[float]]] = [] missed_indices: List[int] = [] for i, text in enumerate(texts): emb = self.get_embedding(text) if emb is not None: embeddings.append(emb) else: embeddings.append(None) missed_indices.append(i) return embeddings, missed_indices # ==================== Rerank Cache 方法 ==================== @staticmethod def _make_rerank_key(query: str, doc_ids: List[str]) -> str: sorted_ids = sorted(doc_ids) return hashlib.md5( f"rerank:{query}:{':'.join(sorted_ids)}".encode() ).hexdigest() def get_rerank_scores(self, query: str, doc_ids: List[str]) -> Optional[Dict[str, float]]: """获取 Rerank 分数缓存 返回 {doc_id: score} 映射（而非位置列表），调用方按当前 doc_ids 顺序查表， 避免同一组文档以不同顺序返回时分数错位。 """ key = self._make_rerank_key(query, doc_ids) return self.rerank_cache.get(key) def set_rerank_scores(self, query: str, doc_ids: List[str], scores: List[float]) -> None: """设置 Rerank 分数缓存 以 {doc_id: score} 映射存储，保证顺序无关的正确性。 """ key = self._make_rerank_key(query, doc_ids) score_map = {doc_id: float(s) for doc_id, s in zip(doc_ids, scores)} self.rerank_cache.set(key, score_map) # ==================== 统计方法 ==================== def get_all_stats(self) -> Dict[str, CacheStats]: """获取所有缓存的统计信息""" return { "query_cache": self.query_cache.get_stats(), "embedding_cache": self.embedding_cache.get_stats(), "rerank_cache": self.rerank_cache.get_stats() } def clear_all(self) -> None: """清空所有缓存""" self.query_cache.clear() self.embedding_cache.clear() self.rerank_cache.clear() logger.info("所有缓存已清空") # ==================== 全局缓存实例 ==================== _cache_manager: Optional[RAGCacheManager] = None _cache_lock = threading.Lock() def get_cache_manager() -> RAGCacheManager: """获取全局缓存管理器实例（单例模式）""" global _cache_manager if _cache_manager is None: with _cache_lock: if _cache_manager is None: # 尝试从配置加载参数 try: from config import ( QUERY_CACHE_SIZE, QUERY_CACHE_TTL, EMBEDDING_CACHE_SIZE, EMBEDDING_CACHE_TTL, RERANK_CACHE_SIZE, RERANK_CACHE_TTL ) _cache_manager = RAGCacheManager( query_cache_size=QUERY_CACHE_SIZE, query_cache_ttl=QUERY_CACHE_TTL, embedding_cache_size=EMBEDDING_CACHE_SIZE, embedding_cache_ttl=EMBEDDING_CACHE_TTL, rerank_cache_size=RERANK_CACHE_SIZE, rerank_cache_ttl=RERANK_CACHE_TTL ) except ImportError: # 使用默认配置 _cache_manager = RAGCacheManager() return _cache_manager def reset_cache_manager() -> None: """重置全局缓存管理器（主要用于测试）""" global _cache_manager with _cache_lock: if _cache_manager is not None: _cache_manager.clear_all() _cache_manager = None # ==================== 测试 ==================== if __name__ == "__main__": import sys if sys.platform == 'win32': sys.stdout.reconfigure(encoding='utf-8') print("=" * 60) print("缓存模块测试") print("=" * 60) cache = RAGCacheManager() # 测试 Query Cache print("\n1. Query Cache 测试") cache.set_query_result("什么是Python?", "public_kb", {"answer": "Python是一种编程语言"}) result = cache.get_query_result("什么是Python?", "public_kb") print(f" 缓存命中: {result}") # 测试版本号失效 print("\n2. 版本号失效测试") cache.increment_kb_version("public_kb") result = cache.get_query_result("什么是Python?", "public_kb") print(f" 版本更新后缓存失效: {result is None}") # 测试 Embedding Cache print("\n3. Embedding Cache 测试") cache.set_embedding("测试文本", [0.1, 0.2, 0.3]) emb = cache.get_embedding("测试文本") print(f" Embedding 缓存: {emb}") # 测试统计 print("\n4. 缓存统计") stats = cache.get_all_stats() for name, stat in stats.items(): print(f" {name}: hits={stat.hits}, misses={stat.misses}, hit_rate={stat.hit_rate:.2%}")