Mofox-Core/src/common/cache_manager.py

import time
import json
import hashlib
import inspect
import numpy as np
import faiss
import chromadb
from typing import Any, Dict, Optional
from src.common.logger import get_logger
from src.llm_models.utils_model import LLMRequest
from src.config.config import global_config, model_config
from src.common.database.sqlalchemy_models import CacheEntries
from src.common.database.sqlalchemy_database_api import db_query, db_save

logger = get_logger("cache_manager")

class CacheManager:
    """
    一个支持分层和语义缓存的通用工具缓存管理器。
    采用单例模式，确保在整个应用中只有一个缓存实例。
    L1缓存: 内存字典 (KV) + FAISS (Vector)。
    L2缓存: 数据库 (KV) + ChromaDB (Vector)。
    """
    _instance = None

    def __new__(cls, *args, **kwargs):
        if not cls._instance:
            cls._instance = super(CacheManager, cls).__new__(cls)
        return cls._instance

    def __init__(self, default_ttl: int = 3600, chroma_path: str = "data/chroma_db"):
        """
        初始化缓存管理器。
        """
        if not hasattr(self, '_initialized'):
            self.default_ttl = default_ttl
            
            # L1 缓存 (内存)
            self.l1_kv_cache: Dict[str, Dict[str, Any]] = {}
            embedding_dim = global_config.lpmm_knowledge.embedding_dimension
            self.l1_vector_index = faiss.IndexFlatIP(embedding_dim)
            self.l1_vector_id_to_key: Dict[int, str] = {}
            
            # 语义缓存 (ChromaDB)
                
            self.chroma_client = chromadb.PersistentClient(path=chroma_path)
            self.chroma_collection = self.chroma_client.get_or_create_collection(name="semantic_cache")

            
            # 嵌入模型
            self.embedding_model = LLMRequest(model_config.model_task_config.embedding)

            self._initialized = True
            logger.info("缓存管理器已初始化: L1 (内存+FAISS), L2 (数据库+ChromaDB)")

    def _validate_embedding(self, embedding_result: Any) -> Optional[np.ndarray]:
        """
        验证和标准化嵌入向量格式
        """
        try:
            if embedding_result is None:
                return None
                
            # 确保embedding_result是一维数组或列表
            if isinstance(embedding_result, (list, tuple, np.ndarray)):
                # 转换为numpy数组进行处理
                embedding_array = np.array(embedding_result)
                
                # 如果是多维数组，展平它
                if embedding_array.ndim > 1:
                    embedding_array = embedding_array.flatten()
                
                # 检查维度是否符合预期
                expected_dim = global_config.lpmm_knowledge.embedding_dimension
                if embedding_array.shape[0] != expected_dim:
                    logger.warning(f"嵌入向量维度不匹配: 期望 {expected_dim}, 实际 {embedding_array.shape[0]}")
                    return None
                
                # 检查是否包含有效的数值
                if np.isnan(embedding_array).any() or np.isinf(embedding_array).any():
                    logger.warning("嵌入向量包含无效的数值 (NaN 或 Inf)")
                    return None
                
                return embedding_array.astype('float32')
            else:
                logger.warning(f"嵌入结果格式不支持: {type(embedding_result)}")
                return None
                
        except Exception as e:
            logger.error(f"验证嵌入向量时发生错误: {e}")
            return None

    def _generate_key(self, tool_name: str, function_args: Dict[str, Any], tool_class: Any) -> str:
        """生成确定性的缓存键，包含代码哈希以实现自动失效。"""
        try:
            source_code = inspect.getsource(tool_class)
            code_hash = hashlib.md5(source_code.encode()).hexdigest()
        except (TypeError, OSError) as e:
            code_hash = "unknown"
            logger.warning(f"无法获取 {tool_class.__name__} 的源代码，代码哈希将为 'unknown'。错误: {e}")
        try:
            sorted_args = json.dumps(function_args, sort_keys=True)
        except TypeError:
            sorted_args = repr(sorted(function_args.items()))
        return f"{tool_name}::{sorted_args}::{code_hash}"

    async def get(self, tool_name: str, function_args: Dict[str, Any], tool_class: Any, semantic_query: Optional[str] = None) -> Optional[Any]:
        """
        从缓存获取结果，查询顺序: L1-KV -> L1-Vector -> L2-KV -> L2-Vector。
        """
        # 步骤 1: L1 精确缓存查询
        key = self._generate_key(tool_name, function_args, tool_class)
        logger.debug(f"生成的缓存键: {key}")
        if semantic_query:
            logger.debug(f"使用的语义查询: '{semantic_query}'")

        if key in self.l1_kv_cache:
            entry = self.l1_kv_cache[key]
            if time.time() < entry["expires_at"]:
                logger.info(f"命中L1键值缓存: {key}")
                return entry["data"]
            else:
                del self.l1_kv_cache[key]

        # 步骤 2: L1/L2 语义和L2精确缓存查询
        query_embedding = None
        if semantic_query and self.embedding_model:
            embedding_result = await self.embedding_model.get_embedding(semantic_query)
            if embedding_result:
                validated_embedding = self._validate_embedding(embedding_result)
                if validated_embedding is not None:
                    query_embedding = np.array([validated_embedding], dtype='float32')

        # 步骤 2a: L1 语义缓存 (FAISS)
        if query_embedding is not None and self.l1_vector_index.ntotal > 0:
            faiss.normalize_L2(query_embedding)
            distances, indices = self.l1_vector_index.search(query_embedding, 1)
            if indices.size > 0 and distances[0][0] > 0.75: # IP 越大越相似
                hit_index = indices[0][0]
                l1_hit_key = self.l1_vector_id_to_key.get(hit_index)
                if l1_hit_key and l1_hit_key in self.l1_kv_cache:
                    logger.info(f"命中L1语义缓存: {l1_hit_key}")
                    return self.l1_kv_cache[l1_hit_key]["data"]

        # 步骤 2b: L2 精确缓存 (数据库)
        cache_results = await db_query(
            model_class=CacheEntries,
            query_type="get",
            filters={"cache_key": key},
            single_result=True
        )
        
        if cache_results:
            expires_at = cache_results["expires_at"]
            if time.time() < expires_at:
                logger.info(f"命中L2键值缓存: {key}")
                data = json.loads(cache_results["cache_value"])
                
                # 更新访问统计
                await db_query(
                    model_class=CacheEntries,
                    query_type="update",
                    filters={"cache_key": key},
                    data={
                        "last_accessed": time.time(),
                        "access_count": cache_results["access_count"] + 1
                    }
                )
                
                # 回填 L1
                self.l1_kv_cache[key] = {"data": data, "expires_at": expires_at}
                return data
            else:
                # 删除过期的缓存条目
                await db_query(
                    model_class=CacheEntries,
                    query_type="delete",
                    filters={"cache_key": key}
                )

        # 步骤 2c: L2 语义缓存 (ChromaDB)
        if query_embedding is not None and self.chroma_collection:
            try:
                results = self.chroma_collection.query(query_embeddings=query_embedding.tolist(), n_results=1)
                if results and results['ids'] and results['ids'][0]:
                    distance = results['distances'][0][0] if results['distances'] and results['distances'][0] else 'N/A'
                    logger.debug(f"L2语义搜索找到最相似的结果: id={results['ids'][0]}, 距离={distance}")
                    if distance != 'N/A' and distance < 0.75:
                        l2_hit_key = results['ids'][0][0] if isinstance(results['ids'][0], list) else results['ids'][0]
                        logger.info(f"命中L2语义缓存: key='{l2_hit_key}', 距离={distance:.4f}")
                        
                        # 从数据库获取缓存数据
                        semantic_cache_results = await db_query(
                            model_class=CacheEntries,
                            query_type="get",
                            filters={"cache_key": l2_hit_key},
                            single_result=True
                        )
                        
                        if semantic_cache_results:
                            expires_at = semantic_cache_results["expires_at"]
                            if time.time() < expires_at:
                                data = json.loads(semantic_cache_results["cache_value"])
                                logger.debug(f"L2语义缓存返回的数据: {data}")
                                
                                # 回填 L1
                                self.l1_kv_cache[key] = {"data": data, "expires_at": expires_at}
                                if query_embedding is not None:
                                    try:
                                        new_id = self.l1_vector_index.ntotal
                                        faiss.normalize_L2(query_embedding)
                                        self.l1_vector_index.add(x=query_embedding)
                                        self.l1_vector_id_to_key[new_id] = key
                                    except Exception as e:
                                        logger.error(f"回填L1向量索引时发生错误: {e}")
                                return data
            except Exception as e:
                logger.warning(f"ChromaDB查询失败: {e}")

        logger.debug(f"缓存未命中: {key}")
        return None

    async def set(self, tool_name: str, function_args: Dict[str, Any], tool_class: Any, data: Any, ttl: Optional[int] = None, semantic_query: Optional[str] = None):
        """将结果存入所有缓存层。"""
        if ttl is None:
            ttl = self.default_ttl
        if ttl <= 0:
            return

        key = self._generate_key(tool_name, function_args, tool_class)
        expires_at = time.time() + ttl
        
        # 写入 L1
        self.l1_kv_cache[key] = {"data": data, "expires_at": expires_at}

        # 写入 L2 (数据库)
        cache_data = {
            "cache_key": key,
            "cache_value": json.dumps(data, ensure_ascii=False),
            "expires_at": expires_at,
            "tool_name": tool_name,
            "created_at": time.time(),
            "last_accessed": time.time(),
            "access_count": 1
        }
        
        await db_save(
            model_class=CacheEntries,
            data=cache_data,
            key_field="cache_key",
            key_value=key
        )

        # 写入语义缓存
        if semantic_query and self.embedding_model and self.chroma_collection:
            try:
                embedding_result = await self.embedding_model.get_embedding(semantic_query)
                if embedding_result:
                    validated_embedding = self._validate_embedding(embedding_result)
                    if validated_embedding is not None:
                        embedding = np.array([validated_embedding], dtype='float32')
                        # 写入 L1 Vector
                        new_id = self.l1_vector_index.ntotal
                        faiss.normalize_L2(embedding)
                        self.l1_vector_index.add(x=embedding)
                        self.l1_vector_id_to_key[new_id] = key
                        # 写入 L2 Vector
                        self.chroma_collection.add(embeddings=embedding.tolist(), ids=[key])
            except Exception as e:
                logger.warning(f"语义缓存写入失败: {e}")

        logger.info(f"已缓存条目: {key}, TTL: {ttl}s")

    def clear_l1(self):
        """清空L1缓存。"""
        self.l1_kv_cache.clear()
        self.l1_vector_index.reset()
        self.l1_vector_id_to_key.clear()
        logger.info("L1 (内存+FAISS) 缓存已清空。")

    async def clear_l2(self):
        """清空L2缓存。"""
        # 清空数据库缓存
        await db_query(
            model_class=CacheEntries,
            query_type="delete",
            filters={}  # 删除所有记录
        )
        
        # 清空ChromaDB
        if self.chroma_collection:
            try:
                self.chroma_client.delete_collection(name="semantic_cache")
                self.chroma_collection = self.chroma_client.get_or_create_collection(name="semantic_cache")
            except Exception as e:
                logger.warning(f"清空ChromaDB失败: {e}")
        
        logger.info("L2 (数据库 & ChromaDB) 缓存已清空。")

    async def clear_all(self):
        """清空所有缓存。"""
        self.clear_l1()
        await self.clear_l2()
        logger.info("所有缓存层级已清空。")

    async def clean_expired(self):
        """清理过期的缓存条目"""
        current_time = time.time()
        
        # 清理L1过期条目
        expired_keys = []
        for key, entry in self.l1_kv_cache.items():
            if current_time >= entry["expires_at"]:
                expired_keys.append(key)
        
        for key in expired_keys:
            del self.l1_kv_cache[key]
        
        # 清理L2过期条目
        await db_query(
            model_class=CacheEntries,
            query_type="delete",
            filters={"expires_at": {"$lt": current_time}}
        )
        
        if expired_keys:
            logger.info(f"清理了 {len(expired_keys)} 个过期的L1缓存条目")

# 全局实例
tool_cache = CacheManager()