mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-08-28 05:25:04 +00:00
196 lines
8.3 KiB
Python
196 lines
8.3 KiB
Python
"""测试 ReMeLight
|
||
|
||
演示 ReMeLight 的完整功能，并使用 AsMsgHandler 跟踪每步 Token 变化：
|
||
1. compact_tool_result - 压缩超长工具输出
|
||
2. compact_memory - 生成压缩摘要
|
||
3. summary_memory - 生成完整摘要并写入文件
|
||
4. pre_reasoning_hook - 推理前预处理钩子
|
||
5. memory_search - 语义搜索记忆
|
||
6. ReMeInMemoryMemory.estimate_tokens - 估算 Token 使用
|
||
7. ReMeInMemoryMemory.get_history_str - 获取格式化历史记录
|
||
"""
|
||
|
||
import asyncio
|
||
import logging
|
||
from test_utils import build_sample_messages, get_msg_handler
|
||
from reme.reme_light import ReMeLight
|
||
|
||
|
||
def print_token_change(_step_name: str, before: int, after: int):
|
||
"""打印 Token 变化统计。"""
|
||
change = after - before
|
||
change_pct = (change / before * 100) if before > 0 else 0
|
||
print(f" 📊 Token 统计: {before:,} → {after:,} (变化: {change:+,}, {change_pct:+.1f}%)")
|
||
|
||
|
||
# ==================== 主测试流程 ====================
|
||
async def main():
|
||
"""测试 ReMeLight 的完整功能，并跟踪每步 Token 变化。"""
|
||
# 初始化 AsMsgHandler 用于 Token 统计
|
||
msg_handler = get_msg_handler()
|
||
|
||
# 初始化 ReMeLight
|
||
reme = ReMeLight(
|
||
default_as_llm_config={"model_name": "qwen3.5-35b-a3b"},
|
||
# default_embedding_model_config={"model_name": "text-embedding-v4"},
|
||
default_file_store_config={"fts_enabled": True, "vector_enabled": False},
|
||
enable_load_env=True,
|
||
)
|
||
logging.getLogger("reme").setLevel(logging.WARNING)
|
||
await reme.start()
|
||
print("=" * 70)
|
||
print("ReMeLight 已启动")
|
||
print("=" * 70)
|
||
|
||
# 构建模拟对话历史（包含超长 tool_result，确保超过 128K token）
|
||
original_messages = build_sample_messages(include_large_tool_result=True)
|
||
initial_tokens = await msg_handler.count_msgs_token(original_messages)
|
||
|
||
print(f"\n[原始消息]: {len(original_messages)} 条, {initial_tokens:,} tokens")
|
||
print(f" 目标阈值: 128K = {128 * 1024:,} tokens")
|
||
print(f" 超出阈值: {initial_tokens > 128 * 1024}")
|
||
|
||
# ==================== 1. compact_tool_result ====================
|
||
print("\n" + "=" * 70)
|
||
print("[步骤 1] compact_tool_result - 压缩超长工具输出")
|
||
print("=" * 70)
|
||
|
||
# 重新获取原始消息
|
||
messages = build_sample_messages(include_large_tool_result=True)
|
||
tokens_before = await msg_handler.count_msgs_token(messages)
|
||
messages_after_step1 = await reme.compact_tool_result(messages)
|
||
tokens_after = await msg_handler.count_msgs_token(messages_after_step1)
|
||
|
||
print(f" 消息数量: {len(messages)} → {len(messages_after_step1)}")
|
||
print_token_change("compact_tool_result", tokens_before, tokens_after)
|
||
|
||
# ==================== 2. compact_memory ====================
|
||
print("\n" + "=" * 70)
|
||
print("[步骤 2] compact_memory - 生成结构化压缩摘要")
|
||
print("=" * 70)
|
||
|
||
# 重新获取原始消息
|
||
messages = build_sample_messages(include_large_tool_result=True)
|
||
tokens_before = await msg_handler.count_msgs_token(messages)
|
||
compact_summary = await reme.compact_memory(
|
||
messages=messages,
|
||
previous_summary="",
|
||
)
|
||
summary_tokens = await msg_handler.count_str_token(compact_summary)
|
||
|
||
print(f" 输入消息 tokens: {tokens_before:,}")
|
||
print(f" 压缩摘要长度: {len(compact_summary)} 字符, {summary_tokens:,} tokens")
|
||
print(f" 压缩比: {summary_tokens / tokens_before * 100:.1f}%" if tokens_before > 0 else " 压缩比: N/A")
|
||
print(f" 摘要预览: {compact_summary[:200]}..." if len(compact_summary) > 200 else f" 摘要: {compact_summary}")
|
||
|
||
# ==================== 3. summary_memory ====================
|
||
print("\n" + "=" * 70)
|
||
print("[步骤 3] summary_memory - 生成完整摘要并写入文件")
|
||
print("=" * 70)
|
||
|
||
# 重新获取原始消息
|
||
messages = build_sample_messages(include_large_tool_result=True)
|
||
tokens_before = await msg_handler.count_msgs_token(messages)
|
||
summary_result = await reme.summary_memory(messages=messages)
|
||
|
||
print(f" 输入消息 tokens: {tokens_before:,}")
|
||
print(f" 摘要结果长度: {len(summary_result)} 字符")
|
||
print(f" 摘要预览: {summary_result[:200]}..." if len(summary_result) > 200 else f" 摘要: {summary_result}")
|
||
|
||
# ==================== 4. pre_reasoning_hook ====================
|
||
print("\n" + "=" * 70)
|
||
print("[步骤 4] pre_reasoning_hook - 推理前预处理")
|
||
print("=" * 70)
|
||
|
||
# 重新获取原始消息
|
||
messages = build_sample_messages(include_large_tool_result=True)
|
||
tokens_before = await msg_handler.count_msgs_token(messages)
|
||
processed_messages, compressed_summary = await reme.pre_reasoning_hook(
|
||
messages=messages,
|
||
system_prompt="你是一个有帮助的 AI 助手。",
|
||
compressed_summary="",
|
||
max_input_length=128000,
|
||
compact_ratio=0.7,
|
||
memory_compact_reserve=10000,
|
||
enable_tool_result_compact=True,
|
||
tool_result_compact_keep_n=3,
|
||
)
|
||
tokens_after = await msg_handler.count_msgs_token(processed_messages)
|
||
compressed_summary_tokens = await msg_handler.count_str_token(compressed_summary)
|
||
|
||
print(f" 消息数量: {len(messages)} → {len(processed_messages)}")
|
||
print_token_change("pre_reasoning_hook", tokens_before, tokens_after)
|
||
print(f" 压缩摘要: {len(compressed_summary)} 字符, {compressed_summary_tokens:,} tokens")
|
||
print(f" 总上下文: {tokens_after + compressed_summary_tokens:,} tokens")
|
||
|
||
# ==================== 5. memory_search ====================
|
||
print("\n" + "=" * 70)
|
||
print("[步骤 5] memory_search - 语义搜索记忆")
|
||
print("=" * 70)
|
||
|
||
search_result = await reme.memory_search(query="Python 版本偏好", max_results=5)
|
||
if search_result.content:
|
||
print(f" 搜索结果: {search_result.content}")
|
||
else:
|
||
print(" 未找到相关记忆")
|
||
|
||
# ==================== 6 & 7. ReMeInMemoryMemory ====================
|
||
print("\n" + "=" * 70)
|
||
print("[步骤 6] ReMeInMemoryMemory - 会话内存管理")
|
||
print("=" * 70)
|
||
|
||
# 重新获取原始消息
|
||
messages = build_sample_messages(include_large_tool_result=True)
|
||
memory = reme.get_in_memory_memory()
|
||
for msg in messages:
|
||
await memory.add(msg)
|
||
print(f" 已添加 {len(messages)} 条原始消息到内存")
|
||
|
||
# 6.1 estimate_tokens
|
||
print("\n[6.1] estimate_tokens - 估算 Token 使用:")
|
||
token_stats = await memory.estimate_tokens(max_input_length=128000)
|
||
print(f" - 总消息数: {token_stats['total_messages']}")
|
||
print(f" - 消息 Token 数: {token_stats['messages_tokens']:,}")
|
||
print(f" - 压缩摘要 Token 数: {token_stats['compressed_summary_tokens']:,}")
|
||
print(f" - 预估总 Token 数: {token_stats['estimated_tokens']:,}")
|
||
print(f" - 最大输入长度: {token_stats['max_input_length']:,}")
|
||
print(f" - 上下文使用率: {token_stats['context_usage_ratio']:.2f}%")
|
||
|
||
# 6.2 get_history_str
|
||
print("\n[6.2] get_history_str - 格式化历史记录:")
|
||
history_str = await memory.get_history_str(max_input_length=128000)
|
||
print(history_str[:1000] + "..." if len(history_str) > 1000 else history_str)
|
||
|
||
# ==================== 等待后台任务完成 ====================
|
||
print("\n" + "=" * 70)
|
||
print("[步骤 7] 等待后台任务完成")
|
||
print("=" * 70)
|
||
await_result = await reme.await_summary_tasks()
|
||
print(f" 后台任务完成，结果长度: {len(await_result)} 字符")
|
||
|
||
# ==================== 总结 ====================
|
||
print("\n" + "=" * 70)
|
||
print("📊 Token 变化总结")
|
||
print("=" * 70)
|
||
print(f" 原始消息: {initial_tokens:,} tokens")
|
||
print(f" Step 1 compact_tool_result 后: {(await msg_handler.count_msgs_token(messages_after_step1)):,} tokens")
|
||
print(f" Step 2 compact_memory 摘要: {summary_tokens:,} tokens")
|
||
print(
|
||
f" Step 4 pre_reasoning_hook 后: {tokens_after:,} tokens + 摘要 {compressed_summary_tokens:,} "
|
||
f"tokens = {tokens_after + compressed_summary_tokens:,} tokens",
|
||
)
|
||
print(
|
||
f" 最大节省: {initial_tokens - tokens_after:,} "
|
||
f"tokens ({(initial_tokens - tokens_after) / initial_tokens * 100:.1f}%)",
|
||
)
|
||
print(f" 目标阈值: {128 * 1024:,} tokens")
|
||
|
||
# 关闭 ReMeLight
|
||
await reme.close()
|
||
print("\n" + "=" * 70)
|
||
print("ReMeLight 已关闭")
|
||
print("=" * 70)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
asyncio.run(main())
|