kg_agent/knowledge_v2/multi_search_knowledge.py

'''
多渠道获取知识，当前有两个渠道 llm_search_knowledge.py 和 xhs_search_knowledge.py
1. 输入：问题
2. 判断选择哪些渠道获取知识，目录默认返回 llm_search 和 xhs_search 两个渠道
3. 根据选择的结果调用对应的渠道获取知识
4. 合并多个渠道返回知识文本，返回知识文本，使用大模型合并，prompt在 prompt/multi_search_merge_knowledge_prompt.md 中
补充：暂时将xhs_search_knowledge渠道的调用注释掉，后续完成xhs_search_knowledge的实现
'''

import os
import sys
import json
from typing import List, Dict
from loguru import logger

# 设置路径以便导入工具类
current_dir = os.path.dirname(os.path.abspath(__file__))
root_dir = os.path.dirname(current_dir)
sys.path.insert(0, root_dir)

from utils.gemini_client import generate_text
from knowledge_v2.llm_search_knowledge import get_knowledge as get_llm_knowledge
from knowledge_v2.cache_manager import CacheManager
# from knowledge_v2.xhs_search_knowledge import get_knowledge as get_xhs_knowledge

class MultiSearchKnowledge:
    """多渠道知识获取类"""
    
    def __init__(self, use_cache: bool = True):
        """
        初始化
        
        Args:
            use_cache: 是否启用缓存，默认启用
        """
        logger.info("=" * 60)
        logger.info("初始化 MultiSearchKnowledge")
        self.prompt_dir = os.path.join(current_dir, "prompt")
        self.use_cache = use_cache
        self.cache = CacheManager() if use_cache else None
        
        # 执行详情收集
        self.execution_detail = {
            "sources": {},
            "merge_detail": None,
            "execution_time": 0,
            "cache_hits": []
        }
        
        logger.info(f"缓存状态: {'启用' if use_cache else '禁用'}")
        logger.info("=" * 60)
        
    def _load_prompt(self, filename: str) -> str:
        """
        加载prompt文件内容
        
        Args:
            filename: prompt文件名
            
        Returns:
            str: prompt内容
        """
        prompt_path = os.path.join(self.prompt_dir, filename)
        
        if not os.path.exists(prompt_path):
            error_msg = f"Prompt文件不存在: {prompt_path}"
            logger.error(error_msg)
            raise FileNotFoundError(error_msg)
        
        try:
            with open(prompt_path, 'r', encoding='utf-8') as f:
                content = f.read().strip()
                if not content:
                    error_msg = f"Prompt文件内容为空: {prompt_path}"
                    logger.error(error_msg)
                    raise ValueError(error_msg)
                return content
        except Exception as e:
            logger.error(f"读取prompt文件 {filename} 失败: {e}")
            raise

    def merge_knowledge(self, question: str, knowledge_map: Dict[str, str]) -> str:
        """
        合并多个渠道的知识文本
        
        Args:
            question: 用户问题  
            knowledge_map: 渠道名到知识文本的映射
            
        Returns:
            str: 合并后的知识文本
        """
        logger.info(f"[Multi-Search] 合并多渠道知识 - {len(knowledge_map)} 个渠道")
        
        # 尝试从缓存读取
        if self.use_cache:
            cached_merged = self.cache.get(question, 'multi_search', 'merged_knowledge.txt')
            if cached_merged:
                logger.info(f"✓ 使用缓存的合并知识 (长度: {len(cached_merged)})")
                # 记录缓存命中
                self.execution_detail["merge_detail"].update({
                    "cached": True,
                    "sources_count": len(knowledge_map),
                    "result_length": len(cached_merged)
                })
                self.execution_detail["cache_hits"].append("merged_knowledge")
                return cached_merged
        
        try:
            # 过滤空文本
            valid_knowledge = {k: v for k, v in knowledge_map.items() if v and v.strip()}
            logger.info(f"  有效渠道: {list(valid_knowledge.keys())}")
            
            if not valid_knowledge:
                logger.warning("  ⚠ 所有渠道的知识文本都为空")
                return ""
            
            # 如果只有一个渠道有内容，也经过LLM整理以保证输出风格一致
            
            # 加载prompt
            prompt_template = self._load_prompt("multi_search_merge_knowledge_prompt.md")
            
            # 构建知识文本部分
            knowledge_texts_str = ""
            for source, text in valid_knowledge.items():
                knowledge_texts_str += f"【来源：{source}】\n{text}\n\n"
            
            # 填充prompt
            prompt = prompt_template.format(question=question, knowledge_texts=knowledge_texts_str)
            
            # 调用大模型
            logger.info("  → 调用Gemini合并多渠道知识...")
            merged_text = generate_text(prompt=prompt)
            
            logger.info(f"✓ 多渠道知识合并完成 (长度: {len(merged_text)})")
            
            # 记录合并详情
            self.execution_detail["merge_detail"].update({
                "cached": False,
                "prompt": prompt,
                "response": merged_text,
                "sources_count": len(knowledge_map),
                "valid_sources_count": len(valid_knowledge),
                "result_length": len(merged_text)
            })
            
            # 写入缓存
            if self.use_cache:
                self.cache.set(question, 'multi_search', 'merged_knowledge.txt', merged_text.strip())
            
            return merged_text.strip()
            
        except Exception as e:
            logger.error(f"✗ 合并知识失败: {e}")
            raise
    
    def _save_execution_detail(self, cache_key: str):
        """保存执行详情到缓存（支持合并旧记录）"""
        if not self.use_cache or not self.cache:
            return
        
        try:
            import hashlib
            question_hash = hashlib.md5(cache_key.encode('utf-8')).hexdigest()[:12]
            detail_dir = os.path.join(
                self.cache.base_cache_dir,
                question_hash,
                'multi_search'
            )
            os.makedirs(detail_dir, exist_ok=True)
            
            detail_file = os.path.join(detail_dir, 'execution_detail.json')
            
            # 准备最终要保存的数据
            final_detail = self.execution_detail.copy()
            
            # 尝试读取旧文件进行合并
            if os.path.exists(detail_file):
                try:
                    with open(detail_file, 'r', encoding='utf-8') as f:
                        old_detail = json.load(f)
                    
                    # 合并 merge_detail
                    new_merge = self.execution_detail.get("merge_detail")
                    old_merge = old_detail.get("merge_detail")
                    
                    if (new_merge and isinstance(new_merge, dict) and 
                        new_merge.get("cached") is True and 
                        old_merge and isinstance(old_merge, dict) and 
                        "prompt" in old_merge):
                        final_detail["merge_detail"] = old_merge
                        
                except Exception as e:
                    logger.warning(f"  ⚠ 读取旧详情失败: {e}")

            with open(detail_file, 'w', encoding='utf-8') as f:
                json.dump(final_detail, f, ensure_ascii=False, indent=2)
            
            logger.info(f"✓ 执行详情已保存: {detail_file}")
            
        except Exception as e:
            logger.error(f"✗ 保存执行详情失败: {e}")

    def get_knowledge(self, question: str, cache_key: str = None) -> str:
        """
        获取知识的主方法
        
        Args:
            question: 问题字符串
            cache_key: 可选的缓存键，用于与主流程共享同一缓存目录
            
        Returns:
            str: 最终的知识文本
        """
        #使用cache_key或question作为缓存键
        actual_cache_key = cache_key if cache_key is not None else question
        
        import time
        start_time = time.time()
        
        logger.info(f"{'='*60}")
        logger.info(f"Multi-Search - 开始处理问题: {question[:50]}...")
        logger.info(f"{'='*60}")
        
        # 检查整体缓存
        if self.use_cache:
            cached_final = self.cache.get(actual_cache_key, 'multi_search', 'final_knowledge.txt')
            if cached_final:
                logger.info(f"✓ 使用缓存的最终知识 (长度: {len(cached_final)})")
                logger.info(f"{'='*60}\n")
                # 记录缓存命中
                self.execution_detail["cache_hits"].append("final_knowledge")
                self.execution_detail["execution_time"] = time.time() - start_time
                self._save_execution_detail(actual_cache_key)
                return cached_final
        
        knowledge_map = {}
        
        # 1. 获取 LLM Search 知识
        try:
            logger.info("[渠道1] 调用 LLM Search...")
            llm_knowledge = get_llm_knowledge(question, cache_key=actual_cache_key)
            knowledge_map["LLM Search"] = llm_knowledge
            logger.info(f"✓ LLM Search 完成 (长度: {len(llm_knowledge)})")
            # 记录来源详情
            self.execution_detail["sources"]["llm_search"] = {
                "success": True,
                "knowledge_length": len(llm_knowledge)
            }
        except Exception as e:
            logger.error(f"✗ LLM Search 失败: {e}")
            knowledge_map["LLM Search"] = ""
            self.execution_detail["sources"]["llm_search"] = {
                "success": False,
                "error": str(e)
            }
            
        # 2. 获取 XHS Search 知识 (暂时注释)
        # try:
        #     logger.info("[渠道2] 调用 XHS Search...")
        #     xhs_knowledge = get_xhs_knowledge(question)
        #     knowledge_map["XHS Search"] = xhs_knowledge
        # except Exception as e:
        #     logger.error(f"✗ XHS Search 失败: {e}")
        #     knowledge_map["XHS Search"] = ""
        
        # 3. 合并知识
        final_knowledge = self.merge_knowledge(actual_cache_key, knowledge_map)
        
        # 保存最终缓存
        if self.use_cache and final_knowledge:
            self.cache.set(actual_cache_key, 'multi_search', 'final_knowledge.txt', final_knowledge)
        
        logger.info(f"{'='*60}")
        logger.info(f"✓ Multi-Search 完成 (最终长度: {len(final_knowledge)})")
        logger.info(f"{'='*60}\n")
        
        # 计算执行时间并保存详情
        self.execution_detail["execution_time"] = time.time() - start_time
        self._save_execution_detail(actual_cache_key)
        
        return final_knowledge

def get_knowledge(question: str, cache_key: str = None) -> str:
    """
    便捷调用函数
    
    Args:
        question: 问题
        cache_key: 可选的缓存键
    """
    agent = MultiSearchKnowledge()
    return agent.get_knowledge(question, cache_key=cache_key)

if __name__ == "__main__":
    # 测试代码
    test_question = "如何评价最近的国产3A游戏黑神话悟空？"
    try:
        result = get_knowledge(test_question)
        print("=" * 50)
        print("最终整合知识:")
        print("=" * 50)
        print(result)
    except Exception as e:
        logger.error(f"测试失败: {e}")