kg_agent/knowledge_v2/function_knowledge.py

'''
方法知识获取模块
1. 输入：问题 + 帖子信息 + 账号人设信息
2. 将输入的问题转化成query，调用大模型，prompt在 function_knowledge_generate_query_prompt.md 中
3. 从已有方法工具库中尝试选择合适的方法工具（调用大模型执行，prompt在 function_knowledge_select_tools_prompt.md 中），如果有，则返回选择的方法工具，否则：
    - 调用 multi_search_knowledge.py 获取知识
    - 返回新的方法工具知识
    - 异步从新方法知识中获取新工具(调用大模型执行，prompt在 function_knowledge_generate_new_tool_prompt.md 中)，调用工具库系统，接入新的工具
4. 调用选择的方法工具执行验证，返回工具执行结果
'''

import os
import sys
import json
import threading
from loguru import logger
import re

# 设置路径以便导入工具类
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.tools_library import call_tool, save_tool_info, get_all_tool_infos, get_tool_info, get_tool_params
from knowledge_v2.multi_search_knowledge import get_knowledge as get_multi_search_knowledge
from knowledge_v2.cache_manager import CacheManager

class FunctionKnowledge:
    """方法知识获取类"""
    
    def __init__(self, use_cache: bool = True):
        """
        初始化
        
        Args:
            use_cache: 是否启用缓存，默认启用
        """
        logger.info("=" * 80)
        logger.info("初始化 FunctionKnowledge - 方法知识获取入口")
        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_query": {},
            "select_tool": {},
            "extract_params": {},
            "execution_time": 0,
            "cache_hits": []
        }
        
        logger.info(f"缓存状态: {'启用' if use_cache else '禁用'}")
        logger.info("=" * 80)
        
    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,
                'function_knowledge'
            )
            os.makedirs(detail_dir, exist_ok=True)
            
            detail_file = os.path.join(detail_dir, 'execution_detail.json')
            with open(detail_file, 'w', encoding='utf-8') as f:
                json.dump(self.execution_detail, f, ensure_ascii=False, indent=2)
            
            logger.info(f"✓ 执行详情已保存: {detail_file}")
            
        except Exception as e:
            logger.error(f"✗ 保存执行详情失败: {e}")

    def _load_prompt(self, filename: str) -> str:
        """加载prompt文件内容"""
        prompt_path = os.path.join(self.prompt_dir, filename)
        if not os.path.exists(prompt_path):
            raise FileNotFoundError(f"Prompt文件不存在: {prompt_path}")
        with open(prompt_path, 'r', encoding='utf-8') as f:
            return f.read().strip()

    def generate_query(self, question: str, post_info: str, persona_info: str) -> str:
        """
        生成查询语句
        
        Returns:
            str: 生成的查询语句
        """
        logger.info(f"[步骤1] 生成Query...")
        
        # 组合问题的唯一标识
        combined_question = f"{question}||{post_info}||{persona_info}"
        
        try:
            prompt_template = self._load_prompt("function_generate_query_prompt.md")
            prompt = prompt_template.format(
                question=question,
                post_info=post_info,
                persona_info=persona_info
            )

            # 尝试从缓存读取
            if self.use_cache:
                cached_query = self.cache.get(combined_question, 'function_knowledge', 'generated_query.txt')
                if cached_query:
                    logger.info(f"✓ 使用缓存的Query: {cached_query}")
                    # 记录缓存命中
                    self.execution_detail["generate_query"].update({"cached": True, "query": cached_query})
                    return cached_query
            
            logger.info("→ 调用Gemini生成Query...")
            query = generate_text(prompt=prompt)
            query = query.strip()
            
            logger.info(f"✓ 生成Query: {query}")
            
            # 写入缓存
            if self.use_cache:
                self.cache.set(combined_question, 'function_knowledge', 'generated_query.txt', query)
            
            # 记录详情
            self.execution_detail["generate_query"] = {
                "cached": False,
                "prompt": prompt,
                "response": query,
                "query": query
            }
            
            return query
        except Exception as e:
            logger.error(f"✗ 生成Query失败: {e}")
            return question  # 降级使用原问题

    def select_tool(self, combined_question: str, query: str) -> str:
        """
        选择合适的工具
        
        Returns:
            str: 工具名称，如果没有合适的工具则返回"None"
        """
        logger.info(f"[步骤2] 选择工具...")
        
        try:
            all_tool_infos = self._load_prompt("all_tools_infos.md")
            if not all_tool_infos:
                logger.info("  工具库为空，无可用工具")
                return "None"
            
            prompt_template = self._load_prompt("function_knowledge_select_tools_prompt.md")
            prompt = prompt_template.replace("{all_tool_infos}", all_tool_infos).replace("query", query)

            # 尝试从缓存读取
            if self.use_cache:
                cached_tool = self.cache.get(combined_question, 'function_knowledge', 'selected_tool.txt')
                if cached_tool:
                    logger.info(f"✓ 使用缓存的工具: {cached_tool}")
                    # 记录缓存命中
                    self.execution_detail["select_tool"].update({
                        "cached": True,
                        "response": json.loads(cached_tool),
                        "prompt": prompt,
                    })
                    return json.loads(cached_tool)
            
            logger.info("→ 调用Gemini选择工具...")
            result = generate_text(prompt=prompt)
            result = self.extract_and_validate_json(result)
            if not result:
                logger.error("✗ 选择工具失败: 无法提取有效JSON")
                return "None"

            result_json = json.loads(result)

            logger.info(f"✓ 选择结果: {result_json.get('工具名', 'None')}")
            
            # 写入缓存
            if self.use_cache:
                self.cache.set(combined_question, 'function_knowledge', 'selected_tool.txt', result)
            
            # 记录详情
            self.execution_detail["select_tool"] = {
                "cached": False,
                "prompt": prompt,
                "response": result_json,
            }
            
            return result_json
        except Exception as e:
            logger.error(f"✗ 选择工具失败: {e}")
            return "None"

    def extract_and_validate_json(self, text: str):
        """
        从字符串中提取 JSON 部分，并返回标准的 JSON 字符串。
        如果无法提取或解析失败，返回 None (或者你可以改为抛出异常)。
        """
        # 1. 使用正则表达式寻找最大的 JSON 块
        # r"(\{[\s\S]*\}|\[[\s\S]*\])" 的含义：
        # - \{[\s\S]*\} : 匹配以 { 开头，} 结尾的最长字符串（[\s\S] 包含换行符）
        # - | : 或者
        # - \[[\s\S]*\] : 匹配以 [ 开头，] 结尾的最长字符串（处理 JSON 数组）
        match = re.search(r"(\{[\s\S]*\}|\[[\s\S]*\])", text)

        if match:
            json_str = match.group(0)
            try:
                # 2. 尝试解析提取出的字符串，验证是否为合法 JSON
                parsed_json = json.loads(json_str)

                # 3. 重新转储为标准字符串 (去除原本可能存在的缩进、多余空格等)
                # ensure_ascii=False 保证中文不会变成 \uXXXX
                return json.dumps(parsed_json, ensure_ascii=False)

            except json.JSONDecodeError as e:
                print(f"提取到了类似JSON的片段，但解析失败: {e}")
                return None
        else:
            print("未在文本中发现 JSON 结构")
            return None

    def extract_tool_params(self, combined_question: str, query: str, tool_id: str, tool_instructions: str) -> dict:
        """
        根据工具信息和查询提取调用参数
        
        Args:
            combined_question: 组合问题（用于缓存）
            tool_name: 工具名称
            query: 查询内容
            
        Returns:
            dict: 提取的参数字典
        """
        logger.info(f"[步骤3] 提取工具参数...")
        
        # 尝试从缓存读取
        if self.use_cache:
            cached_params = self.cache.get(combined_question, 'function_knowledge', 'tool_params.json')
            if cached_params:
                logger.info(f"✓ 使用缓存的参数: {cached_params}")
                # 记录缓存命中
                self.execution_detail["extract_params"].update({
                    "cached": True,
                    "params": cached_params
                })
                return cached_params
        
        try:
            # 获取工具信息
            tool_params = get_tool_params(tool_id)
            if not tool_params:
                logger.warning(f"  ⚠ 未找到工具 {tool_id} 的信息，使用默认参数")
                return {"keyword": query}
            
            # 加载prompt
            prompt_template = self._load_prompt("function_knowledge_extract_tool_params_prompt.md")
            prompt = prompt_template.format(
                query=query,
                all_tool_params=tool_params
            )
            
            # 调用LLM提取参数
            logger.info("  → 调用Gemini提取参数...")
            response_text = generate_text(prompt=prompt)
            
            # 解析JSON
            logger.info("  → 解析参数JSON...")
            try:
                # 清理可能的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()
                
                params = json.loads(response_text)
                
                logger.info(f"✓ 提取参数成功: {params}")
                
                # 写入缓存
                if self.use_cache:
                    self.cache.set(combined_question, 'function_knowledge', 'tool_params.json', params)
                
                # 记录详情
                self.execution_detail["extract_params"].update({
                    "cached": False,
                    "prompt": prompt,
                    "response": response_text,
                    "params": params
                })
                
                return params
                
            except json.JSONDecodeError as e:
                logger.error(f"  ✗ 解析JSON失败: {e}")
                logger.error(f"  响应内容: {response_text}")
                # 降级：使用query作为keyword
                default_params = {"keyword": query}
                logger.warning(f"  使用默认参数: {default_params}")
                return default_params
                
        except Exception as e:
            logger.error(f"✗ 提取工具参数失败: {e}")
            # 降级：使用query作为keyword
            return {"keyword": query}

    def save_knowledge_to_file(self, knowledge: str, combined_question: str):
        """保存获取到的知识到文件"""
        try:
            logger.info("[保存知识] 开始保存知识到文件...")

            # 获取问题hash
            import hashlib
            question_hash = hashlib.md5(combined_question.encode('utf-8')).hexdigest()[:12]

            # 获取缓存目录（和execution_record.json同级）
            if self.use_cache and self.cache:
                cache_dir = os.path.join(self.cache.base_cache_dir, question_hash)
            else:
                cache_dir = os.path.join(os.path.dirname(__file__), '.cache', question_hash)

            os.makedirs(cache_dir, exist_ok=True)

            # 保存到knowledge.txt
            knowledge_file = os.path.join(cache_dir, 'knowledge.txt')
            with open(knowledge_file, 'w', encoding='utf-8') as f:
                f.write(knowledge)

            logger.info(f"✓ 知识已保存到: {knowledge_file}")
            logger.info(f"  知识长度: {len(knowledge)} 字符")

        except Exception as e:
            logger.error(f"✗ 保存知识失败: {e}")

    def get_knowledge(self, question: str, post_info: str, persona_info: str) -> dict:
        """
        获取方法知识的主流程（重构后）
        
        Returns:
            dict: 完整的执行记录
        """
        import time
        timestamp = time.strftime("%Y-%m-%d %H:%M:%S")
        start_time = time.time()
        
        logger.info("=" * 80)
        logger.info(f"Function Knowledge - 开始处理")
        logger.info(f"问题: {question}")
        logger.info(f"帖子信息: {post_info}")
        logger.info(f"人设信息: {persona_info}")
        logger.info("=" * 80)
        
        # 组合问题的唯一标识
        combined_question = f"{question}||{post_info}||{persona_info}"
        
        try:
            # 步骤1: 生成Query
            query = self.generate_query(question, post_info, persona_info)
            
            # 步骤2: 选择工具
            tool_info = self.select_tool(combined_question, query)
            # tool_name = tool_info.get("工具名")
            tool_id = tool_info.get("工具调用ID")
            tool_instructions = tool_info.get("使用方法")
            if tool_id and tool_instructions:
                # 路径A: 使用工具
                # 步骤3: 提取参数
                arguments = self.extract_tool_params(combined_question, query, tool_id, tool_instructions)

                # 步骤4: 调用工具
                logger.info(f"[步骤4] 调用工具: {tool_id}")
                
                # 检查工具调用缓存
                if self.use_cache:
                    cached_tool_result = self.cache.get(combined_question, 'function_knowledge', 'tool_result.json')
                    if cached_tool_result:
                        logger.info(f"✓ 使用缓存的工具调用结果")
                        tool_result = cached_tool_result
                    else:
                        logger.info(f"  → 调用工具，参数: {arguments}")
                        tool_result = call_tool(tool_id, arguments)
                        # 缓存工具调用结果
                        self.cache.set(combined_question, 'function_knowledge', 'tool_result.json', tool_result)
                else:
                    logger.info(f"  → 调用工具，参数: {arguments}")
                    tool_result = call_tool(tool_id, arguments)

                logger.info(f"✓ 工具调用完成")
                
            else:
                # 路径B: 知识搜索
                logger.info("[步骤4] 未找到合适工具，调用 MultiSearch...")
                
                knowledge = get_multi_search_knowledge(query, cache_key=combined_question)
                
                # 异步保存知识到文件
                logger.info("[后台任务] 保存知识到文件...")
                threading.Thread(target=self.save_knowledge_to_file, args=(knowledge, combined_question)).start()
            
            # 计算执行时间并保存详情
            self.execution_detail["execution_time"] = time.time() - start_time
            self._save_execution_detail(combined_question)
            
            # 收集所有执行记录
            logger.info("=" * 80)
            logger.info("收集执行记录...")
            logger.info("=" * 80)
            
            from knowledge_v2.execution_collector import collect_and_save_execution_record
            
            execution_record = collect_and_save_execution_record(
                combined_question,
                {
                    "question": question,
                    "post_info": post_info,
                    "persona_info": persona_info,
                    "timestamp": timestamp
                }
            )
            
            logger.info("=" * 80)
            logger.info(f"✓ Function Knowledge 完成")
            logger.info(f"  执行时间: {execution_record.get('metadata', {}).get('execution_time', 0):.2f}秒")
            logger.info("=" * 80 + "\n")
            
            return execution_record
            
        except Exception as e:
            logger.error(f"✗ 执行失败: {e}")
            import traceback
            logger.error(traceback.format_exc())
            
            # 即使失败也尝试保存详情和收集记录
            try:
                self.execution_detail["execution_time"] = time.time() - start_time
                self._save_execution_detail(combined_question)
                
                from knowledge_v2.execution_collector import collect_and_save_execution_record
                execution_record = collect_and_save_execution_record(
                    combined_question,
                    {
                        "question": question,
                        "post_info": post_info,
                        "persona_info": persona_info,
                        "timestamp": timestamp
                    }
                )
                return execution_record
            except Exception as collect_error:
                logger.error(f"收集执行记录也失败: {collect_error}")
                # 返回基本错误信息
                return {
                    "input": {
                        "question": question,
                        "post_info": post_info,
                        "persona_info": persona_info,
                        "timestamp": timestamp
                    },
                    "result": {
                        "type": "error",
                        "content": f"执行失败: {str(e)}"
                    },
                    "metadata": {
                        "errors": [str(e)]
                    }
                }

if __name__ == "__main__":
    # 测试代码
    question = "教资查分这个信息怎么来的"
    post_info = "发帖时间：2025.11.07"
    persona_info = ""
    
    try:
        agent = FunctionKnowledge()
        execution_result = agent.get_knowledge(question, post_info, persona_info)
        print("=" * 50)
        print("执行结果:")
        print("=" * 50)
        print(json.dumps(execution_result, ensure_ascii=False, indent=2))
        print(f"\n完整JSON已保存到缓存目录")
    except Exception as e:
        logger.error(f"测试失败: {e}")