mirror of https://github.com/Mai-with-u/MaiBot.git
508 lines
19 KiB
Python
508 lines
19 KiB
Python
import argparse
|
||
import asyncio
|
||
import os
|
||
import sys
|
||
import time
|
||
import json
|
||
import importlib
|
||
from typing import 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
|
||
|
||
logger = get_logger("compare_finish_search_token")
|
||
|
||
|
||
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 _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,
|
||
)
|
||
|
||
# 如果模块已经在 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）
|
||
# 如果它们已经导入，就确保它们完全初始化
|
||
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,
|
||
)
|
||
except (ImportError, AttributeError) as e:
|
||
logger.error(f"导入 memory_retrieval 模块失败: {e}", exc_info=True)
|
||
raise
|
||
|
||
|
||
def _init_tools_without_finish_search():
|
||
"""初始化工具但不注册 finish_search"""
|
||
from src.memory_system.retrieval_tools import (
|
||
register_query_chat_history,
|
||
register_query_person_info,
|
||
register_query_words,
|
||
)
|
||
from src.memory_system.retrieval_tools.tool_registry import get_tool_registry
|
||
from src.config.config import global_config
|
||
|
||
# 清空工具注册器
|
||
tool_registry = get_tool_registry()
|
||
tool_registry.tools.clear()
|
||
|
||
# 注册除 finish_search 外的所有工具
|
||
register_query_chat_history()
|
||
register_query_person_info()
|
||
register_query_words()
|
||
|
||
# 如果启用 LPMM agent 模式，也注册 LPMM 工具
|
||
if global_config.lpmm_knowledge.lpmm_mode == "agent":
|
||
from src.memory_system.retrieval_tools.query_lpmm_knowledge import register_tool as register_lpmm_knowledge
|
||
register_lpmm_knowledge()
|
||
|
||
logger.info("已初始化工具（不包含 finish_search）")
|
||
|
||
|
||
def _init_tools_with_finish_search():
|
||
"""初始化工具并注册 finish_search"""
|
||
from src.memory_system.retrieval_tools.tool_registry import get_tool_registry
|
||
from src.memory_system.retrieval_tools import init_all_tools
|
||
|
||
# 清空工具注册器
|
||
tool_registry = get_tool_registry()
|
||
tool_registry.tools.clear()
|
||
|
||
# 初始化所有工具（包括 finish_search）
|
||
init_all_tools()
|
||
logger.info("已初始化工具（包含 finish_search）")
|
||
|
||
|
||
async def get_prompt_tokens_for_tools(
|
||
question: str,
|
||
chat_id: str,
|
||
use_finish_search: bool,
|
||
) -> Dict[str, Any]:
|
||
"""获取使用不同工具配置时的prompt token消耗
|
||
|
||
Args:
|
||
question: 要查询的问题
|
||
chat_id: 聊天ID
|
||
use_finish_search: 是否使用 finish_search 工具
|
||
|
||
Returns:
|
||
包含prompt token信息的字典
|
||
"""
|
||
# 先初始化 prompt（如果还未初始化）
|
||
# 注意：init_memory_retrieval_prompt 会调用 init_all_tools，所以我们需要在它之后重新设置工具
|
||
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, _ = _import_memory_retrieval()
|
||
init_memory_retrieval_prompt()
|
||
|
||
# 初始化工具（根据参数决定是否包含 finish_search）
|
||
# 必须在 init_memory_retrieval_prompt 之后调用，因为它会调用 init_all_tools
|
||
if use_finish_search:
|
||
_init_tools_with_finish_search()
|
||
else:
|
||
_init_tools_without_finish_search()
|
||
|
||
# 获取工具注册器
|
||
from src.memory_system.retrieval_tools.tool_registry import get_tool_registry
|
||
tool_registry = get_tool_registry()
|
||
tool_definitions = tool_registry.get_tool_definitions()
|
||
|
||
# 验证工具列表（调试用）
|
||
tool_names = [tool["name"] for tool in tool_definitions]
|
||
if use_finish_search:
|
||
if "finish_search" not in tool_names:
|
||
logger.warning("期望包含 finish_search 工具，但工具列表中未找到")
|
||
else:
|
||
if "finish_search" in tool_names:
|
||
logger.warning("期望不包含 finish_search 工具，但工具列表中找到了，将移除")
|
||
# 移除 finish_search 工具
|
||
tool_registry.tools.pop("finish_search", None)
|
||
tool_definitions = tool_registry.get_tool_definitions()
|
||
tool_names = [tool["name"] for tool in tool_definitions]
|
||
|
||
# 构建第一次调用的prompt（模拟_react_agent_solve_question的第一次调用）
|
||
from src.config.config import global_config
|
||
bot_name = global_config.bot.nickname
|
||
time_now = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
|
||
|
||
# 构建head_prompt
|
||
head_prompt = await global_prompt_manager.format_prompt(
|
||
"memory_retrieval_react_prompt_head",
|
||
bot_name=bot_name,
|
||
time_now=time_now,
|
||
question=question,
|
||
collected_info="",
|
||
current_iteration=1,
|
||
remaining_iterations=global_config.memory.max_agent_iterations - 1,
|
||
max_iterations=global_config.memory.max_agent_iterations,
|
||
)
|
||
|
||
# 构建消息列表（只包含system message，模拟第一次调用）
|
||
from src.llm_models.payload_content.message import MessageBuilder, RoleType
|
||
messages = []
|
||
system_builder = MessageBuilder()
|
||
system_builder.set_role(RoleType.System)
|
||
system_builder.add_text_content(head_prompt)
|
||
messages.append(system_builder.build())
|
||
|
||
# 调用LLM API来计算token（只调用一次，不实际执行）
|
||
from src.llm_models.utils_model import LLMRequest, RequestType
|
||
from src.config.config import model_config
|
||
|
||
# 创建LLM请求对象
|
||
llm_request = LLMRequest(model_set=model_config.model_task_config.tool_use, request_type="memory.react.compare")
|
||
|
||
# 构建工具选项
|
||
tool_built = llm_request._build_tool_options(tool_definitions)
|
||
|
||
# 直接调用 _execute_request 以获取完整的响应对象（包含 usage）
|
||
response, model_info = await llm_request._execute_request(
|
||
request_type=RequestType.RESPONSE,
|
||
message_factory=lambda _client, *, _messages=messages: _messages,
|
||
temperature=None,
|
||
max_tokens=None,
|
||
tool_options=tool_built,
|
||
)
|
||
|
||
# 从响应中获取token使用情况
|
||
prompt_tokens = 0
|
||
completion_tokens = 0
|
||
total_tokens = 0
|
||
|
||
if response and hasattr(response, 'usage') and response.usage:
|
||
prompt_tokens = response.usage.prompt_tokens or 0
|
||
completion_tokens = response.usage.completion_tokens or 0
|
||
total_tokens = response.usage.total_tokens or 0
|
||
|
||
return {
|
||
"use_finish_search": use_finish_search,
|
||
"prompt_tokens": prompt_tokens,
|
||
"completion_tokens": completion_tokens,
|
||
"total_tokens": total_tokens,
|
||
"tool_count": len(tool_definitions),
|
||
"tool_names": [tool["name"] for tool in tool_definitions],
|
||
}
|
||
|
||
|
||
async def compare_prompt_tokens(
|
||
question: str,
|
||
chat_id: str = "compare_finish_search",
|
||
) -> Dict[str, Any]:
|
||
"""对比使用 finish_search 工具与否的输入 token 差异
|
||
|
||
只运行一次，只计算输入 token 的差异，确保除了工具定义外其他内容一致
|
||
|
||
Args:
|
||
question: 要查询的问题
|
||
chat_id: 聊天ID
|
||
|
||
Returns:
|
||
包含对比结果的字典
|
||
"""
|
||
print("\n" + "=" * 80)
|
||
print("finish_search 工具 输入 Token 消耗对比测试")
|
||
print("=" * 80)
|
||
print(f"\n[测试问题] {question}")
|
||
print(f"[聊天ID] {chat_id}")
|
||
print("\n注意: 只对比第一次LLM调用的输入token差异，不运行完整迭代流程")
|
||
|
||
# 第一次测试：不使用 finish_search
|
||
print("\n" + "-" * 80)
|
||
print("[测试 1/2] 不使用 finish_search 工具")
|
||
print("-" * 80)
|
||
result_without = await get_prompt_tokens_for_tools(
|
||
question=question,
|
||
chat_id=f"{chat_id}_without",
|
||
use_finish_search=False,
|
||
)
|
||
|
||
print(f"\n[结果]")
|
||
print(f" 工具数量: {result_without['tool_count']}")
|
||
print(f" 工具列表: {', '.join(result_without['tool_names'])}")
|
||
print(f" 输入 Prompt Tokens: {result_without['prompt_tokens']:,}")
|
||
|
||
# 等待一下，确保数据库记录已写入
|
||
await asyncio.sleep(1)
|
||
|
||
# 第二次测试：使用 finish_search
|
||
print("\n" + "-" * 80)
|
||
print("[测试 2/2] 使用 finish_search 工具")
|
||
print("-" * 80)
|
||
result_with = await get_prompt_tokens_for_tools(
|
||
question=question,
|
||
chat_id=f"{chat_id}_with",
|
||
use_finish_search=True,
|
||
)
|
||
|
||
print(f"\n[结果]")
|
||
print(f" 工具数量: {result_with['tool_count']}")
|
||
print(f" 工具列表: {', '.join(result_with['tool_names'])}")
|
||
print(f" 输入 Prompt Tokens: {result_with['prompt_tokens']:,}")
|
||
|
||
# 对比结果
|
||
print("\n" + "=" * 80)
|
||
print("[对比结果]")
|
||
print("=" * 80)
|
||
|
||
prompt_token_diff = result_with['prompt_tokens'] - result_without['prompt_tokens']
|
||
prompt_token_diff_percent = (prompt_token_diff / result_without['prompt_tokens'] * 100) if result_without['prompt_tokens'] > 0 else 0
|
||
|
||
tool_count_diff = result_with['tool_count'] - result_without['tool_count']
|
||
|
||
print(f"\n[输入 Prompt Token 对比]")
|
||
print(f" 不使用 finish_search: {result_without['prompt_tokens']:,} tokens")
|
||
print(f" 使用 finish_search: {result_with['prompt_tokens']:,} tokens")
|
||
print(f" 差异: {prompt_token_diff:+,} tokens ({prompt_token_diff_percent:+.2f}%)")
|
||
|
||
print(f"\n[工具数量对比]")
|
||
print(f" 不使用 finish_search: {result_without['tool_count']} 个工具")
|
||
print(f" 使用 finish_search: {result_with['tool_count']} 个工具")
|
||
print(f" 差异: {tool_count_diff:+d} 个工具")
|
||
|
||
print(f"\n[工具列表对比]")
|
||
without_tools = set(result_without['tool_names'])
|
||
with_tools = set(result_with['tool_names'])
|
||
only_with = with_tools - without_tools
|
||
only_without = without_tools - with_tools
|
||
|
||
if only_with:
|
||
print(f" 仅在 '使用 finish_search' 中的工具: {', '.join(only_with)}")
|
||
if only_without:
|
||
print(f" 仅在 '不使用 finish_search' 中的工具: {', '.join(only_without)}")
|
||
if not only_with and not only_without:
|
||
print(f" 工具列表相同（除了 finish_search）")
|
||
|
||
# 显示其他token信息
|
||
print(f"\n[其他 Token 信息]")
|
||
print(f" Completion Tokens (不使用 finish_search): {result_without.get('completion_tokens', 0):,}")
|
||
print(f" Completion Tokens (使用 finish_search): {result_with.get('completion_tokens', 0):,}")
|
||
print(f" 总 Tokens (不使用 finish_search): {result_without.get('total_tokens', 0):,}")
|
||
print(f" 总 Tokens (使用 finish_search): {result_with.get('total_tokens', 0):,}")
|
||
|
||
print("\n" + "=" * 80)
|
||
|
||
return {
|
||
"question": question,
|
||
"without_finish_search": result_without,
|
||
"with_finish_search": result_with,
|
||
"comparison": {
|
||
"prompt_token_diff": prompt_token_diff,
|
||
"prompt_token_diff_percent": prompt_token_diff_percent,
|
||
"tool_count_diff": tool_count_diff,
|
||
},
|
||
}
|
||
|
||
|
||
def main() -> None:
|
||
parser = argparse.ArgumentParser(
|
||
description="对比使用 finish_search 工具与否的 token 消耗差异"
|
||
)
|
||
parser.add_argument(
|
||
"--chat-id",
|
||
default="compare_finish_search",
|
||
help="测试用的聊天ID（默认: compare_finish_search）",
|
||
)
|
||
parser.add_argument(
|
||
"--output",
|
||
"-o",
|
||
help="将结果保存到JSON文件（可选）",
|
||
)
|
||
|
||
args = parser.parse_args()
|
||
|
||
# 初始化日志（使用较低的详细程度，避免输出过多日志）
|
||
initialize_logging(verbose=False)
|
||
|
||
# 交互式输入问题
|
||
print("\n" + "=" * 80)
|
||
print("finish_search 工具 Token 消耗对比测试工具")
|
||
print("=" * 80)
|
||
question = input("\n请输入要查询的问题: ").strip()
|
||
if not question:
|
||
print("错误: 问题不能为空")
|
||
return
|
||
|
||
# 连接数据库
|
||
try:
|
||
db.connect(reuse_if_open=True)
|
||
except Exception as e:
|
||
logger.error(f"数据库连接失败: {e}")
|
||
print(f"错误: 数据库连接失败: {e}")
|
||
return
|
||
|
||
# 运行对比测试
|
||
try:
|
||
result = asyncio.run(
|
||
compare_prompt_tokens(
|
||
question=question,
|
||
chat_id=args.chat_id,
|
||
)
|
||
)
|
||
|
||
# 如果指定了输出文件，保存结果
|
||
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()
|
||
|