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()