kg_agent/knowledge_v2/cache_manager.py

"""
缓存管理模块
提供统一的缓存读写接口，支持基于问题的分级缓存
"""

import os
import json
import hashlib
from typing import Any, Optional
from loguru import logger


class CacheManager:
    """缓存管理器"""
    
    def __init__(self, base_cache_dir: str = None):
        """
        初始化缓存管理器
        
        Args:
            base_cache_dir: 缓存根目录，默认为当前目录下的 .cache
        """
        if base_cache_dir is None:
            current_dir = os.path.dirname(os.path.abspath(__file__))
            base_cache_dir = os.path.join(current_dir, '.cache')
        
        self.base_cache_dir = base_cache_dir
        os.makedirs(base_cache_dir, exist_ok=True)
        logger.info(f"缓存管理器初始化，缓存目录: {base_cache_dir}")
    
    def _get_question_hash(self, question: str) -> str:
        """
        获取问题的hash值，用作文件夹名
        
        Args:
            question: 问题文本
            
        Returns:
            str: hash值（MD5的前12位）
        """
        return hashlib.md5(question.encode('utf-8')).hexdigest()[:12]
    
    def _get_cache_path(self, question: str, cache_type: str, filename: str) -> str:
        """
        获取缓存文件的完整路径
        
        Args:
            question: 问题文本
            cache_type: 缓存类型（如 'function_knowledge', 'llm_search', 'multi_search'）
            filename: 缓存文件名
            
        Returns:
            str: 缓存文件完整路径
        """
        question_hash = self._get_question_hash(question)
        cache_dir = os.path.join(self.base_cache_dir, question_hash, cache_type)
        os.makedirs(cache_dir, exist_ok=True)
        
        # 同时保存原始问题文本以便查看
        question_file = os.path.join(self.base_cache_dir, question_hash, 'question.txt')
        if not os.path.exists(question_file):
            with open(question_file, 'w', encoding='utf-8') as f:
                f.write(question)
        
        return os.path.join(cache_dir, filename)
    
    def get(self, question: str, cache_type: str, filename: str) -> Optional[Any]:
        """
        读取缓存
        
        Args:
            question: 问题文本
            cache_type: 缓存类型
            filename: 缓存文件名
            
        Returns:
            缓存内容，如果缓存不存在则返回 None
        """
        cache_path = self._get_cache_path(question, cache_type, filename)
        
        if not os.path.exists(cache_path):
            logger.debug(f"缓存未命中: {cache_type}/{filename}")
            return None
        
        try:
            with open(cache_path, 'r', encoding='utf-8') as f:
                content = f.read()
            
            # 尝试解析为JSON
            if filename.endswith('.json'):
                content = json.loads(content)
            
            logger.info(f"✓ 缓存命中: {cache_type}/{filename}")
            return content
            
        except Exception as e:
            logger.error(f"读取缓存失败 {cache_type}/{filename}: {e}")
            return None
    
    def set(self, question: str, cache_type: str, filename: str, content: Any) -> bool:
        """
        写入缓存
        
        Args:
            question: 问题文本
            cache_type: 缓存类型
            filename: 缓存文件名
            content: 缓存内容
            
        Returns:
            bool: 是否写入成功
        """
        cache_path = self._get_cache_path(question, cache_type, filename)
        
        try:
            # 如果是字典或列表，转换为JSON
            if isinstance(content, (dict, list)):
                content = json.dumps(content, ensure_ascii=False, indent=2)
            
            with open(cache_path, 'w', encoding='utf-8') as f:
                f.write(str(content))
            
            logger.debug(f"缓存已保存: {cache_type}/{filename}")
            return True
            
        except Exception as e:
            logger.error(f"写入缓存失败 {cache_type}/{filename}: {e}")
            return False
    
    def clear(self, question: str = None):
        """
        清除缓存
        
        Args:
            question: 如果指定，只清除该问题的缓存；否则清除所有缓存
        """
        if question:
            question_hash = self._get_question_hash(question)
            cache_dir = os.path.join(self.base_cache_dir, question_hash)
            if os.path.exists(cache_dir):
                import shutil
                shutil.rmtree(cache_dir)
                logger.info(f"已清除问题缓存: {question[:30]}...")
        else:
            import shutil
            if os.path.exists(self.base_cache_dir):
                shutil.rmtree(self.base_cache_dir)
                os.makedirs(self.base_cache_dir)
                logger.info("已清除所有缓存")