Agent/agent/trace/compaction.py

"""
Context 压缩 — 两级压缩策略

Level 1: GoalTree 过滤（确定性，零成本）
  - 跳过 completed/abandoned goals 的消息（信息已在 GoalTree summary 中）
  - 始终保留：system prompt、第一条 user message、当前 focus goal 的消息

Level 2: LLM 总结（仅在 Level 1 后仍超限时触发）
  - 在消息列表末尾追加压缩 prompt → 主模型回复 → summary 存为新消息
  - summary 的 parent_sequence 跳过被压缩的范围

压缩不修改存储：原始消息永远保留在 messages/，通过 parent_sequence 树结构实现跳过。
"""

import logging
from dataclasses import dataclass
from typing import List, Dict, Any, Optional, Set

from .goal_models import GoalTree
from .models import Message

logger = logging.getLogger(__name__)


# ===== 模型 Context Window（tokens）=====

MODEL_CONTEXT_WINDOWS: Dict[str, int] = {
    # Anthropic Claude
    "claude-sonnet-4": 200_000,
    "claude-opus-4": 200_000,
    "claude-3-5-sonnet": 200_000,
    "claude-3-5-haiku": 200_000,
    "claude-3-opus": 200_000,
    "claude-3-sonnet": 200_000,
    "claude-3-haiku": 200_000,
    # OpenAI
    "gpt-4o": 128_000,
    "gpt-4o-mini": 128_000,
    "gpt-4-turbo": 128_000,
    "gpt-4": 8_192,
    "o1": 200_000,
    "o3-mini": 200_000,
    # Google Gemini
    "gemini-2.5-pro": 1_000_000,
    "gemini-2.5-flash": 1_000_000,
    "gemini-2.0-flash": 1_000_000,
    "gemini-1.5-pro": 2_000_000,
    "gemini-1.5-flash": 1_000_000,
    # DeepSeek
    "deepseek-chat": 64_000,
    "deepseek-r1": 64_000,
}

DEFAULT_CONTEXT_WINDOW = 200_000


def get_context_window(model: str) -> int:
    """
    根据模型名称获取 context window 大小。

    支持带 provider 前缀的模型名（如 "anthropic/claude-sonnet-4.5"）和
    带版本后缀的名称（如 "claude-3-5-sonnet-20241022"）。
    """
    # 去掉 provider 前缀
    name = model.split("/")[-1].lower()

    # 精确匹配
    if name in MODEL_CONTEXT_WINDOWS:
        return MODEL_CONTEXT_WINDOWS[name]

    # 前缀匹配（处理版本后缀）
    for key, window in MODEL_CONTEXT_WINDOWS.items():
        if name.startswith(key):
            return window

    return DEFAULT_CONTEXT_WINDOW


# ===== 配置 =====

@dataclass
class CompressionConfig:
    """压缩配置"""
    max_tokens: int = 0                # 最大 token 数（0 = 自动：context_window * 0.5）
    threshold_ratio: float = 0.5       # 触发压缩的阈值 = context_window 的比例
    keep_recent_messages: int = 10     # Level 1 中始终保留最近 N 条消息

    def get_max_tokens(self, model: str) -> int:
        """获取实际的 max_tokens（如果为 0 则自动计算）"""
        if self.max_tokens > 0:
            return self.max_tokens
        window = get_context_window(model)
        return int(window * self.threshold_ratio)


# ===== Level 1: GoalTree 过滤 =====

def filter_by_goal_status(
    messages: List[Message],
    goal_tree: Optional[GoalTree],
) -> List[Message]:
    """
    Level 1 过滤：跳过 completed/abandoned goals 的消息

    始终保留：
    - goal_id 为 None 的消息（system prompt、初始 user message）
    - 当前 focus goal 及其祖先链上的消息
    - in_progress 和 pending goals 的消息

    跳过：
    - completed 且不在焦点路径上的 goals 的消息
    - abandoned goals 的消息

    Args:
        messages: 主路径上的有序消息列表
        goal_tree: GoalTree 实例

    Returns:
        过滤后的消息列表
    """
    if not goal_tree or not goal_tree.goals:
        return messages

    # 构建焦点路径（当前焦点 + 父链 + 直接子节点）
    focus_path = _get_focus_path(goal_tree)

    # 构建需要跳过的 goal IDs
    skip_goal_ids: Set[str] = set()
    for goal in goal_tree.goals:
        if goal.id in focus_path:
            continue  # 焦点路径上的 goal 始终保留
        if goal.status in ("completed", "abandoned"):
            skip_goal_ids.add(goal.id)

    # 过滤消息
    result = []
    for msg in messages:
        if msg.goal_id is None:
            result.append(msg)  # 无 goal 的消息始终保留
        elif msg.goal_id not in skip_goal_ids:
            result.append(msg)  # 不在跳过列表中的消息保留

    return result


