test_reme_light.py

"""测试 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())