context_compaction.py

"""
KKCode 上下文压缩示例
====================

演示 KKCode 的两层渐进式上下文压缩机制：
- Layer 1 (落盘): 将早期消息序列化到磁盘，保留结构化摘要
- Layer 2 (LLM 摘要): 对落盘内容生成高密度语义摘要
- 熔断器: 防止过度压缩导致信息丢失
- PTL 重试: 压缩失败时自动重试
- 恢复段: 保留最近 N 轮对话不被压缩

本示例模拟长会话场景下的上下文管理。

运行方式：
    python examples/context_compaction.py

前提：
    - 已安装依赖：uv sync
"""

import sys
import os

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from kkcode.context.manager import (
    CompactCircuitBreaker,
    RecoveryState,
)
from kkcode.conversation import ConversationManager


def demonstrate_compaction_mechanism() -> None:
    """演示上下文压缩机制（无需 LLM 连接）。"""

    print("=" * 60)
    print("KKCode 两层上下文压缩机制演示")
    print("=" * 60)

    # 1. 展示熔断器机制
    print("\n--- 熔断器 (CompactCircuitBreaker) ---")
    breaker = CompactCircuitBreaker()
    print(f"初始连续失败计数: {breaker.consecutive_failures}")
    print(f"熔断阈值 (max_failures): {breaker.max_failures}")
    print(f"初始 is_open: {breaker.is_open()}")

    # 模拟连续压缩失败
    for i in range(3):
        breaker.record_failure()
        print(
            f"  记录第 {i+1} 次失败 → 连续失败={breaker.consecutive_failures} "
            f"is_open={breaker.is_open()}"
        )

    # 模拟成功恢复
    breaker.record_success()
    print(
        f"  记录成功 → 连续失败={breaker.consecutive_failures} "
        f"is_open={breaker.is_open()}"
    )

    # 2. 展示恢复段（RecoveryState）
    print("\n--- 恢复段 (RecoveryState) ---")
    print("RecoveryState 在 Layer 2 压缩中存活：记录已读文件/已调用技能，")
    print("即便对话历史被清空，模型仍保有工作上下文。")
    recovery = RecoveryState()
    recovery.record_file_read("auth.py", "class User: ...")
    recovery.record_file_read("models/base.py", "class BaseModel: ...")
    recovery.record_skill_invocation("refactor-sop", "拆类步骤: 1.提取基类 ...")
    files = recovery.snapshot_files(limit=10)
    skills = recovery.snapshot_skills()
    print(f"已记录文件快照: {[f.path for f in files]}")
    print(f"已记录技能快照: {[s.name for s in skills]}")
    print("压缩时这些快照会被重新附加到摘要 user 消息上 → 信息零丢失")

    # 3. 模拟长会话的上下文演变
    print("\n--- 长会话上下文演变模拟 ---")
    conv = ConversationManager()

    # 模拟 20 轮对话
    messages = [
        ("user", "帮我重构 auth 模块"),
        ("assistant", "好的，我先看一下当前的 auth.py 结构..."),
        ("tool", "[ReadFile] auth.py → 245 行"),
        ("assistant", "auth.py 有 4 个类: User, Role, Permission, TokenManager"),
        ("user", "把 User 类拆成 BaseModel + UserProfile"),
        ("tool", "[WriteFile] models/base.py → 创建"),
        ("tool", "[WriteFile] models/profile.py → 创建"),
        ("tool", "[Edit] auth.py → 改为继承 BaseModel"),
        ("assistant", "拆分完成，运行测试验证..."),
        ("tool", "[Bash] pytest tests/test_auth.py → 12 passed"),
        ("user", "加上 JWT token 刷新逻辑"),
        ("assistant", "我来在 TokenManager 中添加 refresh 方法"),
        ("tool", "[Edit] auth.py → 添加 refresh_token()"),
        ("tool", "[Bash] pytest tests/test_auth.py → 14 passed"),
        ("user", "写一个中间件版本"),
        ("tool", "[WriteFile] middleware/auth_middleware.py → 创建"),
        ("assistant", "中间件已创建，支持 token 验证和自动刷新"),
        ("user", "更新 README 文档"),
        ("tool", "[ReadFile] README.md"),
        ("tool", "[Edit] README.md → 更新 API 文档"),
        ("assistant", "文档已更新"),
    ]

    total_chars = 0
    for role, content in messages:
        if role == "user":
            conv.add_user_message(content)
        elif role == "assistant":
            # 模拟 assistant 消息
            pass
        total_chars += len(content)

    print(f"模拟轮次: {len([m for m in messages if m[0] == 'user'])} 轮用户对话")
    print(f"总字符数: ~{total_chars:,}")
    print(f"估算 Token: ~{total_chars // 3:,} (按 3 字符/token)")

    # 4. 展示压缩策略
    print("\n--- 两层压缩策略 ---")
    print("""
    ┌────────────────────────────────────┐
    │           完整上下文                │
    │  (Token 用量 < 50% 窗口)           │
    │  → 不压缩                          │
    ├────────────────────────────────────┤
    │           完整上下文                │
    │  (Token 用量 50%~75%)              │
    │  → Layer 1: 早期消息落盘            │
    │  → 保留结构化索引 + 恢复段快照       │
    ├────────────────────────────────────┤
    │  Layer 1 摘要 + 恢复段             │
    │  (Token 用量 > 75%)                │
    │  → Layer 2: LLM 语义摘要           │
    │  → 进一步压缩 + 重新附加文件/技能快照 │
    ├────────────────────────────────────┤
    │  Layer 2 摘要 + 恢复段             │
    │  (接近窗口上限 / 连续压缩失败)      │
    │  → 熔断器介入                      │
    │  → 停止自动压缩，保留已读上下文      │
    └────────────────────────────────────┘

    关键设计:
    - PTL 重试: 压缩失败自动重试最多 3 次
    - 恢复段: 已读文件/已调用技能的快照在压缩后重新附加，工作上下文不丢失
    - 信息保留率测试: 通过 SWE-bench 验证
      压缩后 Agent 仍能正确定位和修复 bug
    """)

    # 5. 自动触发条件
    print("--- 自动触发条件 ---")
    context_window = 128000  # 以 deepseek-chat 为例
    usage_ratio = (total_chars // 3) / context_window * 100
    print(f"当前上下文使用率: {usage_ratio:.1f}%")
    print(f"触发阈值:")
    print(f"  Layer 1 (落盘):   > 50% ({context_window * 0.5 / 1000:.0f}K tokens)")
    print(f"  Layer 2 (摘要):   > 75% ({context_window * 0.75 / 1000:.0f}K tokens)")
    print(f"  熔断器介入:       连续失败 >= {CompactCircuitBreaker().max_failures} 次")


if __name__ == "__main__":
    demonstrate_compaction_mechanism()