feat：优化记忆查询，添加时间信息

2025-12-31 14:06:14 +08:00
parent ce9e17df25
commit b296f0683f
5 changed files with 1071 additions and 25 deletions
--- a/scripts/test_memory_retrieval.py
+++ b/scripts/test_memory_retrieval.py
@@ -0,0 +1,447 @@
+import argparse
+import asyncio
+import os
+import sys
+import time
+import json
+import importlib
+from typing import Optional, Dict, Any
+from datetime import datetime
+
+# 强制使用 utf-8，避免控制台编码报错
+try:
+    if hasattr(sys.stdout, "reconfigure"):
+        sys.stdout.reconfigure(encoding="utf-8")
+    if hasattr(sys.stderr, "reconfigure"):
+        sys.stderr.reconfigure(encoding="utf-8")
+except Exception:
+    pass
+
+# 确保能导入 src.*
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
+
+from src.common.logger import initialize_logging, get_logger
+from src.common.database.database import db
+from src.common.database.database_model import LLMUsage
+from src.chat.message_receive.chat_stream import ChatStream
+from maim_message import UserInfo, GroupInfo
+
+logger = get_logger("test_memory_retrieval")
+
+# 使用 importlib 动态导入，避免循环导入问题
+def _import_memory_retrieval():
+    """使用 importlib 动态导入 memory_retrieval 模块，避免循环导入"""
+    try:
+        # 先导入 prompt_builder，检查 prompt 是否已经初始化
+        from src.chat.utils.prompt_builder import global_prompt_manager
+        
+        # 检查 memory_retrieval 相关的 prompt 是否已经注册
+        # 如果已经注册，说明模块可能已经通过其他路径初始化过了
+        prompt_already_init = "memory_retrieval_question_prompt" in global_prompt_manager._prompts
+        
+        module_name = "src.memory_system.memory_retrieval"
+        
+        # 如果 prompt 已经初始化，尝试直接使用已加载的模块
+        if prompt_already_init and module_name in sys.modules:
+            existing_module = sys.modules[module_name]
+            if hasattr(existing_module, 'init_memory_retrieval_prompt'):
+                return (
+                    existing_module.init_memory_retrieval_prompt,
+                    existing_module._react_agent_solve_question,
+                    existing_module._process_single_question,
+                )
+        
+        # 如果模块已经在 sys.modules 中但部分初始化，先移除它
+        if module_name in sys.modules:
+            existing_module = sys.modules[module_name]
+            if not hasattr(existing_module, 'init_memory_retrieval_prompt'):
+                # 模块部分初始化，移除它
+                logger.warning(f"检测到部分初始化的模块 {module_name}，尝试重新导入")
+                del sys.modules[module_name]
+                # 清理可能相关的部分初始化模块
+                keys_to_remove = []
+                for key in sys.modules.keys():
+                    if key.startswith('src.memory_system.') and key != 'src.memory_system':
+                        keys_to_remove.append(key)
+                for key in keys_to_remove:
+                    try:
+                        del sys.modules[key]
+                    except KeyError:
+                        pass
+        
+        # 在导入 memory_retrieval 之前，先确保所有可能触发循环导入的模块都已完全加载
+        # 这些模块在导入时可能会触发 memory_retrieval 的导入，所以我们需要先加载它们
+        try:
+            # 先导入可能触发循环导入的模块，让它们完成初始化
+            import src.config.config
+            import src.chat.utils.prompt_builder
+            # 尝试导入可能触发循环导入的模块（这些模块可能在模块级别导入了 memory_retrieval）
+            # 如果它们已经导入，就确保它们完全初始化
+            # 尝试导入可能触发循环导入的模块（这些模块可能在模块级别导入了 memory_retrieval）
+            # 如果它们已经导入，就确保它们完全初始化
+            try:
+                import src.chat.replyer.group_generator  # noqa: F401
+            except (ImportError, AttributeError):
+                pass  # 如果导入失败，继续
+            try:
+                import src.chat.replyer.private_generator  # noqa: F401
+            except (ImportError, AttributeError):
+                pass  # 如果导入失败，继续
+        except Exception as e:
+            logger.warning(f"预加载依赖模块时出现警告: {e}")
+        
+        # 现在尝试导入 memory_retrieval
+        # 如果此时仍然触发循环导入，说明有其他模块在模块级别导入了 memory_retrieval
+        memory_retrieval_module = importlib.import_module(module_name)
+        
+        return (
+            memory_retrieval_module.init_memory_retrieval_prompt,
+            memory_retrieval_module._react_agent_solve_question,
+            memory_retrieval_module._process_single_question,
+        )
+    except (ImportError, AttributeError) as e:
+        logger.error(f"导入 memory_retrieval 模块失败: {e}", exc_info=True)
+        raise
+
+
+def create_test_chat_stream(chat_id: str = "test_memory_retrieval") -> ChatStream:
+    """创建一个测试用的 ChatStream 对象"""
+    user_info = UserInfo(
+        platform="test",
+        user_id="test_user",
+        user_nickname="测试用户",
+    )
+    group_info = GroupInfo(
+        platform="test",
+        group_id="test_group",
+        group_name="测试群组",
+    )
+    return ChatStream(
+        stream_id=chat_id,
+        platform="test",
+        user_info=user_info,
+        group_info=group_info,
+    )
+
+
+def get_token_usage_since(start_time: float) -> Dict[str, Any]:
+    """获取从指定时间开始的token使用情况
+    
+    Args:
+        start_time: 开始时间戳
+        
+    Returns:
+        包含token使用统计的字典
+    """
+    try:
+        start_datetime = datetime.fromtimestamp(start_time)
+        
+        # 查询从开始时间到现在的所有memory相关的token使用记录
+        records = (
+            LLMUsage.select()
+            .where(
+                (LLMUsage.timestamp >= start_datetime)
+                & (
+                    (LLMUsage.request_type.like("%memory%"))
+                    | (LLMUsage.request_type == "memory.question")
+                    | (LLMUsage.request_type == "memory.react")
+                    | (LLMUsage.request_type == "memory.react.final")
+                )
+            )
+            .order_by(LLMUsage.timestamp.asc())
+        )
+        
+        total_prompt_tokens = 0
+        total_completion_tokens = 0
+        total_tokens = 0
+        total_cost = 0.0
+        request_count = 0
+        model_usage = {}  # 按模型统计
+        
+        for record in records:
+            total_prompt_tokens += record.prompt_tokens or 0
+            total_completion_tokens += record.completion_tokens or 0
+            total_tokens += record.total_tokens or 0
+            total_cost += record.cost or 0.0
+            request_count += 1
+            
+            # 按模型统计
+            model_name = record.model_name or "unknown"
+            if model_name not in model_usage:
+                model_usage[model_name] = {
+                    "prompt_tokens": 0,
+                    "completion_tokens": 0,
+                    "total_tokens": 0,
+                    "cost": 0.0,
+                    "request_count": 0,
+                }
+            model_usage[model_name]["prompt_tokens"] += record.prompt_tokens or 0
+            model_usage[model_name]["completion_tokens"] += record.completion_tokens or 0
+            model_usage[model_name]["total_tokens"] += record.total_tokens or 0
+            model_usage[model_name]["cost"] += record.cost or 0.0
+            model_usage[model_name]["request_count"] += 1
+        
+        return {
+            "total_prompt_tokens": total_prompt_tokens,
+            "total_completion_tokens": total_completion_tokens,
+            "total_tokens": total_tokens,
+            "total_cost": total_cost,
+            "request_count": request_count,
+            "model_usage": model_usage,
+        }
+    except Exception as e:
+        logger.error(f"获取token使用情况失败: {e}")
+        return {
+            "total_prompt_tokens": 0,
+            "total_completion_tokens": 0,
+            "total_tokens": 0,
+            "total_cost": 0.0,
+            "request_count": 0,
+            "model_usage": {},
+        }
+
+
+def format_thinking_steps(thinking_steps: list) -> str:
+    """格式化思考步骤为可读字符串"""
+    if not thinking_steps:
+        return "无思考步骤"
+    
+    lines = []
+    for step in thinking_steps:
+        iteration = step.get("iteration", "?")
+        thought = step.get("thought", "")
+        actions = step.get("actions", [])
+        observations = step.get("observations", [])
+        
+        lines.append(f"\n--- 迭代 {iteration} ---")
+        if thought:
+            lines.append(f"思考: {thought[:200]}...")
+        
+        if actions:
+            lines.append("行动:")
+            for action in actions:
+                action_type = action.get("action_type", "unknown")
+                action_params = action.get("action_params", {})
+                lines.append(f"  - {action_type}: {json.dumps(action_params, ensure_ascii=False)}")
+        
+        if observations:
+            lines.append("观察:")
+            for obs in observations:
+                obs_str = str(obs)[:200]
+                if len(str(obs)) > 200:
+                    obs_str += "..."
+                lines.append(f"  - {obs_str}")
+    
+    return "\n".join(lines)
+
+
+async def test_memory_retrieval(
+    question: str,
+    chat_id: str = "test_memory_retrieval",
+    context: str = "",
+    max_iterations: Optional[int] = None,
+) -> Dict[str, Any]:
+    """测试记忆检索功能
+    
+    Args:
+        question: 要查询的问题
+        chat_id: 聊天ID
+        context: 上下文信息
+        max_iterations: 最大迭代次数
+        
+    Returns:
+        包含测试结果的字典
+    """
+    print("\n" + "=" * 80)
+    print(f"[测试] 记忆检索测试")
+    print(f"[问题] {question}")
+    print("=" * 80)
+    
+    # 记录开始时间
+    start_time = time.time()
+    
+    # 延迟导入并初始化记忆检索prompt（这会自动加载 global_config）
+    # 注意：必须在函数内部调用，避免在模块级别触发循环导入
+    try:
+        init_memory_retrieval_prompt, _react_agent_solve_question, _ = _import_memory_retrieval()
+        
+        # 检查 prompt 是否已经初始化，避免重复初始化
+        from src.chat.utils.prompt_builder import global_prompt_manager
+        if "memory_retrieval_question_prompt" not in global_prompt_manager._prompts:
+            init_memory_retrieval_prompt()
+        else:
+            logger.debug("记忆检索 prompt 已经初始化，跳过重复初始化")
+    except Exception as e:
+        logger.error(f"初始化记忆检索模块失败: {e}", exc_info=True)
+        raise
+    
+    # 获取 global_config（此时应该已经加载）
+    from src.config.config import global_config
+    
+    # 直接调用 _react_agent_solve_question 来获取详细的迭代信息
+    if max_iterations is None:
+        max_iterations = global_config.memory.max_agent_iterations
+    
+    timeout = global_config.memory.agent_timeout_seconds
+    
+    print(f"\n[配置]")
+    print(f"  最大迭代次数: {max_iterations}")
+    print(f"  超时时间: {timeout}秒")
+    print(f"  聊天ID: {chat_id}")
+    
+    # 执行检索
+    print(f"\n[开始检索] {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())}")
+    
+    found_answer, answer, thinking_steps, is_timeout = await _react_agent_solve_question(
+        question=question,
+        chat_id=chat_id,
+        max_iterations=max_iterations,
+        timeout=timeout,
+        initial_info="",
+    )
+    
+    # 记录结束时间
+    end_time = time.time()
+    elapsed_time = end_time - start_time
+    
+    # 获取token使用情况
+    token_usage = get_token_usage_since(start_time)
+    
+    # 构建结果
+    result = {
+        "question": question,
+        "found_answer": found_answer,
+        "answer": answer,
+        "is_timeout": is_timeout,
+        "elapsed_time": elapsed_time,
+        "thinking_steps": thinking_steps,
+        "iteration_count": len(thinking_steps),
+        "token_usage": token_usage,
+    }
+    
+    # 输出结果
+    print(f"\n[检索完成] {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())}")
+    print(f"\n[结果]")
+    print(f"  是否找到答案: {'是' if found_answer else '否'}")
+    if found_answer and answer:
+        print(f"  答案: {answer}")
+    else:
+        print(f"  答案: (未找到答案)")
+    print(f"  是否超时: {'是' if is_timeout else '否'}")
+    print(f"  迭代次数: {len(thinking_steps)}")
+    print(f"  总耗时: {elapsed_time:.2f}秒")
+    
+    print(f"\n[Token使用情况]")
+    print(f"  总请求数: {token_usage['request_count']}")
+    print(f"  总Prompt Tokens: {token_usage['total_prompt_tokens']:,}")
+    print(f"  总Completion Tokens: {token_usage['total_completion_tokens']:,}")
+    print(f"  总Tokens: {token_usage['total_tokens']:,}")
+    print(f"  总成本: ${token_usage['total_cost']:.6f}")
+    
+    if token_usage['model_usage']:
+        print(f"\n[按模型统计]")
+        for model_name, usage in token_usage['model_usage'].items():
+            print(f"  {model_name}:")
+            print(f"    请求数: {usage['request_count']}")
+            print(f"    Prompt Tokens: {usage['prompt_tokens']:,}")
+            print(f"    Completion Tokens: {usage['completion_tokens']:,}")
+            print(f"    总Tokens: {usage['total_tokens']:,}")
+            print(f"    成本: ${usage['cost']:.6f}")
+    
+    print(f"\n[迭代详情]")
+    print(format_thinking_steps(thinking_steps))
+    
+    print("\n" + "=" * 80)
+    
+    return result
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(
+        description="测试记忆检索功能。可以输入一个问题，脚本会使用记忆检索的逻辑进行检索，并记录迭代信息、时间和token总消耗。"
+    )
+    parser.add_argument(
+        "--chat-id",
+        default="test_memory_retrieval",
+        help="测试用的聊天ID（默认: test_memory_retrieval）",
+    )
+    parser.add_argument(
+        "--context",
+        default="",
+        help="上下文信息（可选）",
+    )
+    parser.add_argument(
+        "--output",
+        "-o",
+        help="将结果保存到JSON文件（可选）",
+    )
+    
+    args = parser.parse_args()
+    
+    # 初始化日志（使用较低的详细程度，避免输出过多日志）
+    initialize_logging(verbose=False)
+    
+    # 交互式输入问题
+    print("\n" + "=" * 80)
+    print("记忆检索测试工具")
+    print("=" * 80)
+    question = input("\n请输入要查询的问题: ").strip()
+    if not question:
+        print("错误: 问题不能为空")
+        return
+    
+    # 交互式输入最大迭代次数
+    max_iterations_input = input("\n请输入最大迭代次数（直接回车使用配置默认值）: ").strip()
+    max_iterations = None
+    if max_iterations_input:
+        try:
+            max_iterations = int(max_iterations_input)
+            if max_iterations <= 0:
+                print("警告: 迭代次数必须大于0，将使用配置默认值")
+                max_iterations = None
+        except ValueError:
+            print("警告: 无效的迭代次数，将使用配置默认值")
+            max_iterations = None
+    
+    # 连接数据库
+    try:
+        db.connect(reuse_if_open=True)
+    except Exception as e:
+        logger.error(f"数据库连接失败: {e}")
+        print(f"错误: 数据库连接失败: {e}")
+        return
+    
+    # 运行测试
+    try:
+        result = asyncio.run(
+            test_memory_retrieval(
+                question=question,
+                chat_id=args.chat_id,
+                context=args.context,
+                max_iterations=max_iterations,
+            )
+        )
+        
+        # 如果指定了输出文件，保存结果
+        if args.output:
+            # 将thinking_steps转换为可序列化的格式
+            output_result = result.copy()
+            with open(args.output, "w", encoding="utf-8") as f:
+                json.dump(output_result, f, ensure_ascii=False, indent=2)
+            print(f"\n[结果已保存] {args.output}")
+        
+    except KeyboardInterrupt:
+        print("\n\n[中断] 用户中断测试")
+    except Exception as e:
+        logger.error(f"测试失败: {e}", exc_info=True)
+        print(f"\n[错误] 测试失败: {e}")
+    finally:
+        try:
+            db.close()
+        except Exception:
+            pass
+
+
+if __name__ == "__main__":
+    main()
+