kg_agent/knowledge_v2/llm_search_knowledge.py

'''
基于LLM+search的知识获取模块
1. 输入：问题
2. 输出：知识文本
3. 处理流程：
- 3.1 根据问题构建query，调用大模型生成多个query，prompt 在 llm_search_generate_query_prompt.md 中
- 3.2 根据query调用 utils/qwen_client.py 的 search_and_chat 方法（使用返回中的 'content' 字段即可），获取知识文本
- 3.3 用大模型合并多个query的知识文本，prompt在 llm_search_merge_knowledge_prompt.md 中
- 3.4 返回知识文本
4. 大模型调用使用uitls/gemini_client.py 的 generate_text 方法
5. 考虑复用性，尽量把每个步骤封装在一个方法中
'''

import os
import sys
import json
from typing import List
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 utils.qwen_client import QwenClient
from knowledge_v2.cache_manager import CacheManager


class LLMSearchKnowledge:
    """基于LLM+search的知识获取类"""
    
    def __init__(self, use_cache: bool = True):
        """
        初始化
        
        Args:
            use_cache: 是否启用缓存，默认启用
        """
        logger.info("=" * 60)
        logger.info("初始化 LLMSearchKnowledge")
        self.qwen_client = QwenClient()
        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 = {
            "generate_queries": None,
            "search_results": [],
            "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内容
            
        Raises:
            FileNotFoundError: 文件不存在时抛出
            ValueError: 文件内容为空时抛出
        """
        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:
            error_msg = f"读取prompt文件 {filename} 失败: {e}"
            logger.error(error_msg)
            raise
    
    def generate_queries(self, question: str) -> List[str]:
        """
        根据问题生成多个搜索query
        
        Args:
            question: 问题字符串
            
        Returns:
            List[str]: query列表
            
        Raises:
            Exception: 生成query失败时抛出异常
        """
        logger.info(f"[步骤1] 生成搜索Query - 问题: {question[:50]}...")
        
        # 尝试从缓存读取
        if self.use_cache:
            cached_queries = self.cache.get(question, 'llm_search', 'generated_queries.json')
            if cached_queries:
                logger.info(f"✓ 使用缓存的queries: {cached_queries}")
                # 记录缓存命中
                self.execution_detail["generate_queries"].update({
                    "cached": True,
                    "queries_count": len(cached_queries)
                })
                return cached_queries
        
        try:
            # 加载prompt
            prompt_template = self._load_prompt("llm_search_generate_query_prompt.md")
            
            # 构建prompt，使用 {question} 作为占位符
            prompt = prompt_template.format(question=question)
            
            # 调用gemini生成query
            logger.info("→ 调用Gemini生成query...")
            response_text = generate_text(prompt=prompt)
            
            # 解析JSON响应
            logger.info("→ 解析生成的query...")
            try:
                # 尝试提取JSON部分（去除可能的markdown代码块标记）
                response_text = response_text.strip()
                if response_text.startswith("```json"):
                    response_text = response_text[7:]
                if response_text.startswith("```"):
                    response_text = response_text[3:]
                if response_text.endswith("```"):
                    response_text = response_text[:-3]
                response_text = response_text.strip()
                
                result = json.loads(response_text)
                queries = result.get("queries", [])
                
                if not queries:
                    raise ValueError("生成的query列表为空")
                
                logger.info(f"✓ 成功生成 {len(queries)} 个query:")
                for i, q in enumerate(queries, 1):
                    logger.info(f"  {i}. {q}")
                
                # 记录执行详情
                self.execution_detail["generate_queries"].update({
                    "cached": False,
                    "prompt": prompt,
                    "response": response_text,
                    "queries_count": len(queries),
                    "queries": queries
                })
                
                # 写入缓存
                if self.use_cache:
                    self.cache.set(question, 'llm_search', 'generated_queries.json', queries)
                
                return queries
                
            except json.JSONDecodeError as e:
                logger.error(f"✗ 解析JSON失败: {e}")
                logger.error(f"响应内容: {response_text}")
                raise ValueError(f"无法解析模型返回的JSON: {e}")
                
        except Exception as e:
            logger.error(f"✗ 生成query失败: {e}")
            raise
    
    def search_knowledge(self, question: str, query: str, query_index: int = 0) -> str:
        """
        根据单个query搜索知识
        
        Args:
            question: 原始问题（用于缓存）
            query: 搜索query
            query_index: query索引（用于缓存文件名）
            
        Returns:
            str: 搜索到的知识文本（content字段）
            
        Raises:
            Exception: 搜索失败时抛出异常
        """
        logger.info(f"  [{query_index}] 搜索Query: {query}")
        
        # 尝试从缓存读取
        if self.use_cache:
            cache_filename = f"search_result_{query_index:03d}.txt"
            cached_result = self.cache.get(question, 'llm_search/search_results', cache_filename)
            if cached_result:
                logger.info(f"  ✓ 使用缓存结果 (长度: {len(cached_result)})")
                # 记录缓存命中
                self.execution_detail["search_results"].append({
                    "query": query,
                    "query_index": query_index,
                    "cached": True,
                    "result_length": len(cached_result)
                })
                self.execution_detail["cache_hits"].append(f"search_result_{query_index:03d}")
                return cached_result
        
        try:
            # 调用qwen_client的search_and_chat方法
            logger.info(f"  → 调用搜索引擎...")
            result = self.qwen_client.search_and_chat(
                user_prompt=query,
                search_strategy="agent"
            )
            
            # 提取content字段
            knowledge_text = result.get("content", "")
            
            if not knowledge_text:
                logger.warning(f"  ⚠ query '{query}' 的搜索结果为空")
                return ""
            
            logger.info(f"  ✓ 获取知识文本 (长度: {len(knowledge_text)})")
            
            # 记录搜索结果详情
            self.execution_detail["search_results"].append({
                "query": query,
                "query_index": query_index,
                "cached": False,
                "result_length": len(knowledge_text)
            })
            
            # 写入缓存
            if self.use_cache:
                cache_filename = f"search_result_{query_index:03d}.txt"
                self.cache.set(question, 'llm_search/search_results', cache_filename, knowledge_text)
            
            return knowledge_text
            
        except Exception as e:
            logger.error(f"  ✗ 搜索知识失败，query: {query}, 错误: {e}")
            raise
    
    def search_knowledge_batch(self, question: str, queries: List[str]) -> List[str]:
        """
        批量搜索知识
        
        Args:
            question: 原始问题（用于缓存）
            queries: query列表
            
        Returns:
            List[str]: 知识文本列表
        """
        logger.info(f"[步骤2] 批量搜索 - 共 {len(queries)} 个Query")
        
        knowledge_texts = []
        for i, query in enumerate(queries, 1):
            try:
                knowledge_text = self.search_knowledge(question, query, i)
                knowledge_texts.append(knowledge_text)
            except Exception as e:
                logger.error(f"  ✗ 搜索第 {i} 个query失败，跳过: {e}")
                # 失败时添加空字符串，保持索引对应
                knowledge_texts.append("")
        
        logger.info(f"✓ 批量搜索完成，获得 {len([k for k in knowledge_texts if k])} 个有效结果")
        return knowledge_texts
    
    def merge_knowledge(self, question: str, knowledge_texts: List[str]) -> str:
        """
        合并多个知识文本
        
        Args:
            question: 原始问题（用于缓存）
            knowledge_texts: 知识文本列表
            
        Returns:
            str: 合并后的知识文本
            
        Raises:
            Exception: 合并失败时抛出异常
        """
        logger.info(f"[步骤3] 合并知识 - 共 {len(knowledge_texts)} 个文本")
        
        # 尝试从缓存读取
        if self.use_cache:
            cached_merged = self.cache.get(question, 'llm_search', 'merged_knowledge.txt')
            if cached_merged:
                logger.info(f"✓ 使用缓存的合并知识 (长度: {len(cached_merged)})")
                # 记录缓存命中
                self.execution_detail["merge_detail"].update({
                    "cached": True,
                    "knowledge_count": len(knowledge_texts),
                    "result_length": len(cached_merged)
                })
                return cached_merged
        
        try:
            # 过滤空文本
            valid_texts = [text for text in knowledge_texts if text.strip()]
            logger.info(f"  有效文本数量: {len(valid_texts)}/{len(knowledge_texts)}")
            
            if not valid_texts:
                logger.warning("  ⚠ 所有知识文本都为空，返回空字符串")
                return ""
            
            if len(valid_texts) == 1:
                logger.info("  只有一个有效知识文本，直接返回")
                result = valid_texts[0]
                if self.use_cache:
                    self.cache.set(question, 'llm_search', 'merged_knowledge.txt', result)
                return result
            
            # 加载prompt
            prompt_template = self._load_prompt("llm_search_merge_knowledge_prompt.md")
            
            # 构建prompt，将多个知识文本格式化
            knowledge_sections = []
            for i, text in enumerate(valid_texts, 1):
                knowledge_sections.append(f"【知识文本 {i}】\n{text}")
            
            knowledge_texts_str = "\n\n".join(knowledge_sections)
            prompt = prompt_template.format(knowledge_texts=knowledge_texts_str)
            
            # 调用gemini合并知识
            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,
                "knowledge_count": len(knowledge_texts),
                "result_length": len(merged_text)
            })
            
            # 写入缓存
            if self.use_cache:
                self.cache.set(question, 'llm_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):
        """
        保存执行详情到缓存（支持合并旧记录）
        
        Args:
            cache_key: 缓存键
        """
        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,
                'llm_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)
                    
                    # 1. 合并 generate_queries
                    new_gen = self.execution_detail.get("generate_queries")
                    old_gen = old_detail.get("generate_queries")
                    if (new_gen and isinstance(new_gen, dict) and 
                        new_gen.get("cached") is True and 
                        old_gen and isinstance(old_gen, dict) and 
                        "prompt" in old_gen):
                        final_detail["generate_queries"] = old_gen
                        
                    # 2. 合并 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
                        
                    # 3. 合并 search_results (列表)
                    new_results = self.execution_detail.get("search_results", [])
                    old_results = old_detail.get("search_results", [])
                    
                    if new_results and old_results:
                        merged_results = []
                        # 建立旧结果的索引：(query, index) -> item
                        old_map = {(item.get("query"), item.get("query_index")): item 
                                  for item in old_results if isinstance(item, dict)}
                        
                        for item in new_results:
                            if item.get("cached") is True:
                                key = (item.get("query"), item.get("query_index"))
                                if key in old_map:
                                    # 如果旧项包含更多信息（例如非cached状态），则使用旧项
                                    old_item = old_map[key]
                                    if old_item.get("cached") is False:
                                        merged_results.append(old_item)
                                        continue
                            merged_results.append(item)
                        final_detail["search_results"] = merged_results
                        
                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: 最终的知识文本
            
        Raises:
            Exception: 处理过程中出现错误时抛出异常
        """
        # 使用cache_key或question作为缓存键
        actual_cache_key = cache_key if cache_key is not None else question
        
        import time
        start_time = time.time()
        
        try:
            logger.info(f"{'='*60}")
            logger.info(f"LLM Search - 开始处理问题: {question[:50]}...")
            logger.info(f"{'='*60}")
            
            # 步骤1: 生成多个query
            queries = self.generate_queries(actual_cache_key)
            
            # 步骤2: 对每个query搜索知识
            knowledge_texts = self.search_knowledge_batch(actual_cache_key, queries)
            
            # 步骤3: 合并多个知识文本
            merged_knowledge = self.merge_knowledge(actual_cache_key, knowledge_texts)
            
            logger.info(f"{'='*60}")
            logger.info(f"✓ LLM Search 完成 (最终长度: {len(merged_knowledge)})")
            logger.info(f"{'='*60}\n")
            
            # 计算执行时间并保存详情
            self.execution_detail["execution_time"] = time.time() - start_time
            self._save_execution_detail(actual_cache_key)
            
            return merged_knowledge
            
        except Exception as e:
            logger.error(f"✗ 获取知识文本失败，问题: {question[:50]}..., 错误: {e}")
            # 即使失败也保存执行详情
            self.execution_detail["execution_time"] = time.time() - start_time
            self._save_execution_detail(actual_cache_key)
            raise


def get_knowledge(question: str, cache_key: str = None) -> str:
    """
    便捷函数：根据问题获取知识文本
    
    Args:
        question: 问题字符串
        cache_key: 可选的缓存键
        
    Returns:
        str: 最终的知识文本
    """
    agent = LLMSearchKnowledge()
    return agent.get_knowledge(question, cache_key=cache_key)


if __name__ == "__main__":
    # 测试代码
    test_question = "关于猫咪和墨镜的服装造型元素"
    
    try:
        result = get_knowledge(test_question)
        print("=" * 50)
        print("最终知识文本:")
        print("=" * 50)
        print(result)
    except Exception as e:
        logger.error(f"测试失败: {e}")