def _get_focus_path(goal_tree: GoalTree) -> Set[str]:
    """获取焦点路径上的所有 goal IDs（焦点 + 父链 + 直接子节点）"""
    focus_ids: Set[str] = set()

    if not goal_tree.current_id:
        return focus_ids

    # 焦点自身
    focus_ids.add(goal_tree.current_id)

    # 父链
    goal = goal_tree.find(goal_tree.current_id)
    while goal and goal.parent_id:
        focus_ids.add(goal.parent_id)
        goal = goal_tree.find(goal.parent_id)

    # 直接子节点
    children = goal_tree.get_children(goal_tree.current_id)
    for child in children:
        focus_ids.add(child.id)

    return focus_ids


# ===== Token 估算 =====

def estimate_tokens(messages: List[Dict[str, Any]]) -> int:
    """
    估算消息列表的 token 数量

    简单估算：字符数 / 4。实际使用时应该用 tiktoken 或 API 返回的 token 数。
    """
    total_chars = 0
    for msg in messages:
        content = msg.get("content", "")
        if isinstance(content, str):
            total_chars += len(content)
        elif isinstance(content, list):
            for part in content:
                if isinstance(part, dict) and part.get("type") == "text":
                    total_chars += len(part.get("text", ""))
        # tool_calls
        tool_calls = msg.get("tool_calls")
        if tool_calls and isinstance(tool_calls, list):
            for tc in tool_calls:
                if isinstance(tc, dict):
                    func = tc.get("function", {})
                    total_chars += len(func.get("name", ""))
                    args = func.get("arguments", "")
                    if isinstance(args, str):
                        total_chars += len(args)

    return total_chars // 4


def estimate_tokens_from_messages(messages: List[Message]) -> int:
    """从 Message 对象列表估算 token 数"""
    return estimate_tokens([msg.to_llm_dict() for msg in messages])


def needs_level2_compression(
    token_count: int,
    config: CompressionConfig,
    model: str = "",
) -> bool:
    """判断是否需要触发 Level 2 压缩"""
    limit = config.get_max_tokens(model) if model else config.max_tokens
    return token_count > limit


# ===== Level 2: 压缩 Prompt =====

COMPRESSION_PROMPT = """请对以上对话历史进行压缩总结。

要求：
1. 保留关键决策、结论和产出（如创建的文件、修改的代码、得出的分析结论）
2. 保留重要的上下文（如用户的要求、约束条件、之前的讨论结果）
3. 省略中间探索过程、重复的工具调用细节
4. 使用结构化格式（标题 + 要点）
5. 控制在 2000 字以内

当前 GoalTree 状态（完整版，含 summary）：
{goal_tree_prompt}
"""

REFLECT_PROMPT = """请回顾以上整个执行过程，提取有价值的经验教训。

关注以下方面：
1. **人工干预**：如果有用户中途修改了指令或纠正了方向，说明之前的决策哪里有问题
2. **弯路**：哪些尝试是不必要的，有没有更直接的方法
3. **好的决策**：哪些判断和选择是正确的，值得记住
4. **工具使用**：哪些工具用法是高效的，哪些可以改进

请以简洁的规则列表形式输出，每条规则格式为：
- 当遇到 [条件] 时，应该 [动作]（原因：[简短说明]）
"""


def build_compression_prompt(goal_tree: Optional[GoalTree]) -> str:
    """构建 Level 2 压缩 prompt"""
    goal_prompt = ""
    if goal_tree:
        goal_prompt = goal_tree.to_prompt(include_summary=True)
    return COMPRESSION_PROMPT.format(goal_tree_prompt=goal_prompt)


def build_reflect_prompt() -> str:
    """构建反思 prompt"""
    return REFLECT_PROMPT