436 lines
15 KiB
Python
436 lines
15 KiB
Python
"""
|
||
语义缓存服务 - 基于问答历史的缓存方案
|
||
|
||
功能：
|
||
1. 记录用户通过RAG返回的对话历史（问题+答案）
|
||
2. 当用户后续提问时，基于语义先搜索缓存
|
||
3. 如果命中缓存，直接返回答案
|
||
4. 如果未命中，调用RAG，并将结果存入缓存
|
||
|
||
技术方案：
|
||
- 使用Redis存储问答对（问题文本 + 答案文本）
|
||
- 使用简单的文本相似度算法（因为没有现成的Embedding模型）
|
||
- 缓存键：rag:cache:{chat_id}:{question_hash}
|
||
- 支持多级匹配：精确匹配 → 模糊匹配
|
||
|
||
作者：AI Assistant
|
||
日期：2024
|
||
"""
|
||
|
||
from utils.log_util import logger
|
||
|
||
import hashlib
|
||
import json
|
||
import time
|
||
import re
|
||
from typing import Optional, Tuple, Dict, Any, List
|
||
from dataclasses import dataclass, asdict
|
||
from .match_service import get_match_service
|
||
|
||
|
||
@dataclass
|
||
class CacheEntry:
|
||
"""缓存条目"""
|
||
question: str
|
||
answer: str
|
||
created_at: float
|
||
hit_count: int = 0
|
||
chat_id: str = ""
|
||
|
||
def to_dict(self) -> dict:
|
||
return asdict(self)
|
||
|
||
@classmethod
|
||
def from_dict(cls, data: dict) -> 'CacheEntry':
|
||
return cls(**data)
|
||
|
||
|
||
class SemanticCacheService:
|
||
"""
|
||
语义缓存服务
|
||
|
||
支持三种匹配策略：
|
||
1. 精确匹配：问题完全相同
|
||
2. 关键语义匹配：提取问题核心关键词，判断语义相似
|
||
3. 模糊匹配：计算文本相似度
|
||
"""
|
||
|
||
# 缓存配置
|
||
CACHE_PREFIX = "rag:semantic:cache"
|
||
MAX_CACHE_SIZE = 1000 # 每个chat_id最大缓存条数
|
||
SIMILARITY_THRESHOLD = 0.6 # 相似度阈值
|
||
CACHE_TTL_HOURS = 720 # 缓存过期时间（小时，一个月）
|
||
|
||
def __init__(self, redis_client=None):
|
||
"""
|
||
初始化语义缓存服务
|
||
|
||
Args:
|
||
redis_client: Redis客户端实例，如果为None则使用全局redis
|
||
"""
|
||
self.redis = redis_client
|
||
self._default_redis = None
|
||
|
||
def _get_redis(self):
|
||
"""获取Redis客户端"""
|
||
if self.redis is not None:
|
||
return self.redis
|
||
if self._default_redis is None:
|
||
from config.get_redis import RedisUtil
|
||
# 注意：这里需要异步获取，实际使用时应该传入redis实例
|
||
logger.warning("未提供Redis客户端，语义缓存功能不可用")
|
||
return self._default_redis
|
||
|
||
def _build_cache_key(self, chat_id: str, question_hash: str) -> str:
|
||
"""构建缓存键"""
|
||
return f"{self.CACHE_PREFIX}:{chat_id}:{question_hash}"
|
||
|
||
def _normalize_question(self, question: str) -> str:
|
||
"""标准化问题文本"""
|
||
match_service = get_match_service()
|
||
return match_service.preprocess_text(question)
|
||
|
||
def _extract_keywords(self, question: str) -> List[str]:
|
||
"""提取问题关键词（用于语义匹配）"""
|
||
match_service = get_match_service()
|
||
return match_service.extract_keywords(question)
|
||
|
||
def _calculate_text_similarity(self, q1: str, q2: str) -> float:
|
||
"""计算两个问题的文本相似度"""
|
||
match_service = get_match_service()
|
||
return match_service.calculate_similarity(q1, q2)
|
||
|
||
async def lookup(
|
||
self,
|
||
chat_id: str,
|
||
question: str,
|
||
redis_client=None
|
||
) -> Optional[Tuple[str, float]]:
|
||
"""
|
||
在缓存中查找相似问题
|
||
|
||
Args:
|
||
chat_id: 聊天会话ID
|
||
question: 用户问题
|
||
redis_client: Redis客户端（可选）
|
||
|
||
Returns:
|
||
如果命中缓存，返回 (答案, 相似度)
|
||
如果未命中，返回 None
|
||
"""
|
||
redis = redis_client or self._get_redis()
|
||
if redis is None:
|
||
return None
|
||
|
||
try:
|
||
# 1. 精确匹配
|
||
normalized = question.lower().strip()
|
||
question_hash = hashlib.md5(normalized.encode('utf-8')).hexdigest()[:16]
|
||
exact_key = self._build_cache_key(chat_id, question_hash)
|
||
|
||
logger.info(f"[SemanticCache] 精确查找 | chat_id={chat_id} | question={question} | hash={question_hash} | key={exact_key}")
|
||
|
||
cached_data = await redis.get(exact_key)
|
||
if cached_data:
|
||
entry = CacheEntry.from_dict(json.loads(cached_data))
|
||
entry.hit_count += 1
|
||
# 更新命中次数（异步，不阻塞返回）
|
||
await redis.set(exact_key, json.dumps(entry.to_dict()), ex=self.CACHE_TTL_HOURS * 3600)
|
||
logger.info(f"[SemanticCache] 精确命中: {question[:30]}...")
|
||
return entry.answer, 1.0
|
||
|
||
# 2. 语义模糊匹配（扫描同chat_id的所有缓存）
|
||
# 获取该chat_id的所有缓存键
|
||
pattern = f"{self.CACHE_PREFIX}:{chat_id}:*"
|
||
cache_keys = await redis.keys(pattern)
|
||
|
||
if not cache_keys:
|
||
return None
|
||
|
||
# 遍历缓存，计算相似度
|
||
best_match = None
|
||
best_match_key = None # 记录最佳匹配的key
|
||
best_similarity = 0.0
|
||
|
||
for key in cache_keys:
|
||
# 跳过刚检查过的精确匹配键
|
||
if key == exact_key:
|
||
continue
|
||
|
||
cached_data = await redis.get(key)
|
||
if not cached_data:
|
||
continue
|
||
|
||
try:
|
||
entry = CacheEntry.from_dict(json.loads(cached_data))
|
||
similarity = self._calculate_text_similarity(question, entry.question)
|
||
|
||
if similarity > best_similarity:
|
||
best_similarity = similarity
|
||
best_match = entry
|
||
best_match_key = key # 保存正确的key
|
||
|
||
except (json.JSONDecodeError, KeyError) as e:
|
||
logger.warning(f"[SemanticCache] 解析缓存失败: {key}, {e}")
|
||
continue
|
||
|
||
# 判断是否达到相似度阈值
|
||
if best_match and best_similarity >= self.SIMILARITY_THRESHOLD:
|
||
logger.info(f"[SemanticCache] 语义命中 (相似度={best_similarity:.2f}): {question[:30]}...")
|
||
# 更新命中次数
|
||
best_match.hit_count += 1
|
||
await redis.set(best_match_key, json.dumps(best_match.to_dict()), ex=self.CACHE_TTL_HOURS * 3600)
|
||
return best_match.answer, best_similarity
|
||
|
||
return None
|
||
|
||
except Exception as e:
|
||
logger.error(f"[SemanticCache] 查找失败: {e}")
|
||
return None
|
||
|
||
async def store(
|
||
self,
|
||
chat_id: str,
|
||
question: str,
|
||
answer: str,
|
||
redis_client=None
|
||
) -> bool:
|
||
"""
|
||
将问答对存入缓存
|
||
|
||
Args:
|
||
chat_id: 聊天会话ID
|
||
question: 用户问题
|
||
answer: RAG返回的答案
|
||
redis_client: Redis客户端（可选）
|
||
|
||
Returns:
|
||
是否存储成功
|
||
"""
|
||
redis = redis_client or self._get_redis()
|
||
if redis is None:
|
||
return False
|
||
|
||
try:
|
||
if not answer or len(answer.strip()) < 10:
|
||
logger.info(f"[SemanticCache] 答案太短，不缓存: {answer[:20]}...")
|
||
return False
|
||
|
||
normalized = question.lower().strip()
|
||
question_hash = hashlib.md5(normalized.encode('utf-8')).hexdigest()[:16]
|
||
cache_key = self._build_cache_key(chat_id, question_hash)
|
||
|
||
logger.info(f"[SemanticCache] 存储缓存 | chat_id={chat_id} | question={question} | hash={question_hash} | key={cache_key}")
|
||
|
||
entry = CacheEntry(
|
||
question=question,
|
||
answer=answer,
|
||
created_at=time.time(),
|
||
hit_count=0,
|
||
chat_id=chat_id
|
||
)
|
||
|
||
existing = await redis.get(cache_key)
|
||
if existing:
|
||
logger.debug(f"[SemanticCache] 问题已存在，跳过: {question[:30]}...")
|
||
return True
|
||
|
||
await self._cleanup_old_cache(chat_id, redis)
|
||
|
||
await redis.set(
|
||
cache_key,
|
||
json.dumps(entry.to_dict()),
|
||
ex=self.CACHE_TTL_HOURS * 3600
|
||
)
|
||
|
||
logger.info(f"[SemanticCache] 已缓存 | key={cache_key} | question={question[:30]}...")
|
||
|
||
verify = await redis.get(cache_key)
|
||
logger.info(f"[SemanticCache] 存储验证 | key={cache_key} | found={verify is not None}")
|
||
return True
|
||
|
||
except Exception as e:
|
||
logger.error(f"[SemanticCache] 存储失败: {e}")
|
||
return False
|
||
|
||
async def _cleanup_old_cache(self, chat_id: str, redis_client=None) -> int:
|
||
"""
|
||
清理旧缓存（当缓存过多时）
|
||
|
||
策略：删除最久未命中（创建时间最早）的缓存
|
||
|
||
Returns:
|
||
删除的缓存数量
|
||
"""
|
||
redis = redis_client or self._get_redis()
|
||
if redis is None:
|
||
return 0
|
||
|
||
try:
|
||
pattern = f"{self.CACHE_PREFIX}:{chat_id}:*"
|
||
cache_keys = await redis.keys(pattern)
|
||
|
||
if len(cache_keys) <= self.MAX_CACHE_SIZE:
|
||
return 0
|
||
|
||
cache_info = []
|
||
for key in cache_keys:
|
||
data = await redis.get(key)
|
||
if data:
|
||
try:
|
||
entry = CacheEntry.from_dict(json.loads(data))
|
||
cache_info.append((key, entry.created_at))
|
||
except:
|
||
pass
|
||
|
||
cache_info.sort(key=lambda x: x[1])
|
||
delete_count = len(cache_keys) - self.MAX_CACHE_SIZE
|
||
|
||
deleted = 0
|
||
for key, _ in cache_info[:delete_count]:
|
||
await redis.delete(key)
|
||
deleted += 1
|
||
|
||
if deleted > 0:
|
||
logger.info(f"[SemanticCache] 清理了 {deleted} 条旧缓存")
|
||
|
||
return deleted
|
||
|
||
except Exception as e:
|
||
logger.error(f"[SemanticCache] 清理失败: {e}")
|
||
return 0
|
||
|
||
async def clear(self, chat_id: str = None, redis_client=None) -> bool:
|
||
"""
|
||
清除缓存
|
||
|
||
Args:
|
||
chat_id: 要清除的chat_id，如果为None则清除所有
|
||
redis_client: Redis客户端（可选）
|
||
|
||
Returns:
|
||
是否清除成功
|
||
"""
|
||
redis = redis_client or self._get_redis()
|
||
if redis is None:
|
||
return False
|
||
|
||
try:
|
||
if chat_id:
|
||
pattern = f"{self.CACHE_PREFIX}:{chat_id}:*"
|
||
else:
|
||
pattern = f"{self.CACHE_PREFIX}:*"
|
||
|
||
cache_keys = await redis.keys(pattern)
|
||
|
||
if cache_keys:
|
||
await redis.delete(*cache_keys)
|
||
logger.info(f"[SemanticCache] 清除 {len(cache_keys)} 条缓存")
|
||
|
||
return True
|
||
|
||
except Exception as e:
|
||
logger.error(f"[SemanticCache] 清除失败: {e}")
|
||
return False
|
||
|
||
async def get_stats(self, chat_id: str = None, redis_client=None) -> Dict[str, Any]:
|
||
"""
|
||
获取缓存统计信息
|
||
|
||
Args:
|
||
chat_id: 要统计的chat_id，如果为None则统计所有
|
||
redis_client: Redis客户端（可选）
|
||
|
||
Returns:
|
||
统计信息字典
|
||
"""
|
||
redis = redis_client or self._get_redis()
|
||
if redis is None:
|
||
return {}
|
||
|
||
try:
|
||
if chat_id:
|
||
pattern = f"{self.CACHE_PREFIX}:{chat_id}:*"
|
||
else:
|
||
pattern = f"{self.CACHE_PREFIX}:*"
|
||
|
||
cache_keys = await redis.keys(pattern)
|
||
|
||
stats = {
|
||
"total_entries": len(cache_keys),
|
||
"chat_id": chat_id or "all",
|
||
"config": {
|
||
"max_cache_size": self.MAX_CACHE_SIZE,
|
||
"similarity_threshold": self.SIMILARITY_THRESHOLD,
|
||
"cache_ttl_hours": self.CACHE_TTL_HOURS
|
||
}
|
||
}
|
||
|
||
# 统计命中次数分布
|
||
if cache_keys:
|
||
total_hits = 0
|
||
for key in cache_keys[:100]: # 只采样前100个
|
||
data = await redis.get(key)
|
||
if data:
|
||
try:
|
||
entry = CacheEntry.from_dict(json.loads(data))
|
||
total_hits += entry.hit_count
|
||
except:
|
||
pass
|
||
|
||
stats["total_hits"] = total_hits
|
||
stats["avg_hits_per_entry"] = total_hits / min(len(cache_keys), 100)
|
||
|
||
return stats
|
||
|
||
except Exception as e:
|
||
logger.error(f"[SemanticCache] 统计失败: {e}")
|
||
return {}
|
||
|
||
|
||
# 全局服务实例
|
||
_semantic_cache_service: Optional[SemanticCacheService] = None
|
||
|
||
|
||
def get_semantic_cache_service() -> SemanticCacheService:
|
||
"""获取语义缓存服务单例"""
|
||
global _semantic_cache_service
|
||
if _semantic_cache_service is None:
|
||
_semantic_cache_service = SemanticCacheService()
|
||
return _semantic_cache_service
|
||
|
||
|
||
# 便捷函数
|
||
async def lookup_question(chat_id: str, question: str, redis_client=None) -> Optional[Tuple[str, float]]:
|
||
"""查找缓存问题"""
|
||
try:
|
||
logger.info(f"[SemanticCache.lookup_question] 🔍 函数被调用 | chat_id={chat_id} | question={question} | redis_client={redis_client is not None}")
|
||
service = get_semantic_cache_service()
|
||
logger.info(f"[SemanticCache.lookup_question] 获取服务实例完成 | service={service is not None}")
|
||
result = await service.lookup(chat_id, question, redis_client)
|
||
logger.info(f"[SemanticCache.lookup_question] 返回 | result={result is not None}")
|
||
if result:
|
||
logger.info(f"[SemanticCache.lookup_question] 命中结果 | answer_length={len(result[0])} | similarity={result[1]:.2f}")
|
||
return result
|
||
except Exception as e:
|
||
logger.error(f"[SemanticCache.lookup_question] 异常: {e}")
|
||
return None
|
||
|
||
|
||
async def store_qa_pair(chat_id: str, question: str, answer: str, redis_client=None) -> bool:
|
||
"""存储问答对"""
|
||
service = get_semantic_cache_service()
|
||
return await service.store(chat_id, question, answer, redis_client)
|
||
|
||
|
||
def get_question_hash(question: str) -> str:
|
||
"""计算问题的hash值（公共函数，确保存储和查找一致）"""
|
||
normalized = question.lower().strip()
|
||
return hashlib.md5(normalized.encode('utf-8')).hexdigest()[:16]
|
||
|
||
|
||
async def clear_cache(chat_id: str = None, redis_client=None) -> bool:
|
||
"""清除缓存"""
|
||
service = get_semantic_cache_service()
|
||
return await service.clear(chat_id, redis_client)
|